Update app.py

app.py

@@ -1,23 +1,28 @@
-\
 import os, re, json
 from functools import lru_cache
 
 import gradio as gr
 import torch
 
+# ---------- Env/cache (quiet deprecation) ----------
 os.environ.setdefault("HF_HOME", "/data/.cache/huggingface")
 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")
+os.environ.pop("TRANSFORMERS_CACHE", None)  # silence v5 deprecation note
 for p in ["/data/.cache/huggingface/hub", "/data/gradio"]:
-    try: os.makedirs(p, exist_ok=True)
-    except Exception: pass
+    try:
+        os.makedirs(p, exist_ok=True)
+    except Exception:
+        pass
 
+# ---------- Optional timezone ----------
 try:
-    from zoneinfo import ZoneInfo
+    from zoneinfo import ZoneInfo  # noqa: F401
 except Exception:
-    ZoneInfo = None
+    ZoneInfo = None  # noqa: N816
 
+# ---------- Optional Cohere ----------
 try:
     import cohere
     _HAS_COHERE = True
@@ -27,6 +32,7 @@ except Exception:
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from huggingface_hub import login
 
+# ---------- ClarityOps modules ----------
 from safety import safety_filter, refusal_reply
 from retriever import init_retriever, retrieve_context
 from decision_math import compute_operational_numbers
@@ -35,27 +41,40 @@ from upload_ingest import extract_text_from_files
 from session_rag import SessionRAG
 from mdsi_analysis import capacity_projection, cost_estimate, outcomes_summary
 
+# ---------- Config ----------
 MODEL_ID = os.getenv("MODEL_ID", "CohereLabs/c4ai-command-r7b-12-2024")
 HF_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN") or os.getenv("HF_TOKEN")
 COHERE_API_KEY = os.getenv("COHERE_API_KEY")
 USE_HOSTED_COHERE = bool(COHERE_API_KEY and _HAS_COHERE)
 
+# ---------- Helpers ----------
 def pick_dtype_and_map():
-    if torch.cuda.is_available(): return torch.float16, "auto"
-    if torch.backends.mps.is_available(): return torch.float16, {"": "mps"}
+    if torch.cuda.is_available():
+        return torch.float16, "auto"
+    if torch.backends.mps.is_available():
+        return torch.float16, {"": "mps"}
     return torch.float32, "cpu"
 
 def is_identity_query(message, history):
     patterns = [
-        r"\bwho\s+are\s+you\b", r"\bwhat\s+are\s+you\b", r"\bwhat\s+is\s+your\s+name\b",
-        r"\bwho\s+is\s+this\b", r"\bidentify\s+yourself\b", r"\btell\s+me\s+about\s+yourself\b",
-        r"\bdescribe\s+yourself\b", r"\band\s+you\s*\?\b", r"\byour\s+name\b", r"\bwho\s+am\s+i\s+chatting\s+with\b"
+        r"\bwho\s+are\s+you\b",
+        r"\bwhat\s+are\s+you\b",
+        r"\bwhat\s+is\s+your\s+name\b",
+        r"\bwho\s+is\s+this\b",
+        r"\bidentify\s+yourself\b",
+        r"\btell\s+me\s+about\s+yourself\b",
+        r"\bdescribe\s+yourself\b",
+        r"\band\s+you\s*\?\b",
+        r"\byour\s+name\b",
+        r"\bwho\s+am\s+i\s+chatting\s+with\b",
     ]
     def match(t): return any(re.search(p, (t or "").strip().lower()) for p in patterns)
-    if match(message): return True
+    if match(message):
+        return True
     if history:
         last_user = history[-1][0] if isinstance(history[-1], (list, tuple)) else None
-        if match(last_user): return True
+        if match(last_user):
+            return True
     return False
 
 def _iter_user_assistant(history):
@@ -74,6 +93,7 @@ def _history_to_prompt(message, history):
     parts.append("Assistant:")
     return "\n".join(parts)
 
+# ---------- Cohere path ----------
 _co_client = None
 if USE_HOSTED_COHERE:
     _co_client = cohere.Client(api_key=COHERE_API_KEY)
@@ -85,7 +105,7 @@ def cohere_chat(message, history):
             model="command-r7b-12-2024",
             message=prompt,
             temperature=0.3,
-            max_tokens=700,
+            max_tokens=900,
         )
         if hasattr(resp, "text") and resp.text: return resp.text.strip()
         if hasattr(resp, "reply") and resp.reply: return resp.reply.strip()
@@ -94,9 +114,11 @@ def cohere_chat(message, history):
     except Exception as e:
         return f"Error calling Cohere API: {e}"
 
