Update app.py

app.py

@@ -1,4 +1,4 @@
-import os, re, json, threading, traceback
+import os, re, json, traceback
 from functools import lru_cache
 
 import gradio as gr
@@ -8,7 +8,7 @@ from settings import SNAPSHOT_PATH, PERSIST_CONTENT
 from audit_log import log_event, hash_summary
 from privacy import redact_text
 
-# ---------- Env/cache ----------
+# ---------- Environment / cache ----------
 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")
@@ -20,12 +20,6 @@ for p in ["/data/.cache/huggingface/hub", "/data/gradio"]:
     except Exception:
         pass
 
-# Optional timezone
-try:
-    from zoneinfo import ZoneInfo  # noqa: F401
-except Exception:
-    ZoneInfo = None  # noqa: N816
-
 # Optional Cohere
 try:
     import cohere
@@ -75,16 +69,13 @@ def is_identity_query(message, history):
         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):
-    # history is a list of (user, assistant) tuples (Chatbot default format)
     for item in (history or []):
         if isinstance(item, (list, tuple)):
             u = item[0] if len(item) > 0 else ""
@@ -100,32 +91,23 @@ def _history_to_prompt(message, history):
     parts.append("Assistant:")
     return "\n".join(parts)
 
-# ---------- Cohere path (default first; fallback to local on failure) ----------
-_co_client = None
-if USE_HOSTED_COHERE:
-    # Avoid passing unsupported args; some SDK builds don't accept timeout=
-    _co_client = cohere.Client(api_key=COHERE_API_KEY)
-
+# ---------- Cohere (default path) ----------
 def cohere_chat(message, history):
-    """
-    Returns text on success, or None to signal fallback to local model.
-    """
-    if not _co_client:
+    if not USE_HOSTED_COHERE:
         return None
     try:
+        # Create client on demand to avoid init errors on some builds
+        client = cohere.Client(api_key=COHERE_API_KEY)
         prompt = _history_to_prompt(message, history)
-        resp = _co_client.chat(
+        resp = client.chat(
             model="command-r7b-12-2024",
             message=prompt,
             temperature=0.3,
             max_tokens=MAX_NEW_TOKENS,
         )
-        if hasattr(resp, "text") and resp.text:
-            return resp.text.strip()
-        if hasattr(resp, "reply") and resp.reply:
-            return resp.reply.strip()
-        if hasattr(resp, "generations") and resp.generations:
-            return resp.generations[0].text.strip()
+        if hasattr(resp, "text") and resp.text: return resp.text.strip()
+        if hasattr(resp, "reply") and resp.reply: return resp.reply.strip()
+        if hasattr(resp, "generations") and resp.generations: return resp.generations[0].text.strip()
         return None
     except Exception:
         return None
@@ -141,7 +123,6 @@ def load_local_model():
         MODEL_ID, token=HF_TOKEN, use_fast=True, model_max_length=8192,
         padding_side="left", trust_remote_code=True,
     )
-    # Try device_map (needs accelerate); fallback to manual .to(device) if it fails.
     try:
         mdl = AutoModelForCausalLM.from_pretrained(
             MODEL_ID, token=HF_TOKEN, device_map=device_map,
@@ -158,7 +139,6 @@ def load_local_model():
     return mdl, tok
 
 def build_inputs(tokenizer, message, history):
-    # Convert tuple history to chat template input for HF models
     msgs = []
     for u, a in _iter_user_assistant(history):
         if u: msgs.append({"role": "user", "content": u})
@@ -187,7 +167,6 @@ def _load_snapshot(path=SNAPSHOT_PATH):
         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,
@@ -199,7 +178,7 @@ def _load_snapshot(path=SNAPSHOT_PATH):
 
 # ---------- Init retrieval engines ----------
 init_retriever()
-_session_rag = SessionRAG()  # in-memory only; lazy-loads embeddings
+_session_rag = SessionRAG()  # in-memory only; embeddings load lazily upon first use
 
 # ---------- Executive pre-compute (MDSi block) ----------
 def _mdsi_block():
@@ -216,48 +195,35 @@ def _mdsi_block():
 
 # ---------- Core chat logic (Cohere-first with fallback) ----------
 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:
-        # Audit (content-free)
         log_event("user_message", None, {"sizes": {"chars": len(user_msg or "")}})
 
-        # Safety (input)
         safe_in, blocked_in, reason_in = safety_filter(user_msg, mode="input")
         if blocked_in:
             ans = refusal_reply(reason_in)
             return history + [(user_msg, ans)]
 
-        # Identity short-circuit
         if is_identity_query(safe_in, history):
             ans = "I am ClarityOps, your strategic decision making AI partner."
             return history + [(user_msg, ans)]
 
-        # Ingest uploads (PHI-redacted in upload_ingest)
         if uploaded_files_paths:
             items = extract_text_from_files(uploaded_files_paths)
             if items:
                 _session_rag.add_docs(items)
                 log_event("uploads_added", None, {"count": len(items)})
 
-        # Retrieve from session uploads
         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 cultural safety data governance cost effectiveness outcomes bed management discharge acceleration ambulance offload"
         )
         computed = compute_operational_numbers(snapshot)
 
-        # Exec scenario detect (MDSi)
         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 ""
 
