Update app.py

app.py

@@ -1,6 +1,7 @@
 # app.py — Biomedical Device Troubleshooting Assistant (Education-only, Agentic)
-# Minimal deps: chainlit==1.0.200, python-dotenv==1.0.1, openai==1.35.13
-# .env: GEMINI_API_KEY=... (preferred) OR OPENAI_API_KEY=...
+# Deps: chainlit==1.0.200, python-dotenv==1.0.1, openai==1.35.13
+# .env: GEMINI_API_KEY=...  (preferred; uses OpenAI-compatible Gemini endpoint)
+#       OR OPENAI_API_KEY=...
 
 import os, re, json
 from dataclasses import dataclass, field
@@ -11,9 +12,9 @@ from dotenv import load_dotenv
 from openai import AsyncOpenAI as _SDKAsyncOpenAI
 
 # =========================
-# Small “agents” shim
+# Minimal "agents" shim
 # =========================
-def set_tracing_disabled(_: bool = True):  # placeholder (no-op)
+def set_tracing_disabled(_: bool = True):  # no-op placeholder
     return True
 
 def function_tool(func: Callable):
@@ -25,10 +26,10 @@ class AsyncOpenAI:
         kwargs = {"api_key": api_key}
         if base_url:
             kwargs["base_url"] = base_url
-        self._c = _SDKAsyncOpenAI(**kwargs)
+        self._client = _SDKAsyncOpenAI(**kwargs)
     @property
     def client(self):
-        return self._c
+        return self._client
 
 class OpenAIChatCompletionsModel:
     def __init__(self, model: str, openai_client: AsyncOpenAI):
@@ -68,51 +69,87 @@ class Agent:
         return self._tool_specs
 
 class Runner:
+    """Robust tool loop with forced finalization (prevents endless Plan→Act cycles)."""
     @staticmethod
-    async def run(agent: Agent, user_input: str, context: Optional[Dict[str, Any]] = None, turns: int = 5):
+    async def run(
+        agent: Agent,
+        user_input: str,
+        context: Optional[Dict[str, Any]] = None,
+        turns: int = 7,
+        max_tools: int = 3,
+    ):
         msgs = [
             {"role": "system", "content": agent.instructions},
+            {
+                "role": "system",
+                "content": (
+                    "You have access to these tools: "
+                    + ", ".join(t.__name__ for t in (agent.tools or []) if getattr(t, "_is_tool", False))
+                    + f". Use at most {max_tools} tool calls, then produce the final answer in the required sections."
+                ),
+            },
             {"role": "user", "content": user_input},
         ]
+
         tool_specs = agent.tool_specs()
         tool_map = {t.__name__: t for t in (agent.tools or []) if getattr(t, "_is_tool", False)}
+        tool_uses = 0
+
+        for step in range(max(1, min(12, turns))):
+            force_final = (tool_uses >= max_tools) or (step == turns - 1)
 
-        for _ in range(max(1, min(8, turns))):
             resp = await agent.model.client.chat.completions.create(
                 model=agent.model.model,
                 messages=msgs,
-                tools=tool_specs or None,
-                tool_choice="auto" if tool_specs else None,
+                tools=None if force_final else tool_specs,
+                tool_choice="none" if force_final else "auto",
             )
             msg = resp.choices[0].message
-            msgs.append({"role": "assistant", "content": msg.content or "", "tool_calls": msg.tool_calls})
-
-            if msg.tool_calls:
-                # Plan→Act→Observe
-                for call in msg.tool_calls:
-                    fn_name = call.function.name
-                    args = json.loads(call.function.arguments or "{}")
-                    result = {"error": f"Unknown tool: {fn_name}"}
-                    tool = tool_map.get(fn_name)
-                    if tool:
-                        try: result = tool(**args)
-                        except Exception as e: result = {"error": str(e)}
-                    msgs.append({
-                        "role": "tool",
-                        "tool_call_id": call.id,
-                        "name": fn_name,
-                        "content": json.dumps(result),
-                    })
+            tool_calls = getattr(msg, "tool_calls", None)
+
+            # Normal completion (no tool calls) => return
+            if not tool_calls:
+                out = type("Result", (), {})()
+                out.final_output = (msg.content or "").strip()
+                out.context = context or {}
+                out.final_output_as = lambda *_: out.final_output
+                return out
+
+            # If we forced finalization but the model still tried tools, nudge and iterate
+            if force_final:
+                msgs.append({
+                    "role": "system",
+                    "content": "Do not call tools anymore. Produce the final answer now in the required sections."
+                })
                 continue
 