+# ---------- Local model ----------
 @lru_cache(maxsize=1)
 def load_local_model():
-    if not HF_TOKEN: raise RuntimeError("HUGGINGFACE_HUB_TOKEN is not set.")
+    if not HF_TOKEN:
+        raise RuntimeError("HUGGINGFACE_HUB_TOKEN is not set.")
     login(token=HF_TOKEN, add_to_git_credential=False)
     dtype, device_map = pick_dtype_and_map()
     tok = AutoTokenizer.from_pretrained(
@@ -116,23 +138,30 @@ def build_inputs(tokenizer, message, history):
         if u: msgs.append({"role": "user", "content": u})
         if a: msgs.append({"role": "assistant", "content": a})
     msgs.append({"role": "user", "content": message})
-    return tokenizer.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True, return_tensors="pt")
+    return tokenizer.apply_chat_template(
+        msgs, tokenize=True, add_generation_prompt=True, return_tensors="pt"
+    )
 
 def local_generate(model, tokenizer, input_ids, max_new_tokens=900):
     input_ids = input_ids.to(model.device)
     with torch.no_grad():
         out = model.generate(
-            input_ids=input_ids, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.3, top_p=0.9,
-            repetition_penalty=1.15, pad_token_id=tokenizer.eos_token_id, eos_token_id=tokenizer.eos_token_id,
+            input_ids=input_ids, max_new_tokens=max_new_tokens,
+            do_sample=True, temperature=0.3, top_p=0.9,
+            repetition_penalty=1.15,
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id,
         )
     gen_only = out[0, input_ids.shape[-1]:]
     return tokenizer.decode(gen_only, skip_special_tokens=True).strip()
 
+# ---------- 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:
+        # Safe fallback if no snapshot present
         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,
@@ -142,11 +171,12 @@ def _load_snapshot(path="snapshots/current.json"):
             "isolation_needs_waiting": {"contact": 3, "airborne": 1}, "telemetry_needed_waiting": 5
         }
 
-# Init retriever & session RAG
+# ---------- Init retrieval engines ----------
 init_retriever()
-_session_rag = SessionRAG()
+_session_rag = SessionRAG()  # ephemeral per-session index for uploaded docs/images
 
-def _mdsi_block() -> str:
+# ---------- Executive pre-compute (MDSi block) ----------
+def _mdsi_block():
     base_capacity = capacity_projection(18, 48, 6)
     cons_capacity = capacity_projection(12, 48, 6)
     opt_capacity = capacity_projection(24, 48, 6)
@@ -160,42 +190,58 @@ def _mdsi_block() -> str:
         "outcomes_summary": outcomes
     }, indent=2)
 
-def chat_fn(message, history, user_tz, uploaded_files, scenario_text):
+# ---------- Core chat logic ----------
+def clarityops_reply(user_msg, history, tz, uploaded_files_paths):
+    """
+    - user_msg: latest message text
+    - history: list[(user, assistant)]
+    - tz: timezone str (unused but kept for future features)
+    - uploaded_files_paths: list[str] absolute paths of uploaded files
+    """
     try:
-        safe_in, blocked_in, reason_in = safety_filter(message, mode="input")
-        if blocked_in: return refusal_reply(reason_in)
-        if is_identity_query(safe_in, history):
-            return "I am ClarityOps, your strategic decision making AI partner."
-
-        # Ingest uploads
-        filepaths = [f.name if hasattr(f, "name") else f for f in (uploaded_files or [])]
-        if filepaths:
-            items = extract_text_from_files(filepaths)
-            if items: _session_rag.add_docs(items)
+        # Safety (input)
+        safe_in, blocked_in, reason_in = safety_filter(user_msg, mode="input")
+        if blocked_in:
+            return history + [(user_msg, refusal_reply(reason_in))]
 
