Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py — Biomedical Device Troubleshooting Assistant (Education-only)
+# app.py — Biomedical Device Troubleshooting Assistant (Education-only, Agentic)
 # Self-contained: no external 'agents' package needed.
 
 import os, io, re, json
@@ -74,7 +74,7 @@ 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(4):
+        for _ in range(6):  # allow a few PLAN→ACT→OBSERVE loops
             resp = await agent.model.client.chat.completions.create(
                 model=agent.model.model,
                 messages=msgs,
@@ -140,7 +140,7 @@ llm_model = OpenAIChatCompletionsModel(model=MODEL_ID, openai_client=ext_client)
 # =========================
 # Topic/safety guardrails
 # =========================
-ALLOWED_COMMANDS = ("/help", "/policy", "/manual", "/clear", "/hits")
+ALLOWED_COMMANDS = ("/help", "/policy", "/manual", "/clear", "/hits", "/agent", "/simple")
 
 TOPIC_KEYWORDS = [
     "biomedical","biomed","device","equipment","oem","service manual",
@@ -213,22 +213,34 @@ def manual_hits(pages, query: str, topk: int = 3):
         start = max(0, min([i for i in idxs if i >= 0], default=0) - window)
         end = min(len(t), start + 2*window)
         snippet = re.sub(r"\s+", " ", t[start:end]).strip()
-        # basic PHI redaction
         snippet = RE_EMAIL.sub("[REDACTED_EMAIL]", snippet)
         snippet = RE_PHONE.sub("[REDACTED_PHONE]", snippet)
         snippet = RE_IDHINT.sub("[ID]", snippet)
         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():
+    m = cl.user_session.get("manual")
+    return (m or {}).get("pages", [])
+
 # =========================
-# Tools (deterministic helpers)
+# Tools (agent can call)
 # =========================
 @function_tool
-def quick_reference(device_desc: str) -> dict:
+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()
+    k = int(topk) if str(topk).isdigit() else 3
+    hits = manual_hits(pages, query, topk=k) if pages else []
+    return {"hits": hits}
+
+@function_tool
+def quick_reference(device_and_symptom: str) -> dict:
     """
     Deterministic education-only bullets based on free-text device description.
     """
-    d = (device_desc or "").lower()
+    d = (device_and_symptom or "").lower()
     life_critical = any(k in d for k in ["ventilator","defibrillator","infusion pump"])
     return {
         "safety": [
@@ -261,34 +273,71 @@ def quick_reference(device_desc: 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."
+        ]
+    }
+
+@function_tool
+def record_step(note: str) -> dict:
+    """Append an action/observation to the agent trace shown to the user."""
+    trace: List[str] = cl.user_session.get("agent_trace") or []
+    clean = re.sub(r"\s+", " ", (note or "")).strip()
+    if clean:
+        trace.append(clean[:500])
+        cl.user_session.set("agent_trace", trace)
+    return {"ok": True, "len": len(trace)}
+
 # =========================
-# LLM troubleshooting plan
+# Agentic Planner
 # =========================
-SYSTEM_PROMPT = (
-    "You are a biomedical device troubleshooting assistant for clinical engineers.\n"
-    "STRICT SCOPE: Education-only device troubleshooting. No diagnosis/treatment, no invasive repair, no alarm bypass, no firmware hacks, "
-    "no collection of personal identifiers. Defer to OEM manuals and local policy if any conflict.\n\n"
-    "Given a device description and symptom, produce concise bullet lists with exactly these sections:\n"
-    "1) Safety First (non-invasive, patient-first)\n"
+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"
+    "- 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"
+    "1) Safety First\n"
     "2) Likely Causes (ranked)\n"
-    "3) Step-by-Step Checks (do-not-open device; do-not-bypass alarms)\n"
-    "4) QC/Calibration (what to verify and with what reference/simulator)\n"
-    "5) Escalate When (clear triggers)\n"
-    "End with a one-line summary.\n"
-    "If the text suggests a life-critical device (ventilator/defibrillator/infusion pump) and patient connected, start with: "
-    "REMOVE FROM SERVICE & USE BACKUP — then proceed with safe checks off-patient."
+    "3) Step-by-Step Checks (non-invasive; no alarm bypass)\n"
+    "4) QC/Calibration\n"
+    "5) Escalate When\n"
+    "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,
+    model=llm_model,
+    tools=[search_manual, quick_reference, safety_policy, record_step],
 )
 
-async def call_llm(user_desc: str, manual_excerpts: str = "") -> str:
+# =========================
+# 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."
+)
+
+async def simple_llm_plan(user_desc: str, manual_excerpts: str = "") -> str:
     user_block = f"Device & symptom:\n{user_desc.strip()}"
     if manual_excerpts:
-        user_block += f"\n\nManual excerpts (for reference; OEM prevails if conflict):\n{manual_excerpts.strip()}"
+        user_block += f"\n\nManual excerpts:\n{manual_excerpts.strip()}"
     resp = await ext_client.client.chat.completions.create(
         model=MODEL_ID,
-        messages=[
-            {"role": "system", "content": SYSTEM_PROMPT},
-            {"role": "user", "content": user_block},
-        ],
+        messages=[{"role":"system","content":SIMPLE_SYSTEM},{"role":"user","content":user_block}],
     )
     return resp.choices[0].message.content or ""
 
@@ -296,10 +345,11 @@ async def call_llm(user_desc: str, manual_excerpts: str = "") -> str:
 # UI strings
 # =========================
 WELCOME = (
-    "🛠️ **Biomedical Device Troubleshooting Assistant**\n"
+    "🛠️ **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, **/policy** view rules."
+    "Commands: **/manual** upload PDF/TXT, **/hits** show top manual matches, **/clear** remove manual,\n"
+    "**/agent** agentic mode (default), **/simple** single-pass mode, **/policy** view rules."
 )
 POLICY = (
     "🛡️ **Safety & Scope Policy**\n"
@@ -321,18 +371,26 @@ 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_trace", [])
     await cl.Message(content=WELCOME).send()
 
 @cl.on_message
 async def main(message: cl.Message):
     text = (message.content or "").strip()
+    low = text.lower()
 
     # Commands
-    low = text.lower()
     if low.startswith("/help"):
         await cl.Message(content=WELCOME).send(); return
     if low.startswith("/policy"):
         await cl.Message(content=POLICY).send(); return
+    if low.startswith("/agent"):
+        cl.user_session.set("agent_mode", "agent")
+        await cl.Message(content="✅ Agentic mode enabled (planner will Plan→Act→Observe).").send(); return
+    if low.startswith("/simple"):
+        cl.user_session.set("agent_mode", "simple")
+        await cl.Message(content="✅ Simple mode enabled (single LLM pass).").send(); return
 
     if low.startswith("/manual"):
         files = await cl.AskFileMessage(
@@ -363,8 +421,8 @@ async def main(message: cl.Message):
         m = cl.user_session.get("manual")
         if not m or not m.get("pages"):
             await cl.Message(content="No manual attached. Use **/manual** to upload one.").send(); return
-        hits = manual_hits(m["pages"], text.replace("/hits", "").strip() or "device fault error alarm", topk=3)
-        if not hits: hits = ["No obvious matches — try different wording."]
+        query = text.replace("/hits", "").strip() or "device fault error alarm"
+        hits = manual_hits(m["pages"], query, topk=3) or ["No obvious matches — try different wording."]
         await cl.Message(content="### 🔎 Manual Matches\n" + "\n\n".join(hits)).send()
         return
 
@@ -380,34 +438,49 @@ async def main(message: cl.Message):
         await cl.Message(content=REFUSAL + "\n\n" + POLICY).send()
         return
 
-    # Manual excerpts (optional)
+    # Optional manual context for the simple path
     manual = cl.user_session.get("manual")
+    manual_name = manual.get("name") if manual else None
     excerpts = ""
     if manual and manual.get("pages"):
-        query = text
-        hits = manual_hits(manual["pages"], query, topk=3)
-        excerpts = "\n".join(hits)
+        excerpts = "\n".join(manual_hits(manual["pages"], text, topk=3))
+
+    mode = cl.user_session.get("agent_mode") or "agent"
+    trace_before = len(cl.user_session.get("agent_trace") or [])
 
-    # Deterministic quick reference (for reliability)
-    ref = quick_reference(text)
+    if mode == "agent":
+        # Agentic Planner run
+        try:
+            # give planner the user text; it 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
+    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)
+    qref = quick_reference(text)
     def bullets(arr: List[str]): return "\n".join(f"- {x}" for x in (arr or [])) or "-"
 
     quick_md = (
         f"### 📘 Quick Reference\n"
-        f"**Safety**\n{bullets(ref.get('safety'))}\n\n"
-        f"**Common Faults**\n{bullets(ref.get('common_faults'))}\n\n"
-        f"**Quick Checks**\n{bullets(ref.get('quick_checks'))}\n\n"
-        f"**QC / Calibration**\n{bullets(ref.get('qc'))}\n\n"
-        f"**Escalate If**\n{bullets(ref.get('escalate'))}\n\n"
+        f"**Safety**\n{bullets(qref.get('safety'))}\n\n"
+        f"**Common Faults**\n{bullets(qref.get('common_faults'))}\n\n"
+        f"**Quick Checks**\n{bullets(qref.get('quick_checks'))}\n\n"
+        f"**QC / Calibration**\n{bullets(qref.get('qc'))}\n\n"
+        f"**Escalate If**\n{bullets(qref.get('escalate'))}\n\n"
         f"> ⚠️ Education-only. Refer to OEM manual & policy. No invasive service.\n"
     )
 
-    # LLM troubleshooting plan
-    try:
-        answer = await call_llm(text, excerpts)
-    except Exception as e:
-        await cl.Message(content=f"{quick_md}\n---\n⚠️ LLM call failed: {e}").send()
-        return
-
-    prefix = f"📄 Using manual: **{manual['name']}**\n\n" if (manual and manual.get('name')) else ""
-    await cl.Message(content=f"{quick_md}\n---\n{prefix}{answer or 'I couldn’t generate a plan.'}").send()
+    prefix = f"📄 Using manual: **{manual_name}**\n\n" if manual_name else ""
+    await cl.Message(content=f"{trace_md}{quick_md}\n---\n{prefix}{answer}").send()