-            out = type("Result", (), {})()
-            out.final_output = msg.content or ""
-            out.context = context or {}
-            out.final_output_as = lambda *_: out.final_output
-            return out
+            # Execute tool calls
+            msgs.append({"role": "assistant", "content": msg.content or "", "tool_calls": tool_calls})
+            for call in tool_calls:
+                fn_name = call.function.name
+                args = json.loads(call.function.arguments or "{}")
+                if fn_name not in tool_map:
+                    result = {"error": f"Unknown tool '{fn_name}'. Available: " + ", ".join(tool_map.keys())}
+                else:
+                    try:
+                        result = tool_map[fn_name](**args)
+                    except Exception as e:
+                        result = {"error": str(e)}
+                msgs.append({
+                    "role": "tool",
+                    "tool_call_id": call.id,
+                    "name": fn_name,
+                    "content": json.dumps(result),
+                })
+                tool_uses += 1
 
+        # One last attempt to finalize with tools disabled
+        msgs.append({"role": "system", "content": "Stop tool use. Provide the final answer now."})
+        resp = await agent.model.client.chat.completions.create(
+            model=agent.model.model, messages=msgs, tools=None, tool_choice="none"
+        )
         out = type("Result", (), {})()
-        out.final_output = "Sorry, I couldn't complete the request."
+        out.final_output = (resp.choices[0].message.content or "I couldn’t generate a plan.").strip()
         out.context = context or {}
         out.final_output_as = lambda *_: out.final_output
         return out
@@ -178,7 +215,7 @@ def _bullets(items: List[str]) -> str:
 @function_tool
 def quick_reference(device_and_symptom: str) -> dict:
     """
-    Generic, non-invasive reference for troubleshooting.
+    Generic, non-invasive reference for troubleshooting (education-only).
     """
     s = (device_and_symptom or "").lower()
     life = any(k in s for k in ["ventilator", "defibrillator", "infusion pump"])
@@ -197,13 +234,13 @@ def quick_reference(device_and_symptom: str) -> dict:
         ],
         "quick_checks": [
             "Note model/serial/firmware and any error codes.",
-            "Verify power source/battery; try another outlet; reseat user-removable battery.",
+            "Verify power or battery; try another outlet; reseat user-removable battery.",
             "Inspect & reseat accessories; swap with a known-good.",
             "Confirm settings; restore defaults if safe.",
-            "Run device self-test and review OEM quick-reference."
+            "Run self-test and review OEM quick-reference."
         ],
         "qc": [
-            "Verify against a reference/simulator where applicable (ECG, SpO₂, NIBP, flow).",
+            "Verify against a reference/simulator (ECG, SpO₂, NIBP, flow) where applicable.",
             "Perform electrical safety tests per policy.",
             "Document & compare with prior QC."
         ],
@@ -216,7 +253,7 @@ def quick_reference(device_and_symptom: str) -> dict:
 
 @function_tool
 def safety_policy() -> dict:
-    """Concise policy enforced by the assistant."""
+    """Concise policy block enforced by the assistant."""
     return {"policy": [
         "Education-only troubleshooting; no clinical advice.",
         "No invasive repairs; no alarm bypass; no firmware tampering.",
@@ -236,13 +273,11 @@ def record_step(note: str) -> dict:
 
 @function_tool
 def error_code_lookup(model: str, code: str) -> dict:
-    """
-    Friendly generic notes for common codes (extendable).
-    """
+    """Friendly generic notes for common codes (extendable)."""
     key = f"{(model or '').strip().lower()}::{(code or '').strip().lower()}"
     KB = {
         "generic::e01": {"meaning": "General startup/self-test failure.",
-                         "notes": ["Safe power-cycle (off-patient). Check mains/battery, connections. Re-run self-test. Escalate if persistent."]},
+                         "notes": ["Safe power-cycle (off-patient). Check mains/battery & connections. Re-run self-test. Escalate if persistent."]},
         "generic::occlusion": {"meaning": "Flow obstruction detected.",
                                "notes": ["Check clamps/kinks/filters. Reload set per OEM. Do not bypass alarms."]},
         "monitor::leads_off": {"meaning": "ECG leads not detected / poor contact.",
@@ -276,7 +311,7 @@ def acronym_expand(term: str) -> dict:
 @function_tool
 def unit_convert(value: str, from_unit: str, to_unit: str) -> dict:
     """
-    Limited, safe conversions used in QC:
+    QC-friendly conversions:
     mmHg↔kPa (1 kPa = 7.50062 mmHg) | L/min↔mL/s (1 L/min = 16.6667 mL/s) | mV↔µV (1 mV = 1000 µV)
     """
     def f(x): return float(str(x).strip())
@@ -300,7 +335,7 @@ def triage_priority(device_and_symptom: str) -> dict:
 
 @function_tool
 def simulator_baselines(device: str) -> dict:
-    """Educational simulator targets (generic, not OEM-specific)."""
+    """Educational simulator targets (generic; not OEM-specific)."""
     d = (device or "").lower()
     bank = {
         "ecg": ["1 mV @ 10 mm/mV (paper 25 mm/s).", "Check lead-fault detection & baseline stability."],
@@ -314,7 +349,7 @@ def simulator_baselines(device: str) -> dict:
 
 @function_tool
 def env_checklist(_: str) -> dict:
-    """Environment/power/EMI checklist for safe troubleshooting."""
+    """Environment/power/EMI checklist (safe)."""
     return {"checks": [
         "Verify mains outlet & ground; power indicator on.",
         "Avoid extension cords; try a known-good outlet.",
@@ -325,7 +360,7 @@ def env_checklist(_: str) -> dict:
 
 @function_tool
 def accessory_checklist(_: str) -> dict:
-    """Generic accessory fit/damage/compatibility checklist."""
+    """Accessory fit/damage/compatibility checklist."""
     return {"checks": [
         "Inspect connectors for bent pins/corrosion; ensure tight seating.",
         "Check cables/tubing for cuts/kinks; confirm model compatibility.",
@@ -335,7 +370,7 @@ def accessory_checklist(_: str) -> dict:
 
 @function_tool
 def build_qc_template(device: str) -> dict:
-    """Markdown QC template students can copy into notes/CMMS (education-only)."""
+    """Markdown QC template for notes/CMMS (education-only)."""
     d = (device or "Device").strip()
     md = f"""### QC/Verification — {d}
 - Date / Tech:
@@ -357,11 +392,8 @@ def build_qc_template(device: str) -> dict:
 PLANNER_INSTRUCTIONS = (
     "You are an **agentic biomedical device troubleshooting planner** for clinical engineers.\n"
     "SCOPE (strict): education-only. No clinical advice. No invasive repairs. No alarm bypass. No firmware tampering. No PHI.\n"
-    "Loop up to 3 times:\n"
-    "- PLAN what you need (e.g., quick_reference, safety_policy, acronym_expand, error_code_lookup, unit_convert, triage_priority,\n"
-    "  simulator_baselines, env_checklist, accessory_checklist, build_qc_template).\n"
-    "- ACT by calling a tool.\n"
-    "- OBSERVE and add a short summary via record_step.\n\n"
+    "Use at most **3 tool calls** and then **produce the final answer with no further tool calls**.\n"
+    "Loop up to 3 times: PLAN → ACT (tool) → OBSERVE (summarize via record_step).\n\n"
     "Finally output exactly these sections:\n"
     "1) Safety First\n"
     "2) Likely Causes (ranked)\n"
@@ -417,7 +449,6 @@ REFUSAL = (
     "🚫 I can’t help with diagnosis/treatment, invasive repair, alarm bypass, firmware hacks, or handling personal identifiers.\n"
     "I can guide safe, non-invasive troubleshooting and educational QC steps."
 )
-
 TOOLS_LIST = (
     "🔧 **Available Tools**\n"
     "- quick_reference(text)\n"
@@ -475,9 +506,10 @@ async def on_message(msg: cl.Message):
     mode = cl.user_session.get("mode") or "agent"
     trace_before = len(cl.user_session.get("agent_trace") or [])
 
+    # Generate answer
     if mode == "agent":
         try:
-            result = await Runner.run(planner_agent, text, turns=5)
+            result = await Runner.run(planner_agent, text, turns=7, max_tools=3)
             answer = result.final_output or "I couldn’t generate a plan."
         except Exception as e:
             await cl.Message(content=f"⚠️ Planner error: {e}\nFalling back to simple mode.").send()
@@ -490,7 +522,7 @@ async def on_message(msg: cl.Message):
     new_trace = trace[trace_before:]
     trace_md = "### 🧭 Agent Trace (summary)\n" + "\n".join(f"- {t}" for t in new_trace) + "\n\n" if new_trace else ""
 
-    # Deterministic quick reference (always helpful)
+    # Deterministic quick reference
     qref = quick_reference(text)
     quick_md = (
         f"### 📘 Quick Reference\n"