-        # Retrieve snippets from session uploads
-        session_snips = "\\n---\\n".join(_session_rag.retrieve(
-            "diabetes screening Indigenous Métis mobile program cost throughput outcomes logistics", k=6
+        # Identity short-circuit
+        if is_identity_query(safe_in, history):
+            return history + [(user_msg, "I am ClarityOps, your strategic decision making AI partner.")]
+
+        # Ingest new uploads into session RAG (ephemeral for this chat)
+        if uploaded_files_paths:
+            items = extract_text_from_files(uploaded_files_paths)
+            if items:
+                _session_rag.add_docs(items)
+
+        # Pull session snippets from uploaded docs/images
+        session_snips = "\n---\n".join(_session_rag.retrieve(
+            "diabetes screening Indigenous Métis mobile program cost throughput outcomes logistics bed flow staffing discharge forecast",
+            k=6
         ))
 
+        # Load daily snapshot + policies + computed ops numbers
         snapshot = _load_snapshot()
         policy_context = retrieve_context(
-            "mobile diabetes screening Indigenous community outreach logistics referral pathways privacy cultural safety data governance cost effectiveness outcomes"
+            "mobile diabetes screening Indigenous community outreach logistics referral pathways cultural safety data governance cost effectiveness outcomes bed management discharge acceleration ambulance offload"
         )
         computed = compute_operational_numbers(snapshot)
 
-        mdsi_extra = _mdsi_block() if ("diabetes" in (scenario_text or "").lower() or "mdsi" in (scenario_text or "").lower()) else ""
+        # Smart scenario detection: if user message itself looks like exec MDSi context, include the pre-compute block
+        user_lower = (safe_in or "").lower()
+        mdsi_extra = _mdsi_block() if ("diabetes" in user_lower or "mdsi" in user_lower or "mobile screening" in user_lower) else ""
 
         system_preamble = build_system_preamble(
             snapshot=snapshot,
             policy_context=policy_context,
             computed_numbers=computed,
-            scenario_text=(scenario_text or "" ) + (f"\\n\\nExecutive Pre-Computed Blocks:\\n{mdsi_extra}" if mdsi_extra else ""),
+            scenario_text=(safe_in if len(safe_in) > 400 else "") + (f"\n\nExecutive Pre-Computed Blocks:\n{mdsi_extra}" if mdsi_extra else ""),
             session_snips=session_snips
         )
 
-        augmented_user = system_preamble + "\\n\\nUser question or request:\\n" + safe_in
+        augmented_user = system_preamble + "\n\nUser question or request:\n" + safe_in
 
+        # Generate
         if USE_HOSTED_COHERE:
             out = cohere_chat(augmented_user, history)
         else:
@@ -203,56 +249,143 @@ def chat_fn(message, history, user_tz, uploaded_files, scenario_text):
             inputs = build_inputs(tokenizer, augmented_user, history)
             out = local_generate(model, tokenizer, inputs, max_new_tokens=900)
 
+        # Tidy echoes
         if isinstance(out, str):
             for tag in ("Assistant:", "System:", "User:"):
-                if out.startswith(tag): out = out[len(tag):].strip()
+                if out.startswith(tag):
+                    out = out[len(tag):].strip()
 
+        # Safety (output)
         safe_out, blocked_out, reason_out = safety_filter(out, mode="output")
-        if blocked_out: return refusal_reply(reason_out)
-        return safe_out
+        if blocked_out:
+            out = refusal_reply(reason_out)
+
+        return history + [(user_msg, safe_out)]
     except Exception as e:
-        return f"Error: {e}"
+        return history + [(user_msg, f"Error: {e}")]
 
+# ---------- Theme & CSS ----------
 theme = gr.themes.Soft(primary_hue="teal", neutral_hue="slate", radius_size=gr.themes.sizes.radius_lg)
 custom_css = """
 :root { --brand-bg: #e6f7f8; --brand-accent: #0d9488; --brand-text: #0f172a; --brand-text-light: #ffffff; }
 .gradio-container { background: var(--brand-bg); }
+
+/* Title */
 h1 { color: var(--brand-text); font-weight: 700; font-size: 28px !important; }
-.chatbot header, .chatbot .label, .chatbot .label-wrap, .chatbot .top, .chatbot .header, .chatbot > .wrap > header { display: none !important; }
-.message.user, .message.bot { background: var(--brand-accent) !important; color: var(--brand-text-light) !important; border-radius: 12px !important; padding: 8px 12px !important; }
+
+/* Hide default Chatbot label */
+.chatbot header, .chatbot .label, .chatbot .label-wrap, .chatbot .top, .chatbot .header, .chatbot > .wrap > header {
+  display: none !important;
+}
+
+/* Chat bubbles */
+.message.user, .message.bot {
+  background: var(--brand-accent) !important;
+  color: var(--brand-text-light) !important;
+  border-radius: 12px !important;
+  padding: 8px 12px !important;
+}
+
+/* Inputs softer */
 textarea, input, .gr-input { border-radius: 12px !important; }
-.examples, .examples .grid { display: flex !important; justify-content: center !important; text-align: center !important; }
 """
 
+# ---------- UI (single integrated window; uploads at bottom) ----------
 with gr.Blocks(theme=theme, css=custom_css) as demo:
+    # timezone capture (hidden)
     tz_box = gr.Textbox(visible=False)
-    demo.load(lambda tz: tz, inputs=[tz_box], outputs=[tz_box],
-              js="() => Intl.DateTimeFormat().resolvedOptions().timeZone")
+    demo.load(
+        lambda tz: tz,
+        inputs=[tz_box],
+        outputs=[tz_box],
+        js="() => Intl.DateTimeFormat().resolvedOptions().timeZone",
+    )
 
+    # extra DOM cleanup for some gradio builds
     hide_label_sink = gr.HTML(visible=False)
-    demo.load(fn=lambda: "", inputs=None, outputs=hide_label_sink, js="""
-        () => { const sel = ['.chatbot header','.chatbot .label','.chatbot .label-wrap','.chatbot .top','.chatbot .header','.chatbot > .wrap > header'];
-                sel.forEach(s => document.querySelectorAll(s).forEach(el => el.style.display = 'none')); return ""; } """)
+    demo.load(
+        fn=lambda: "",
+        inputs=None,
+        outputs=hide_label_sink,
+        js="""
+        () => {
+          const sel = [
+            '.chatbot header','.chatbot .label','.chatbot .label-wrap',
+            '.chatbot .top','.chatbot .header','.chatbot > .wrap > header'
+          ];
+          sel.forEach(s => document.querySelectorAll(s).forEach(el => el.style.display = 'none'));
+          return "";
+        }
+        """,
+    )
 
     gr.Markdown("# ClarityOps Augmented Decision AI")
 
-    uploads = gr.Files(label="Upload docs/images (PDF, DOCX, CSV, PNG, JPG)", file_types=["file"], file_count="multiple")
-    scenario = gr.Textbox(label="Scenario Context (paste case studies or executive briefs here)",
-                          lines=10, placeholder="Paste scenario text...")
-
-    gr.ChatInterface(
-        fn=chat_fn,
-        type="messages",
-        additional_inputs=[tz_box, uploads, scenario],
-        chatbot=gr.Chatbot(label="", show_label=False, type="messages", height=700),
-        examples=[
-            ["What are the symptoms of hypertension?"],
-            ["What are common drug interactions with aspirin?"],
-            ["What are the warning signs of diabetes?"],
-        ],
-        cache_examples=False,
+    # Main chat area
+    chat = gr.Chatbot(label="", show_label=False, type="messages", height=700)
+
+    # ---- Bottom bar: uploads + message box + send/clear ----
+    with gr.Row():
+        uploads = gr.Files(
+            label="Upload docs/images (PDF, DOCX, CSV, PNG, JPG)",
+            file_types=["file"],
+            file_count="multiple",
+            # keep compact footprint
+            height=68
+        )
+
+    with gr.Row():
+        msg = gr.Textbox(placeholder="Type a message… (paste scenarios here too; ClarityOps will adapt)", scale=10)
+        send = gr.Button("Send", scale=1)
+        clear = gr.Button("Clear chat", scale=1)
+
+    # States
+    state_history = gr.State(value=[])
+    state_uploaded = gr.State(value=[])
+
+    # When user selects files, store their paths in state (so they persist across turns)
+    def _store_uploads(files, current):
+        paths = []
+        for f in (files or []):
+            # gradio Files returns tempfile objects with .name
+            paths.append(getattr(f, "name", None) or f)
+        return (current or []) + paths
+
+    uploads.change(fn=_store_uploads, inputs=[uploads, state_uploaded], outputs=state_uploaded)
+
+    # Send message -> compute reply -> update chat
+    def _on_send(user_msg, history, tz, up_paths):
+        if not user_msg or not user_msg.strip():
+            return history, ""  # no-op
+        new_history = clarityops_reply(user_msg.strip(), history or [], tz, up_paths or [])
+        return new_history, ""
+
+    send.click(
+        fn=_on_send,
+        inputs=[msg, state_history, tz_box, state_uploaded],
+        outputs=[chat, msg],
+        queue=True,
     )
 
+    # Also allow pressing Enter inside the textbox
+    msg.submit(
+        fn=_on_send,
+        inputs=[msg, state_history, tz_box, state_uploaded],
+        outputs=[chat, msg],
+        queue=True,
+    )
+
+    # Keep Chatbot history state in sync whenever it updates
+    chat.change(lambda h: h, inputs=chat, outputs=state_history)
+
+    # Clear chat (keeps uploads so you can keep referencing docs)
+    def _clear_chat():
+        return [], []
+    clear.click(lambda: [], None, chat)
+    # If you also want to clear uploads, uncomment below:
+    # clear.click(_clear_chat, None, [chat, state_uploaded])
+
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", "7860"))
     demo.launch(server_name="0.0.0.0", server_port=port, show_api=False, max_threads=8)
+