Update app.py

app.py

@@ -1,5 +1,6 @@
-# app.py — Biomedical Device Troubleshooting Assistant (Education-only, Agentic)
-# Self-contained: no external 'agents' package needed.
+# app.py — Biomedical Device Troubleshooting Assistant
+# Education-only · Manual (PDF/TXT) upload & search · Guardrails · Agentic planner (optional)
+# Deps: chainlit==1.0.200, python-dotenv==1.0.1, openai==1.35.13, pypdf>=4.2.0
 
 import os, io, re, json
 from typing import Any, Callable, Dict, List, Optional
@@ -11,7 +12,7 @@ from openai import AsyncOpenAI as _SDKAsyncOpenAI
 from pypdf import PdfReader
 
 # =========================
-# Minimal "agents" shim
+# Minimal "agents" shim (no external package needed)
 # =========================
 def set_tracing_disabled(disabled: bool = True):
     return disabled
@@ -20,9 +21,6 @@ def function_tool(func: Callable):
     func._is_tool = True
     return func
 
-class InputGuardrailTripwireTriggered(Exception):
-    pass
-
 class AsyncOpenAI:
     def __init__(self, api_key: str, base_url: Optional[str] = None):
         kwargs = {"api_key": api_key}
@@ -74,7 +72,8 @@ class Runner:
         tools = agent.tool_specs()
         tool_map = {t.__name__: t for t in (agent.tools or []) if getattr(t, "_is_tool", False)}
 
-        for _ in range(6):  # allow a few PLAN→ACT→OBSERVE loops
+        # small Plan→Act→Observe budget
+        for _ in range(6):
             resp = await agent.model.client.chat.completions.create(
                 model=agent.model.model,
                 messages=msgs,
@@ -102,18 +101,17 @@ class Runner:
                     })
                 continue
 
-            # Final
-            result_obj = type("Result", (), {})()
-            result_obj.final_output = msg.content or ""
-            result_obj.context = context or {}
-            result_obj.final_output_as = lambda *_: result_obj.final_output
-            return result_obj
+            out = type("Result", (), {})()
+            out.final_output = msg.content or ""
+            out.context = context or {}
+            out.final_output_as = lambda *_: out.final_output
+            return out
 
-        result_obj = type("Result", (), {})()
-        result_obj.final_output = "Sorry, I couldn't complete the request."
-        result_obj.context = context or {}
-        result_obj.final_output_as = lambda *_: result_obj.final_output
-        return result_obj
+        out = type("Result", (), {})()
+        out.final_output = "Sorry, I couldn't complete the request."
+        out.context = context or {}
+        out.final_output_as = lambda *_: out.final_output
+        return out
 
 # =========================
 # Provider setup (Gemini or OpenAI)
@@ -131,14 +129,14 @@ elif OPENAI_API_KEY:
     BASE_URL = None
     MODEL_ID = "gpt-4o-mini"
 else:
-    raise RuntimeError("Missing GEMINI_API_KEY or OPENAI_API_KEY in your environment.")
+    raise RuntimeError("Missing GEMINI_API_KEY or OPENAI_API_KEY in environment.")
 
 set_tracing_disabled(True)
 ext_client = AsyncOpenAI(api_key=API_KEY, base_url=BASE_URL)
 llm_model = OpenAIChatCompletionsModel(model=MODEL_ID, openai_client=ext_client)
 
 # =========================
-# Topic/safety guardrails
+# Guardrails (topic & safety)
 # =========================
 ALLOWED_COMMANDS = ("/help", "/policy", "/manual", "/clear", "/hits", "/agent", "/simple")
 
@@ -178,7 +176,7 @@ def local_guard(text: str) -> List[str]:
     return issues
 
 # =========================
-# Manual handling + search
+# Manual handling & search (robust PDF detection)
 # =========================
 def extract_pdf_pages(data: bytes):
     pages = []
@@ -194,7 +192,18 @@ def extract_pdf_pages(data: bytes):
 def extract_txt_pages(data: bytes, chunk_chars: int = 1600):
     try: txt = data.decode("utf-8", errors="ignore")
     except Exception: txt = ""