@@ -272,29 +238,24 @@ def clarityops_reply(user_msg, history, tz, uploaded_files_paths):
 
         augmented_user = system_preamble + "\n\nUser question or request:\n" + safe_in
 
-        # --- Cohere first ---
-        out = None
-        if USE_HOSTED_COHERE:
-            out = cohere_chat(augmented_user, history)
+        # Cohere first
+        out = cohere_chat(augmented_user, history)
 
-        # --- Fallback to local HF model if Cohere not set or fails ---
+        # Fallback to local HF model if Cohere not set or failed
         if not out:
             model, tokenizer = load_local_model()
             inputs = build_inputs(tokenizer, augmented_user, history)
             out = local_generate(model, tokenizer, inputs, max_new_tokens=MAX_NEW_TOKENS)
 
-        # Tidy echoes
         if isinstance(out, str):
             for tag in ("Assistant:", "System:", "User:"):
                 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:
             safe_out = refusal_reply(reason_out)
 
-        # Audit (content-free fingerprints)
         log_event("assistant_reply", None, {
             **hash_summary("prompt", augmented_user if not PERSIST_CONTENT else ""),
             **hash_summary("reply", safe_out if not PERSIST_CONTENT else ""),
@@ -302,7 +263,6 @@ def clarityops_reply(user_msg, history, tz, uploaded_files_paths):
 
         return history + [(user_msg, safe_out)]
     except Exception as e:
-        # Surface the error in-chat so the websocket doesn’t die silently
         err = f"Error: {e}"
         try:
             traceback.print_exc()
@@ -315,54 +275,18 @@ theme = gr.themes.Soft(primary_hue="teal", neutral_hue="slate", radius_size=gr.t
 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; }
-
-/* 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 */
+.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; }
 textarea, input, .gr-input { border-radius: 12px !important; }
 """
 
-# ---------- UI ----------
-with gr.Blocks(theme=theme, css=custom_css) as demo:
-    tz_box = gr.Textbox(visible=False)
-    demo.load(
-        lambda tz: tz,
-        inputs=[tz_box],
-        outputs=[tz_box],
-        js="() => Intl.DateTimeFormat().resolvedOptions().timeZone",
-    )
-
-    # --- Background warmup so first message doesn't time out ---
-    def _warmup():
-        # IMPORTANT: no return value, because we register with outputs=None
-        def _bg():
-            try:
-                load_local_model()  # Preload local fallback quietly
-            except Exception:
-                pass
-        threading.Thread(target=_bg, daemon=True).start()
-    demo.load(_warmup)  # no inputs, no outputs
-
+# ---------- UI (single window; uploads at bottom) ----------
+with gr.Blocks(theme=theme, css=custom_css, analytics_enabled=False) as demo:
     gr.Markdown("# ClarityOps Augmented Decision AI")
 
-    # Main chat (tuple-format history)
     chat = gr.Chatbot(label="", show_label=False, height=700)
 
-    # Uploads above the input
     with gr.Row():
         uploads = gr.Files(
             label="Upload docs/images (PDF, DOCX, CSV, PNG, JPG)",
@@ -379,11 +303,9 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
         send = gr.Button("Send", scale=1)
         clear = gr.Button("Clear chat", scale=1)
 
-    # State
     state_history = gr.State(value=[])
     state_uploaded = gr.State(value=[])
 
-    # Store uploaded file paths in state (persist through session)
     def _store_uploads(files, current):
         paths = []
         for f in (files or []):
@@ -392,12 +314,11 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
 
     uploads.change(fn=_store_uploads, inputs=[uploads, state_uploaded], outputs=state_uploaded)
 
-    # Send / Enter handlers (defensive wrapper)
-    def _on_send(user_msg, history, tz, up_paths):
+    def _on_send(user_msg, history, up_paths):
         try:
             if not user_msg or not user_msg.strip():
                 return history, "", history
-            new_history = clarityops_reply(user_msg.strip(), history or [], tz, up_paths or [])
+            new_history = clarityops_reply(user_msg.strip(), history or [], None, up_paths or [])
             return new_history, "", new_history
         except Exception as e:
             err = f"Error: {e}"
@@ -405,24 +326,17 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
                 traceback.print_exc()
             except Exception:
                 pass
-            return (history or []) + [(user_msg or "", err)], "", (history or []) + [(user_msg or "", err)]
-
-    send.click(
-        fn=_on_send,
-        inputs=[msg, state_history, tz_box, state_uploaded],
-        outputs=[chat, msg, state_history],
-        concurrency_limit=2,
-        queue=True,
-    )
-    msg.submit(
-        fn=_on_send,
-        inputs=[msg, state_history, tz_box, state_uploaded],
-        outputs=[chat, msg, state_history],
-        concurrency_limit=2,
-        queue=True,
-    )
+            new_hist = (history or []) + [(user_msg or "", err)]
+            return new_hist, "", new_hist
+
+    send.click(_on_send, inputs=[msg, state_history, state_uploaded],
+               outputs=[chat, msg, state_history],
+               concurrency_limit=2, queue=True)
+
+    msg.submit(_on_send, inputs=[msg, state_history, state_uploaded],
+               outputs=[chat, msg, state_history],
+               concurrency_limit=2, queue=True)
 
-    # Clear chat (keep uploads)
     clear.click(lambda: ([], "", []), None, [chat, msg, state_history])
 
 if __name__ == "__main__":