Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
 import re
+import json
 from functools import lru_cache
 
 import gradio as gr
@@ -9,13 +10,11 @@ import torch
 # Writable caches for HF + Gradio (fixes PermissionError in Spaces)
 # -------------------
 os.environ.setdefault("HF_HOME", "/data/.cache/huggingface")
-# Removed TRANSFORMERS_CACHE (deprecated warning in Transformers v5+)
 os.environ.setdefault("HF_HUB_CACHE", "/data/.cache/huggingface/hub")
 os.environ.setdefault("GRADIO_TEMP_DIR", "/data/gradio")
 os.environ.setdefault("GRADIO_CACHE_DIR", "/data/gradio")
 
 for p in [
-    # no transformers-specific cache path needed; HF will use HF_HOME
     "/data/.cache/huggingface/hub",
     "/data/gradio",
 ]:
@@ -38,13 +37,20 @@ except Exception:
     _HAS_COHERE = False
 
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import login, HfApi
+from huggingface_hub import login
 
 # -------------------
-# NEW: Safety imports (from your snippet / safety.py)
+# NEW: Safety imports
 # -------------------
 from safety import safety_filter, refusal_reply
 
+# -------------------
+# NEW: Augmentation imports
+# -------------------
+from retriever import init_retriever, retrieve_context
+from decision_math import compute_operational_numbers
+from prompt_templates import build_system_preamble
+
 # -------------------
 # Config
 # -------------------
@@ -87,18 +93,13 @@ def is_identity_query(message, history):
     return False
 
 def _iter_user_assistant(history):
-    """Yield (user, assistant) pairs from a Gradio history list.
-    Safely handles items that are lists/tuples with >2 elements.
-    """
     for item in (history or []):
         if isinstance(item, (list, tuple)):
             u = item[0] if len(item) > 0 else ""
             a = item[1] if len(item) > 1 else ""
             yield u, a
-    # If dicts ever appear, extend handling here.
 
 def _history_to_prompt(message, history):
-    """Build a simple text prompt for the stable cohere.chat API."""
     parts = []
     for u, a in _iter_user_assistant(history):
         if u:
@@ -193,7 +194,34 @@ def local_generate(model, tokenizer, input_ids, max_new_tokens=350):
     return tokenizer.decode(gen_only, skip_special_tokens=True).strip()
 
 # -------------------
-# Chat Function (with Safety layer)
+# Snapshot Loader
+# -------------------
+def _load_snapshot(path="snapshots/current.json"):
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            return json.load(f)
+    except Exception:
+        return {
+            "timestamp": None,
+            "beds_total": 400,
+            "staffed_ratio": 1.0,
+            "occupied_pct": 0.97,
+            "ed_census": 62,
+            "ed_admits_waiting": 19,
+            "avg_ed_wait_hours": 8,
+            "discharge_ready_today": 11,
+            "discharge_barriers": {"allied_health": 7, "placement": 4},
+            "rn_shortfall": {"med_ward_A": 1, "med_ward_B": 1},
+            "forecast_admits_next_24h": {"respiratory": 14, "other": 9},
+            "isolation_needs_waiting": {"contact": 3, "airborne": 1},
+            "telemetry_needed_waiting": 5
+        }
+
+# Init retriever once
+init_retriever()
+
+# -------------------
+# Chat Function (with Augmentation + Safety)
 # -------------------
 def chat_fn(message, history, user_tz):
     try:
@@ -202,16 +230,31 @@ def chat_fn(message, history, user_tz):
         if blocked_in:
             return refusal_reply(reason_in)
 
-        # Identity short-circuit (use sanitized input)
+        # Identity short-circuit
         if is_identity_query(safe_in, history):
             return "I am ClarityOps, your strategic decision making AI partner."
 