-    return [{"page": i+1, "text": txt[i:i+chunk_chars]} for i in range(0, len(txt), chunk_chars)] or [{"page": 1, "text": ""}]
+    chunks = [{"page": i+1, "text": txt[i:i+chunk_chars]} for i in range(0, len(txt), chunk_chars)]
+    return chunks or [{"page": 1, "text": ""}]
+
+def _is_pdf(file_obj, data: bytes) -> bool:
+    # Chainlit may expose .type or nothing; avoid relying on .mime
+    name = getattr(file_obj, "name", "") or getattr(file_obj, "filename", "")
+    typ  = getattr(file_obj, "type", "") or ""
+    return (
+        name.lower().endswith(".pdf")
+        or typ == "application/pdf"
+        or (isinstance(data, (bytes, bytearray)) and data[:5] == b"%PDF-")
+    )
 
 def manual_hits(pages, query: str, topk: int = 3):
     if not pages: return []
@@ -219,18 +228,17 @@ def manual_hits(pages, query: str, topk: int = 3):
         return snippet
     return [f"[p.{h['page']}] {excerpt(h.get('text',''))}" for h in hits]
 
-# Keep manual in session; agent tools will read it live
-def _current_manual():
+# =========================
+# Tools the agent can call
+# =========================
+def _current_manual_pages():
     m = cl.user_session.get("manual")
     return (m or {}).get("pages", [])
 
-# =========================
-# Tools (agent can call)
-# =========================
 @function_tool
 def search_manual(query: str, topk: str = "3") -> dict:
     """Search the uploaded manual (PDF/TXT) for the query. Returns {'hits': [snippets...]}. If no manual, hits=[]"""
-    pages = _current_manual()
+    pages = _current_manual_pages()
     k = int(topk) if str(topk).isdigit() else 3
     hits = manual_hits(pages, query, topk=k) if pages else []
     return {"hits": hits}
@@ -276,14 +284,12 @@ def quick_reference(device_and_symptom: str) -> dict:
 @function_tool
 def safety_policy() -> dict:
     """Return a concise policy block enforced by the assistant."""
-    return {
-        "policy": [
-            "Education-only troubleshooting; no clinical advice.",
-            "No invasive repairs, no alarm bypass, no firmware tampering.",
-            "No collecting/sharing personal identifiers.",
-            "OEM manuals and local policy take precedence."
-        ]
-    }
+    return {"policy": [
+        "Education-only troubleshooting; no clinical advice.",
+        "No invasive repairs, no alarm bypass, no firmware tampering.",
+        "No collecting/sharing personal identifiers.",
+        "OEM manuals and local policy take precedence."
+    ]}
 
 @function_tool
 def record_step(note: str) -> dict:
@@ -296,14 +302,14 @@ def record_step(note: str) -> dict:
     return {"ok": True, "len": len(trace)}
 
 # =========================
-# Agentic Planner
+# Agentic planner & simple fallback
 # =========================
 PLANNER_INSTRUCTIONS = (
     "You are an agentic biomedical device troubleshooting planner for clinical engineers.\n"
     "STRICT SCOPE: Education-only. Do NOT give clinical advice, do NOT instruct invasive repairs, do NOT bypass alarms, "
     "do NOT suggest firmware tampering. Respect policy.\n\n"
     "Loop up to 3 times:\n"
-    "- PLAN: Decide what you need (e.g., search_manual for specific error, or quick_reference).\n"
+    "- PLAN: Decide what you need (e.g., search_manual for error keywords, or quick_reference).\n"
     "- ACT: Call a tool (search_manual, quick_reference, safety_policy, record_step).\n"
     "- OBSERVE: Read tool result. Summarize via record_step.\n"
     "Then produce a final plan with exactly these sections:\n"
@@ -315,7 +321,6 @@ PLANNER_INSTRUCTIONS = (
     "End with a one-line summary. If life-critical & patient-connected is implied, start with: "
     "REMOVE FROM SERVICE & USE BACKUP — then proceed off-patient.\n"
 )
-
 planner_agent = Agent(
     name="Agentic Planner",
     instructions=PLANNER_INSTRUCTIONS,
@@ -323,9 +328,6 @@ planner_agent = Agent(
     tools=[search_manual, quick_reference, safety_policy, record_step],
 )
 
-# =========================
-# Simple (non-agentic) LLM (fallback)
-# =========================
 SIMPLE_SYSTEM = (
     "You are a biomedical device troubleshooting assistant (education-only). "
     "Output sections: Safety First; Likely Causes; Step-by-Step Checks; QC/Calibration; Escalate When; one-line summary."
@@ -348,7 +350,7 @@ WELCOME = (
     "🛠️ **Biomedical Device Troubleshooting Assistant** (Agentic)\n"
     "Education-only. No diagnosis, no invasive service, no alarm bypass. OEM manual & policy rule the day.\n\n"
     "Type your **device & symptom** (e.g., “Infusion pump occlusion alarm”).\n"
-    "Commands: **/manual** upload PDF/TXT, **/hits** show top manual matches, **/clear** remove manual,\n"
+    "Commands: **/manual** upload PDF/TXT, **/hits** show manual matches, **/clear** remove manual,\n"
     "**/agent** agentic mode (default), **/simple** single-pass mode, **/policy** view rules."
 )
 POLICY = (
@@ -371,7 +373,7 @@ REFUSAL = (
 @cl.on_chat_start
 async def start():
     cl.user_session.set("manual", None)
-    cl.user_session.set("agent_mode", "agent")   # 'agent' or 'simple'
+    cl.user_session.set("agent_mode", "agent")   # 'agent' | 'simple'
     cl.user_session.set("agent_trace", [])
     await cl.Message(content=WELCOME).send()
 
@@ -395,21 +397,30 @@ async def main(message: cl.Message):
     if low.startswith("/manual"):
         files = await cl.AskFileMessage(
             content="Upload the **service manual** (PDF or TXT). Max ~20 MB.",
-            accept=["application/pdf", "text/plain"],
+            accept=["application/pdf", "text/plain"],  # UI hint only
             max_files=1, max_size_mb=20, timeout=240
         ).send()
         if not files:
             await cl.Message(content="No file received.").send(); return
+
         f = files[0]
         data = getattr(f, "content", None)
         if data is None and getattr(f, "path", None):
-            with open(f.path, "rb") as fh: data = fh.read()
+            with open(f.path, "rb") as fh:
+                data = fh.read()
+        if not data:
+            await cl.Message(content="Received file but got empty content. Please re-upload.").send(); return
+
         try:
-            pages = extract_pdf_pages(data) if (f.mime == "application/pdf" or f.name.lower().endswith(".pdf")) else extract_txt_pages(data)
+            if _is_pdf(f, data):
+                pages = extract_pdf_pages(data)
+            else:
+                pages = extract_txt_pages(data)
         except Exception as e:
             await cl.Message(content=f"Couldn't read the manual: {e}").send(); return
-        cl.user_session.set("manual", {"name": f.name, "pages": pages})
-        await cl.Message(content=f"✅ Manual indexed: **{f.name}** — {len(pages)} page-chunks.").send()
+
+        cl.user_session.set("manual", {"name": getattr(f, "name", "manual"), "pages": pages})
+        await cl.Message(content=f"✅ Manual indexed: **{getattr(f, 'name', 'manual')}** — {len(pages)} page-chunks.").send()
         return
 
     if low.startswith("/clear"):
@@ -426,7 +437,7 @@ async def main(message: cl.Message):
         await cl.Message(content="### 🔎 Manual Matches\n" + "\n\n".join(hits)).send()
         return
 
-    # Topic & local safety guard
+    # Topic & safety guard
     if not on_topic(text):
         await cl.Message(
             content="I only support **biomedical device troubleshooting**. "
@@ -438,7 +449,7 @@ async def main(message: cl.Message):
         await cl.Message(content=REFUSAL + "\n\n" + POLICY).send()
         return
 
-    # Optional manual context for the simple path
+    # Manual context (for simple fallback)
     manual = cl.user_session.get("manual")
     manual_name = manual.get("name") if manual else None
     excerpts = ""
@@ -449,26 +460,24 @@ async def main(message: cl.Message):
     trace_before = len(cl.user_session.get("agent_trace") or [])
 
     if mode == "agent":
-        # Agentic Planner run
         try:
-            # give planner the user text; it will call tools search_manual/quick_reference/record_step itself
+            # Agentic planner will call tools (search_manual/quick_reference/record_step) itself
             result = await Runner.run(planner_agent, text)
             answer = result.final_output or "I couldn’t generate a plan."
         except Exception as e:
             await cl.Message(content=f"⚠️ Planner failed: {e}\nFalling back to simple mode.").send()
             answer = await simple_llm_plan(text, excerpts)
     else:
-        # Simple single-pass
         answer = await simple_llm_plan(text, excerpts)
 
-    # Show agent trace if any new steps were recorded
+    # Agent trace (new steps only)
     trace = cl.user_session.get("agent_trace") or []
     new_trace = trace[trace_before:]
     trace_md = ""
     if new_trace:
         trace_md = "### 🧭 Agent Trace (summary)\n" + "\n".join(f"- {t}" for t in new_trace) + "\n\n"
 
-    # Deterministic quick reference (always helpful)
+    # Deterministic quick reference
     qref = quick_reference(text)
     def bullets(arr: List[str]): return "\n".join(f"- {x}" for x in (arr or [])) or "-"