-        # ---- GENERATION using sanitized input ----
+        # --- Load snapshot + policies + numbers
+        snapshot = _load_snapshot()
+        policy_context = retrieve_context(
+            "bed management huddle discharge acceleration bed leveling ambulance offload"
+        )
+        computed = compute_operational_numbers(snapshot)
+        system_preamble = build_system_preamble(snapshot, policy_context, computed)
+
+        # Augmented input
+        augmented_user = (
+            system_preamble
+            + "\n\nUser question:\n"
+            + safe_in
+        )
+
+        # ---- GENERATION ----
         if USE_HOSTED_COHERE:
-            out = cohere_chat(safe_in, history)
+            out = cohere_chat(augmented_user, history)
         else:
             model, tokenizer = load_local_model()
-            inputs = build_inputs(tokenizer, safe_in, history)
+            inputs = build_inputs(tokenizer, augmented_user, history)
             out = local_generate(model, tokenizer, inputs, max_new_tokens=350)
 
         # Tidy echoes
@@ -239,28 +282,24 @@ theme = gr.themes.Soft(
 
 custom_css = """
 :root {
-  --brand-bg: #e6f7f8; /* soft medical teal */
-  --brand-accent: #0d9488; /* teal-600 */
+  --brand-bg: #e6f7f8;
+  --brand-accent: #0d9488;
   --brand-text: #0f172a;
   --brand-text-light: #ffffff;
 }
 
-/* Page background */
 .gradio-container { background: var(--brand-bg); }
 
-/* Title */
 h1 {
   color: var(--brand-text);
   font-weight: 700;
   font-size: 28px !important;
 }
 
-/* Hide default Chatbot label (cover most Gradio builds) */
 .chatbot header, .chatbot .label, .chatbot .label-wrap, .chatbot .top, .chatbot .header, .chatbot > .wrap > header {
   display: none !important;
 }
 
-/* Bubble styling */
 .message.user, .message.bot {
   background: var(--brand-accent) !important;
   color: var(--brand-text-light) !important;
@@ -268,10 +307,8 @@ h1 {
   padding: 8px 12px !important;
 }
 
-/* Inputs a bit softer */
 textarea, input, .gr-input { border-radius: 12px !important; }
 
-/* Center examples */
 .examples, .examples .grid {
   display: flex !important;
   justify-content: center !important;
@@ -283,7 +320,6 @@ textarea, input, .gr-input { border-radius: 12px !important; }
 # UI
 # -------------------
 with gr.Blocks(theme=theme, css=custom_css) as demo:
-    # Hidden box to carry timezone (still useful for future features)
     tz_box = gr.Textbox(visible=False)
     demo.load(
         lambda tz: tz,
@@ -292,7 +328,6 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
         js="() => Intl.DateTimeFormat().resolvedOptions().timeZone",
     )
 
-    # Extra JS hard-removal of the Chatbot label to cover all DOM variants
     hide_label_sink = gr.HTML(visible=False)
     demo.load(
         fn=lambda: "",
@@ -314,10 +349,8 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
         """,
     )
 
-    # Title
     gr.Markdown("# ClarityOps Augmented Decision AI")
 
-    # Chat interface (larger chat, no Undo; examples centered & single-column)
     gr.ChatInterface(
         fn=chat_fn,
         type="messages",
@@ -326,7 +359,7 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
             label="",
             show_label=False,
             type="messages",
-            height=700,  # larger chat window
+            height=700,
         ),
         examples=[
             ["What are the symptoms of hypertension?"],
@@ -337,17 +370,17 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
         submit_btn="Submit",
         retry_btn="Retry",
         clear_btn="Clear",
-        undo_btn=None,  # removed Undo button
+        undo_btn=None,
     )
 
 if __name__ == "__main__":
-    # Hugging Face Spaces expects the app to listen on $PORT and 0.0.0.0
     port = int(os.environ.get("PORT", "7860"))
     demo.launch(
         server_name="0.0.0.0",
         server_port=port,
-        show_api=False,  # optional: less overhead
-        max_threads=8,   # optional: avoid thread-starvation on tiny CPUs
+        show_api=False,
+        max_threads=8,
     )
 
 
+