Update app.py

app.py

@@ -1,8 +1,10 @@
 # ---------------------------------------------------------
-# Chainlit app (Method A)
-# Project: Anatomy & Physiology Tutor + /gen Diagram (SD-XL)
+# Ayla — Biomedical Device Troubleshooting Assistant
+# Now with: Upload Device Manual (PDF/TXT) → search relevant pages
+# Refined, on-task, “real-person” tone. Education-only.
+# Flow: device -> symptom -> context/manual -> triage plan (+ optional /gen PNG)
 # ---------------------------------------------------------
-import os, json, time
+import os, json, time, re, io
 from dataclasses import dataclass, field
 from typing import Any, Callable, Dict, List, Optional
 
@@ -12,9 +14,10 @@ from pydantic import BaseModel
 from openai import AsyncOpenAI as _SDKAsyncOpenAI
 from huggingface_hub import InferenceClient
 from PIL import Image
+from pypdf import PdfReader  # ← new: PDF text extraction
 
 # =============================
-# Inline "agents" shim (no extra package needed)
+# Inline lightweight "agents" shim
 # =============================
 def set_tracing_disabled(disabled: bool = True):
     return disabled
@@ -23,29 +26,18 @@ def function_tool(func: Callable):
     func._is_tool = True
     return func
 
-def handoff(*args, **kwargs):
-    return None
-
-class InputGuardrail:
-    def __init__(self, guardrail_function: Callable):
-        self.guardrail_function = guardrail_function
-
 @dataclass
 class GuardrailFunctionOutput:
     output_info: Any
     tripwire_triggered: bool = False
     tripwire_message: str = ""
 
-class InputGuardrailTripwireTriggered(Exception):
-    pass
-
 class AsyncOpenAI:
     def __init__(self, api_key: str, base_url: Optional[str] = None):
         kwargs = {"api_key": api_key}
         if base_url:
             kwargs["base_url"] = base_url
         self._client = _SDKAsyncOpenAI(**kwargs)
-
     @property
     def client(self):
         return self._client
@@ -61,10 +53,6 @@ class Agent:
     instructions: str
     model: OpenAIChatCompletionsModel
     tools: Optional[List[Callable]] = field(default_factory=list)
-    handoff_description: Optional[str] = None
-    output_type: Optional[type] = None
-    input_guardrails: Optional[List[InputGuardrail]] = field(default_factory=list)
-
     def tool_specs(self) -> List[Dict[str, Any]]:
         specs = []
         for t in (self.tools or []):
@@ -95,8 +83,6 @@ class Runner:
         ]
         tools = agent.tool_specs()
         tool_map = {t.__name__: t for t in (agent.tools or []) if getattr(t, "_is_tool", False)}
-
-        # up to 4 tool-use rounds
         for _ in range(4):
             resp = await agent.model.client.chat.completions.create(
                 model=agent.model.model,
@@ -107,41 +93,29 @@ class Runner:
             choice = resp.choices[0]
             msg = choice.message
             msgs.append({"role": "assistant", "content": msg.content or "", "tool_calls": msg.tool_calls})
-
             if msg.tool_calls:
                 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}"}
                     if fn_name in tool_map:
                         try:
                             result = tool_map[fn_name](**args)
                         except Exception as e:
                             result = {"error": str(e)}
-                    else:
-                        result = {"error": f"Unknown tool: {fn_name}"}
                     msgs.append({
                         "role": "tool",
                         "tool_call_id": call.id,
                         "name": fn_name,
                         "content": json.dumps(result),
                     })
-                continue  # next round with tool outputs
-
+                continue
             final_text = msg.content or ""
             final_obj = type("Result", (), {})()
             final_obj.final_output = final_text
             final_obj.context = context or {}
-            if agent.output_type and issubclass(agent.output_type, BaseModel):
-                try:
-                    data = agent.output_type.model_validate_json(final_text)
-                    final_obj.final_output = data.model_dump_json()
-                    final_obj.final_output_as = lambda t: data
-                except Exception:
-                    final_obj.final_output_as = lambda t: final_text
-            else:
-                final_obj.final_output_as = lambda t: final_text
+            final_obj.final_output_as = lambda t: final_text
             return final_obj
-
         final_obj = type("Result", (), {})()
         final_obj.final_output = "Sorry, I couldn't complete the request."
         final_obj.context = context or {}
@@ -149,7 +123,7 @@ class Runner:
         return final_obj
 
 # =============================
-# App configuration
+# Config & model clients
 # =============================
 load_dotenv()
 API_KEY = os.environ.get("GEMINI_API_KEY") or os.environ.get("OPENAI_API_KEY")
@@ -158,230 +132,322 @@ if not API_KEY:
 HF_TOKEN = os.environ.get("HF_TOKEN")  # for SD-XL generation
 
 set_tracing_disabled(True)
-
 external_client: AsyncOpenAI = AsyncOpenAI(
     api_key=API_KEY,
     base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
 )
-llm_model: OpenAIChatCompletionsModel = OpenAIChatCompletionsModel(
-    model="gemini-2.5-flash",
-    openai_client=external_client,
-)
+llm_model = OpenAIChatCompletionsModel("gemini-2.5-flash", external_client)
 
 # =============================
-# A&P tools
+# Catalog & options
 # =============================
-@function_tool
-def topic_reference_guide(topic: str) -> dict:
-    """
-    Educational bullets for common Anatomy & Physiology topics.
-    Returns dict with anatomy, physiology, misconceptions, study_tips.
-    """
-    t = (topic or "").lower()
-
-    def pack(anat, phys, misc, tips):
-        return {"anatomy": anat, "physiology": phys, "misconceptions": misc, "study_tips": tips}
+DEVICE_CATALOG: Dict[str, Dict[str, List[str]]] = {
+    "ECG": {"symptoms": ["No signal", "Noisy baseline", "Leads off", "Paper/feed issue", "Inaccurate HR", "Other"]},
+    "Pulse oximeter": {"symptoms": ["Erratic SpO2", "Low signal", "No reading", "Probe damage", "Other"]},
+    "Infusion pump": {"symptoms": ["Occlusion alarm", "Air-in-line", "Inaccurate volume", "Battery issue", "Other"]},
+    "Ventilator": {"symptoms": ["Leak alarm", "High pressure alarm", "Sensor fault", "Other"]},
+    "Defibrillator": {"symptoms": ["Self-test fail", "Charging issue", "Pads/leads issue", "Other"]},
+    "Patient monitor": {"symptoms": ["NIBP error", "ECG noise", "SpO2 dropout", "Temp probe fault", "Other"]},
+    "Spirometer": {"symptoms": ["Fails 3L syringe", "No flow", "Inaccurate volumes", "Other"]},
+    "Glucometer": {"symptoms": ["Inconsistent readings", "Strip error", "No power", "Other"]},
+    "Suction machine": {"symptoms": ["Low suction", "Motor noise", "Overheating", "Other"]},
+    "Thermometer": {"symptoms": ["No reading", "Inaccurate", "Battery issue", "Other"]},
+    "Other device": {"symptoms": ["Other"]},
+}
+LIFE_CRITICAL = {"Ventilator", "Defibrillator", "Infusion pump"}
 
-    if any(k in t for k in ["cardiac conduction", "sa node", "av node", "bundle", "purkinje", "ecg"]):
+# =============================
+# Tools (education-only; non-invasive)
+# =============================
+@function_tool
+def device_reference_guide(device: str) -> dict:
+    """Safety-first notes, common faults, non-invasive checks, QC reminders, and escalation cues (education-only)."""
+    d = (device or "").lower()
+    def pack(safety, faults, checks, qc, escalate):
+        return {"safety": safety, "common_faults": faults, "quick_checks": checks, "qc_calibration": qc, "escalate_if": escalate}
+    if "ecg" in d:
         return pack(
-            anat=[
-                "SA node (RA) → AV node → His bundle → R/L bundle branches → Purkinje fibers.",
-                "Fibrous skeleton insulates atria from ventricles; AV node is the gateway."
-            ],
-            phys=[
-                "Pacemaker automaticity (If current) sets heart rate.",
-                "AV nodal delay allows ventricular filling; His–Purkinje enables synchronized contraction."
-            ],
-            misc=[
-                "Misconception: all tissues pace equally — SA node dominates.",
-                "PR interval ≈ AV nodal delay."
-            ],
-            tips=[
-                "Map ECG waves to mechanics (P/QRS/T).",
-                "Relate ion channels to nodal vs myocyte AP phases."
-            ],
+            ["Disconnect from mains before handling patient cables. No invasive service steps."],
+            ["Leads off/dry electrodes","Motion/EMI","Incorrect lead placement","Broken wires/high impedance"],
+            ["Verify lead map & skin prep","Replace electrodes","Route away from mains","Known-good lead set"],
+            ["Simulator: 1 mV @ 10 mm/mV","Check paper speed/gain","Electrical safety tests per schedule"],
+            ["Abnormal baseline with simulator","Failed safety tests","Burning smell/liquid ingress"]
         )
-
-    if any(k in t for k in ["nephron", "kidney", "gfr", "raas", "countercurrent"]):
+    if "pulse" in d or "oxi" in d:
         return pack(
-            anat=[
-                "Segments: Bowman’s → PCT → DLH → ALH (thin/thick) → DCT → collecting duct.",
-                "Cortex (glomeruli/PCT/DCT) vs medulla (loops/collecting ducts)."
-            ],
-            phys=[
-                "GFR via Starling forces; PCT bulk reabsorption; ALH generates gradient; DCT/CD fine-tune (ADH/aldosterone).",
-                "Countercurrent multiplication + vasa recta maintain medullary gradient."
-            ],
-            misc=[
-                "Misconception: water reabsorption is equal everywhere — hormones control CD concentration.",
-                "Urea also supports medullary osmolality."
-            ],
-            tips=[
-                "Sketch transporters per segment.",
-                "Practice ‘what if’ with ↑ADH/↑Aldosterone/���GFR."
-            ],
+            ["If patient-connected, assess clinically first; avoid opening casing."],
+            ["Poor perfusion/cold extremities","Motion","Nail polish/ambient light","Probe cable damage"],
+            ["Warm extremity / try another digit","Remove polish / shield light","Inspect/ reseat probe / swap known-good"],
+            ["Daily function check/simulator if available","Document vs reference"],
+            ["Erratic with known-good probe","Cracked housing/liquid ingress"]
         )
-
-    if any(k in t for k in ["alveolar", "gas exchange", "v/q", "ventilation perfusion", "oxygen dissociation"]):
+    if "infusion" in d:
         return pack(
-            anat=[
-                "Conducting vs respiratory zones; Type I (exchange) vs Type II (surfactant) pneumocytes."
-            ],
-            phys=[
-                "V/Q matching optimizes exchange; extremes: dead space (high V/Q) vs shunt (low V/Q).",
-                "O2–Hb curve shifts with pH, CO2, temp, 2,3-BPG (Bohr effect)."
-            ],
-            misc=[
-                "Misconception: uniform V/Q across lung — gravity/disease alter distribution.",
-                "Assuming PaO2–SaO2 linearity (it’s sigmoidal)."
-            ],
-            tips=[
-                "Draw V/Q along lung height.",
-                "Link spirometry patterns to mechanics."
-            ],
+            ["If connected: place patient on backup before checks. Never bypass alarms."],
+            ["Wrong IV set type","Occlusion up/downstream","Door not latched","Air in line"],
+            ["Verify set type","Reload tubing per OEM","Check clamps/kinks","Prime to remove air"],
+            ["Volumetric accuracy & occlusion tests per schedule"],
+            ["Inaccurate on gravimetric test","Alarm malfunction","Physical damage"]
         )
-
-    if any(k in t for k in ["neuron", "synapse", "action potential", "neurotransmitter"]):
+    if "vent" in d:
         return pack(
-            anat=[
-                "Neuron: dendrites, soma, axon hillock, axon; myelin & nodes."
-            ],
-            phys=[
-                "AP: Na+ depolarization → K+ repolarization; refractory periods.",
-                "Chemical synapse: Ca2+-dependent vesicle release; EPSPs/IPSPs summate."
-            ],
-            misc=[
-                "Misconception: any EPSP fires AP — threshold & summation matter."
-            ],
-            tips=[
-                "Relate channel states to AP phases.",
-                "Compare ionotropic vs metabotropic effects."
-            ],
+            ["Life-support. Do NOT troubleshoot while attached to a patient."],
+            ["Circuit leaks","Blocked filters","Sensor faults","Config mismatch"],
+            ["Self-test","Inspect filters/circuit","Swap known-good sensors","Confirm modes/settings"],
+            ["Full analyzer checks per schedule","Electrical safety"],
+            ["Any failed self-test/analyzer","Damage or unknown error code"]
         )
-
-    if any(k in t for k in ["muscle contraction", "excitation contraction", "sarcomere"]):
+    if "defib" in d:
         return pack(
-            anat=[
-                "Sarcomere: Z to Z; thin (actin/troponin/tropomyosin) vs thick (myosin).",
-                "T-tubules & SR triads (skeletal)."
-            ],
-            phys=[
-                "AP → DHPR → RyR → Ca2+ release → troponin C → cross-bridge cycling.",
-                "Force–length & force–velocity relationships; motor unit recruitment."
-            ],
-            misc=[
-                "ATP also needed for detachment & Ca2+ resequestration."
-            ],
-            tips=[
-                "Diagram cross-bridge cycle.",
-                "Predict effects of length on force."
-            ],
+            ["Life-support. Only qualified personnel; no invasive steps."],
+            ["Pads/leads/power issues","Self-test failure"],
+            ["Run self-test","Verify batteries/pads/leads","Inspect connections"],
+            ["Output energy verification with analyzer","Safety tests per schedule"],
+            ["Failed analyzer/self-test","Damage/liquid ingress"]
         )
-
     return pack(
-        anat=["Identify major structures & relationships."],
-        phys=["Describe inputs → key steps → outputs; control loops."],
-        misc=["Clarify commonly conflated terms; distinguish variation from pathology (education-only)."],
-        tips=["Make a labeled sketch; use ‘if-this-then-that’ scenarios to test understanding."],
+        ["If attached to patient, use backup first. Do not open device. Follow OEM policy."],
+        ["Power/battery issues","Loose/ damaged cables","Config mismatch","Environmental interference"],
+        ["Power cycle","Verify mains/battery","Reseat all connectors","Swap accessories","Restore defaults if safe"],
+        ["Scheduled QC/ES tests; verify with simulators if available"],
+        ["Persistent faults","Safety test failures","Burning smell/liquid ingress/cracks"]
     )
 
 @function_tool
-def study_outline(topic: str) -> dict:
-    """Return a suggested outline and practice prompts for the topic."""
-    return {
-        "topic": topic,
-        "subtopics": [
-            "Key structures & relationships",
-            "Mechanism (step-by-step)",
-            "Control & feedback",
-            "Quantitative intuition (flows, pressures, potentials)",
-            "Common misconceptions"
-        ],
-        "practice_prompts": [
-            f"Explain {topic} in 5 steps to a first-year student.",
-            f"Draw a block diagram of {topic} with inputs/outputs.",
-            f"What changes if a key parameter increases/decreases in {topic}?",
-        ],
-    }
+def symptom_checklist(device: str, symptom: str) -> dict:
+    """Targeted, non-invasive checklist for a device & symptom (education-only)."""
+    s = (symptom or "").lower()
+    steps = [
+        "Record model/serial/firmware; note exact error codes/messages.",
+        "Power/battery indicators OK; try another outlet; reseat battery if user-removable.",
+        "Inspect/replace accessories; reseat connectors; swap to known-good if available.",
+        "Confirm configuration/profile matches clinical setup.",
+    ]
+    if any(k in s for k in ["noise", "artifact", "interference"]):
+        steps += ["Minimize motion/EMI; route patient cables away from mains; proper skin prep if applicable."]
+    if any(k in s for k in ["no signal", "no reading"]):
+        steps += ["Try alternate sensor/lead; test with simulator/reference if available."]
+    if any(k in s for k in ["inaccurate", "wrong", "drift"]):
+        steps += ["Verify against simulator/reference; perform user-level zeroing if in OEM user manual."]
+    if any(k in s for k in ["alarm", "error"]):
+        steps += ["Do NOT bypass alarms; log alarm details; check OEM quick-ref."]
+    return {"device": device, "symptom": symptom, "steps": steps}
+
+@function_tool
+def escalation_policy(device: str) -> dict:
+    """Escalation line by device criticality."""
+    name = (device or "")
+    if any(x.lower() in name.lower() for x in LIFE_CRITICAL):
+        return {"policy": "Life-critical pathway: remove from service immediately, tag DO NOT USE, escalate to Biomedical Engineering/OEM."}
+    return {"policy": "If issues persist after basic checks or any safety/QC test fails, remove from service and escalate to Biomedical Engineering."}
 
 # =============================
-# Agents
+# Persona & Guardrails
 # =============================
-tutor_instructions = (
-    "You are an Anatomy & Physiology Tutor. TEACH, do not diagnose.\n"
-    "Given a topic (e.g., 'cardiac conduction', 'nephron physiology'), produce concise bullet points:\n"
-    "1) Anatomy (structures/relationships)\n"
-    "2) Physiology (inputs → steps → outputs)\n"
-    "3) Common misconceptions\n"
-    "4) Study tips\n"
-    "5) Caution: education only, no diagnosis\n"
-    "Use tools (topic_reference_guide, study_outline) to ground the response.\n"
-    "Avoid clinical diagnosis or treatment advice."
+ayla_instructions = (
+    "You are Ayla, a calm, concise biomedical device troubleshooting assistant for clinical engineers. "
+    "Stay strictly on task: only biomedical devices troubleshooting. "
+    "Prohibit diagnosis, treatment, invasive repairs, alarm bypass, firmware hacks, or collecting PHI. "
+    "Speak naturally, tight and actionable. "
+    "Output sections in this order: 1) Safety First 2) Likely Causes 3) Step-by-Step Checks 4) QC/Calibration 5) Escalate When. "
+    "If manual excerpts are provided in the prompt, use them but defer to OEM manual if any conflict. "
+    "End with a one-line summary. Education-only; refer to OEM manuals and policy."
 )
-tutor_agent = Agent(
-    name="A&P Tutor",
-    instructions=tutor_instructions,
+ayla_agent = Agent(
+    name="Ayla",
+    instructions=ayla_instructions,
     model=llm_model,
-    tools=[topic_reference_guide, study_outline],
+    tools=[device_reference_guide, symptom_checklist, escalation_policy],
 )
 
-guardrail_agent = Agent(
-    name="Safety Classifier",
+offtopic_guard = Agent(
+    name="Scope Guard",
     instructions=(
-        "Classify if the user's message requests medical diagnosis or unsafe medical advice, "
-        "and if it includes personal identifiers. Respond as JSON with fields: "
-        "{unsafe_medical_advice: bool, requests_diagnosis: bool, pii_included: bool, reasoning: string}."
+        "Classify if the message is within scope of biomedical device troubleshooting. "
+        "Return JSON: {in_scope: true/false, reason: '...'}"
+    ),
+    model=llm_model,
+)
+
+safety_guard = Agent(
+    name="Safety Guard",
+    instructions=(
+        "If the message asks for diagnosis, treatment, invasive repair, bypassing safety/alarms, hacking firmware, or PHI collection, "
+        "return JSON {unsafe: true, reason: '...'} else {unsafe: false}."
     ),
     model=llm_model,
 )
 
 # =============================
-# SD-XL helper (HF Inference API) -> PNG path
+# SD-XL helper -> PNG path
 # =============================
 def sdxl_png(prompt: str, negative: str = "") -> str:
     token = os.getenv("HF_TOKEN")
     if not token:
         raise RuntimeError("Missing HF_TOKEN. Add it in Space secrets.")
     client = InferenceClient("stabilityai/stable-diffusion-xl-base-1.0", token=token)
-
     img: Image.Image = client.text_to_image(
-        prompt = prompt + " | educational, clean vector style, white background, no labels, safe-for-work",
+        prompt = prompt + " | clean vector style, white background, no text labels, educational, safe-for-work",
         negative_prompt = negative or "text, watermark, logo, gore, photorealistic patient, clutter",
         width = 1024, height = 768,
         guidance_scale = 7.5,
-        num_inference_steps = 30,
+        num_inference_steps = 24,
         seed = 42
     )
     out_dir = os.environ.get("CHAINLIT_FILES_DIR") or os.path.join(os.getcwd(), ".files")
     os.makedirs(out_dir, exist_ok=True)
-    path = os.path.join(out_dir, f"sdxl-{int(time.time())}.png")
+    path = os.path.join(out_dir, f"sdxl-troubleshoot-{int(time.time())}.png")
     img.save(path)
     return path
 
 # =============================
-# Chainlit flows
+# Manual upload, indexing, and search
+# =============================
+def sanitize_short(text: str, maxlen: int = 80) -> str:
+    t = re.sub(r"\s+", " ", (text or "")).strip()
+    return t[:maxlen]
+
+def _extract_pdf_pages(data: bytes) -> List[Dict[str, Any]]:
+    reader = PdfReader(io.BytesIO(data))
+    pages = []
+    for i, pg in enumerate(reader.pages, start=1):
+        try:
+            txt = pg.extract_text() or ""
+        except Exception:
+            txt = ""
+        pages.append({"page": i, "text": txt})
+    return pages
+
+def _extract_txt_pages(data: bytes, chunk_chars: int = 1400) -> List[Dict[str, Any]]:
+    try:
+        txt = data.decode("utf-8", errors="ignore")
+    except Exception:
+        txt = ""
+    chunks = []
+    i = 0
+    while i < len(txt):
+        chunk = txt[i:i+chunk_chars]
+        chunks.append({"page": len(chunks)+1, "text": chunk})
+        i += chunk_chars
+    return chunks
+
+def _manual_search(pages: List[Dict[str, Any]], query: str, topk: int = 3) -> List[Dict[str, Any]]:
+    if not pages:
+        return []
+    terms = [w for w in re.findall(r"\w+", (query or "").lower()) if len(w) > 2]
+    scored = []
+    for p in pages:
+        low = p.get("text","").lower()
+        score = sum(low.count(t) for t in terms)
+        if score > 0:
+            scored.append((score, p))
+    scored.sort(key=lambda x: x[0], reverse=True)
+    return [p for _, p in scored[:topk]] or (pages[:1])  # fallback first page
+
+def _make_excerpt(text: str, terms: List[str], window: int = 380) -> str:
+    t = (text or "")
+    low = t.lower()
+    idxs = [low.find(term) for term in terms if term in low]
+    start = max(0, min([i for i in idxs if i >= 0], default=0) - window)
+    end = min(len(t), start + 2*window)
+    return re.sub(r"\s+", " ", t[start:end]).strip()
+
+# =============================
+# UI helpers
+# =============================
+async def present_device_menu():
+    actions = [cl.Action(name=f"dev_{i}", value=name, label=name) for i, name in enumerate(DEVICE_CATALOG.keys(), 1)]
+    return await cl.AskActionMessage(
+        content="👋 I’m **Ayla**. Let’s keep this tight and practical.\n\n**Pick a device** to start:",
+        actions=actions,
+        timeout=180,
+    ).send()
+
+async def present_symptom_menu(device_name: str):
+    options = DEVICE_CATALOG.get(device_name, {}).get("symptoms", ["Other"])
+    actions = [cl.Action(name=f"sym_{i}", value=s, label=s) for i, s in enumerate(options, 1)]
+    return await cl.AskActionMessage(
+        content=f"**Device:** {device_name}\n\nSelect the **symptom**:",
+        actions=actions,
+        timeout=180,
+    ).send()
+
+async def present_context_menu(device_name: str, symptom: str, manual_loaded: bool):
+    actions = [
+        cl.Action(name="ctx_patient_connected", value="patient_connected", label="Patient currently connected?"),
+        cl.Action(name="ctx_tried_accessories", value="tried_accessories", label="Already swapped accessories?"),
+        cl.Action(name="ctx_env_checked", value="env_checked", label="Checked power/EMI/filters?"),
+        cl.Action(name="ctx_upload_manual", value="upload_manual", label="Upload Manual (PDF/TXT)"),
+    ]
+    if manual_loaded:
+        actions += [
+            cl.Action(name="ctx_view_hits", value="view_hits", label="View Manual Matches"),
+            cl.Action(name="ctx_clear_manual", value="clear_manual", label="Remove Manual"),
+        ]
+    actions += [
+        cl.Action(name="ctx_start", value="start_triage", label="Start Triage ✅"),
+        cl.Action(name="ctx_restart", value="restart", label="Change Device")
+    ]
+    return await cl.AskActionMessage(
+        content=f"**Device:** {device_name}\n**Symptom:** {symptom}\n\nAdd quick context or **upload the device manual**, then **Start Triage**:",
+        actions=actions,
+        timeout=240,
+    ).send()
+
+def session_reset():
+    cl.user_session.set("stage", "await_device")
+    cl.user_session.set("device", None)
+    cl.user_session.set("symptom", None)
+    cl.user_session.set("flags", set())
+    cl.user_session.set("manual", None)  # {"name": str, "pages": List[dict]}
+
+def _manual_info() -> str:
+    m = cl.user_session.get("manual")
+    if not m:
+        return "No manual attached."
+    return f"Manual: **{m.get('name','(unnamed)')}** — {len(m.get('pages',[]))} page-chunks indexed."
+
+# =============================
+# Guards
+# =============================
+def is_off_topic_label(json_str: str) -> bool:
+    try:
+        data = json.loads(json_str) if isinstance(json_str, str) else json_str
+        return not bool(data.get("in_scope", False))
+    except Exception:
+        return False
+
+def is_unsafe(json_str: str) -> bool:
+    try:
+        data = json.loads(json_str) if isinstance(json_str, str) else json_str
+        return bool(data.get("unsafe", False))
+    except Exception:
+        return False
+
+# =============================
+# Chat flow
 # =============================
 WELCOME = (
-    "🧠 **Anatomy & Physiology Tutor**  +  🎨 **/gen Diagram (SD-XL)**\n\n"
-    "• Ask any A&P topic (e.g., *cardiac conduction*, *nephron physiology*, *gas exchange*).\n"
-    "• Generate a clean educational diagram (PNG) with:\n"
-    "  `/gen isometric nephron diagram, flat vector, white background, no labels`\n\n"
-    "⚠️ Education only — no diagnosis or clinical advice."
+    "🛠️ **Ayla — Biomedical Device Troubleshooting Assistant**\n"
+    "Education-only. No diagnosis, no invasive service, no alarm bypass. OEM manual & policy rule the day.\n\n"
+    "Tip: Generate a simple diagram anytime: `/gen flowchart for pulse oximeter troubleshooting`"
 )
 
 @cl.on_chat_start
 async def on_chat_start():
+    session_reset()
     await cl.Message(content=WELCOME).send()
+    await present_device_menu()
 
 @cl.on_message
 async def on_message(message: cl.Message):
     text = (message.content or "").strip()
 
-    # Quick command: /gen <desc>  → SD-XL PNG
+    # Slash command: /gen
     if text.lower().startswith("/gen "):
         desc = text[5:].strip()
         if not desc:
-            await cl.Message(content="Please provide a description after `/gen`.\nExample: `/gen isometric nephron diagram, flat vector`").send()
+            await cl.Message(content="Usage: `/gen flowchart for <device> troubleshooting`").send()
             return
         try:
             path = sdxl_png(desc)
@@ -389,58 +455,250 @@ async def on_message(message: cl.Message):
             await cl.Message(content=f"Generation failed: {e}").send()
             return
         await cl.Message(
-            content="🎨 **Generated diagram** (education only):",
+            content="🎨 **Generated diagram** (education-only):",
             elements=[cl.Image(path=path, name=os.path.basename(path), display="inline")],
         ).send()
         return
 
-    # Light safety check (blocks diagnosis/treatment requests)
+    # Scope & safety checks
     try:
-        safety = await Runner.run(guardrail_agent, text)
-        parsed = safety.final_output
-        try:
-            data = json.loads(parsed) if isinstance(parsed, str) else parsed
-        except Exception:
-            data = {}
-        if isinstance(data, dict) and (data.get("unsafe_medical_advice") or data.get("requests_diagnosis")):
+        scope = await Runner.run(offtopic_guard, text)
+        if is_off_topic_label(scope.final_output):
+            await cl.Message(content="Let’s stay focused on **biomedical device troubleshooting**. Please pick a device below.").send()
+            await present_device_menu()
+            return
+    except Exception:
+        pass
+    try:
+        unsafe = await Runner.run(safety_guard, text)
+        if is_unsafe(unsafe.final_output):
             await cl.Message(
-                content="🚫 I can’t provide diagnosis or treatment advice. I can teach A&P concepts and generate **educational** diagrams."
+                content="🚫 I can’t help with diagnosis, treatment, invasive repairs, alarm bypass, firmware hacks, or collecting personal data. "
+                        "Let’s stick to **safe, non-invasive troubleshooting**. Pick a device again:"
             ).send()
+            await present_device_menu()
             return
     except Exception:
         pass
 
-    # Tutor mode
-    topic = text or "anatomy and physiology overview"
-    result = await Runner.run(tutor_agent, topic)
+    # State machine guardrails
+    stage = cl.user_session.get("stage", "await_device")
+    device = cl.user_session.get("device")
+    if stage == "await_device":
+        await cl.Message(content="Please **choose a device** from the buttons so I can tailor the checks.").send()
+        await present_device_menu()
+        return
+    if stage == "await_symptom":
+        await cl.Message(content=f"**Device:** {device}\nPlease **select a symptom** from the options.").send()
+        await present_symptom_menu(device)
+        return
+    if stage == "await_context":
+        # Free text here just toggles quick context flags
+        low = text.lower()
+        flags: set = cl.user_session.get("flags") or set()
+        if "patient" in low and "connect" in low:
+            flags.add("patient_connected")
+        if "accessor" in low or "probe" in low or "lead" in low:
+            flags.add("tried_accessories")
+        if "emi" in low or "power" in low or "filter" in low:
+            flags.add("env_checked")
+        cl.user_session.set("flags", flags)
+        await cl.Message(content="Noted. You can also **Upload Manual** or tap **Start Triage**.").send()
+        return
 
-    # Tool-based quick reference + outline
-    try:
-        guide = topic_reference_guide(topic)
-    except Exception:
-        guide = {}
-    try:
-        outline = study_outline(topic)
-    except Exception:
-        outline = {}
+# Capture custom symptom typed between steps
+@cl.on_message
+async def capture_custom_symptom(message: cl.Message):
+    if cl.user_session.get("stage") != "await_custom_symptom":
+        return
+    raw = sanitize_short(message.content, 80)
+    if not raw or len(raw) < 3:
+        await cl.Message(content="Please enter a short symptom (≥ 3 chars).").send()
+        return
+    cl.user_session.set("symptom", raw)
+    cl.user_session.set("stage", "await_context")
+    manual_loaded = cl.user_session.get("manual") is not None
+    await present_context_menu(cl.user_session.get("device"), raw, manual_loaded)
 
-    def bullets(arr):
-        return "\n".join([f"- {b}" for b in arr]) if isinstance(arr, list) else "- (n/a)"
+# =============================
+# Action handlers
+# =============================
+@cl.action_callback("*")
+async def on_action(action: cl.Action):
+    name = action.name or ""
+    stage = cl.user_session.get("stage", "await_device")
+    device = cl.user_session.get("device")
+    symptom = cl.user_session.get("symptom")
+    flags: set = cl.user_session.get("flags") or set()
+    manual = cl.user_session.get("manual")
+
+    # Device selection
+    if name.startswith("dev_"):
+        device_name = action.value
+        cl.user_session.set("device", device_name)
+        cl.user_session.set("stage", "await_symptom")
+        await present_symptom_menu(device_name)
+        return
 
-    quick_ref = (
-        f"### 📚 Quick Reference: {topic}\n"
-        f"**Anatomy**\n{bullets(guide.get('anatomy', []))}\n\n"
-        f"**Physiology**\n{bullets(guide.get('physiology', []))}\n\n"
-        f"**Common Misconceptions**\n{bullets(guide.get('misconceptions', []))}\n\n"
-        f"**Study Tips**\n{bullets(guide.get('study_tips', []))}\n\n"
-    )
+    # Symptom selection
+    if name.startswith("sym_"):
+        chosen = action.value
+        if chosen == "Other":
+            await cl.Message(content="Type a short symptom (≤ 80 chars).").send()
+            cl.user_session.set("stage", "await_custom_symptom")
+            return
+        cl.user_session.set("symptom", chosen)
+        cl.user_session.set("stage", "await_context")
+        manual_loaded = cl.user_session.get("manual") is not None
+        await present_context_menu(cl.user_session.get("device"), chosen, manual_loaded)
+        return
+
+    # Context toggles
+    if name == "ctx_patient_connected":
+        flags.add("patient_connected")
+        cl.user_session.set("flags", flags)
+        await cl.Message(content="Marked: patient currently connected.").send()
+        return
+    if name == "ctx_tried_accessories":
+        flags.add("tried_accessories")
+        cl.user_session.set("flags", flags)
+        await cl.Message(content="Marked: accessories already swapped/checked.").send()
+        return
+    if name == "ctx_env_checked":
+        flags.add("env_checked")
+        cl.user_session.set("flags", flags)
+        await cl.Message(content="Marked: power/EMI/filters checked.").send()
+        return
 
-    outline_md = ""
-    if outline:
-        subs = bullets(outline.get("subtopics", []))
-        qs = bullets(outline.get("practice_prompts", []))
-        outline_md = f"### 🗂️ Suggested Outline\n{subs}\n\n### 📝 Practice Prompts\n{qs}\n\n"
+    # Upload manual
+    if name == "ctx_upload_manual":
+        files = await cl.AskFileMessage(
+            content="Upload the **device manual** (PDF or TXT). Max ~20 MB.\n\n" + _manual_info(),
+            accept=["application/pdf", "text/plain"],
+            max_size_mb=20,
+            max_files=1,
+            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()
+        pages = []
+        try:
+            if f.mime == "application/pdf" or f.name.lower().endswith(".pdf"):
+                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.\nYou can **View Manual Matches** or **Start Triage**.").send()
+        await present_context_menu(cl.user_session.get("device"), cl.user_session.get("symptom"), manual_loaded=True)
+        return
+
+    # View manual matches
+    if name == "ctx_view_hits":
+        manual = cl.user_session.get("manual")
+        if not manual:
+            await cl.Message(content="No manual attached yet.").send()
+            return
+        query = f"{cl.user_session.get('device')} {cl.user_session.get('symptom')}"
+        hits = _manual_search(manual.get("pages", []), query, topk=3)
+        terms = [w for w in re.findall(r"\w+", query.lower()) if len(w) > 2]
+        parts = []
+        for h in hits:
+            excerpt = _make_excerpt(h.get("text",""), terms, window=420)
+            parts.append(f"**p.{h.get('page')}** — {excerpt}")
+        if not parts:
+            parts = ["No obvious matches found; try different symptom phrasing."]
+        await cl.Message(content="### 🔎 Manual Matches\n" + "\n\n".join(parts)).send()
+        return
+
+    # Clear manual
+    if name == "ctx_clear_manual":
+        cl.user_session.set("manual", None)
+        await cl.Message(content="Manual removed.").send()
+        await present_context_menu(cl.user_session.get("device"), cl.user_session.get("symptom"), manual_loaded=False)
+        return
+
+    # Restart flow
+    if name == "ctx_restart":
+        session_reset()
+        await present_device_menu()
+        return
+
+    # Start triage
+    if name == "ctx_start":
+        device = cl.user_session.get("device") or "device"
+        symptom = cl.user_session.get("symptom") or "issue"
+        flags = cl.user_session.get("flags") or set()
+        manual = cl.user_session.get("manual")
+
+        context_lines = []
+        if "patient_connected" in flags:
+            context_lines.append("Patient currently connected.")
+        if "tried_accessories" in flags:
+            context_lines.append("Accessories already swapped/checked.")
+        if "env_checked" in flags:
+            context_lines.append("Power/EMI/filters checked.")
+
+        # Manual excerpts (if provided)
+        manual_section = ""
+        if manual and manual.get("pages"):
+            query = f"{device} {symptom}"
+            terms = [w for w in re.findall(r"\w+", query.lower()) if len(w) > 2]
+            hits = _manual_search(manual["pages"], query, topk=3)
+            blocks = []
+            for h in hits:
+                excerpt = _make_excerpt(h.get("text",""), terms, window=420)
+                blocks.append(f"[Manual p.{h.get('page')}] {excerpt}")
+            if blocks:
+                manual_section = "Manual excerpts (for reference; OEM manual prevails if any conflict):\n" + "\n".join(blocks)
+
+        triage_prompt = (
+            f"Device: {device}\n"
+            f"Symptom: {symptom}\n"
+            f"Context: {', '.join(context_lines) if context_lines else 'n/a'}\n\n"
+            f"{manual_section}\n\n"
+            "Produce structured sections exactly:\n"
+            "1) Safety First (specific to device & context; non-invasive)\n"
+            "2) Likely Causes (ranked)\n"
+            "3) Step-by-Step Checks (bullet list, do-not-open, do-not-bypass alarms)\n"
+            "4) QC/Calibration (what to verify and with what tool)\n"
+            "5) Escalate When (clear triggers)\n"
+            "End with a one-line summary.\n"
+            "If device is life-critical (ventilator/defibrillator/infusion pump) and patient_connected, "
+            "start with: REMOVE FROM SERVICE & USE BACKUP — then proceed with safe checks off-patient."
+        )
+
+        result = await Runner.run(ayla_agent, triage_prompt)
+
+        # Tool-sourced quick guide for reliability
+        guide = device_reference_guide(device)
+        checklist = symptom_checklist(device, symptom)
+        policy = escalation_policy(device)
+
+        def bullets(arr): return "\n".join([f"- {b}" for b in (arr or [])]) if arr else "-"
+
+        quick = (
+            f"### 📘 Quick Reference: {device}\n"
+            f"**Safety**\n{bullets(guide.get('safety'))}\n\n"
+            f"**Common Faults**\n{bullets(guide.get('common_faults'))}\n\n"
+            f"**Quick Checks**\n{bullets(guide.get('quick_checks'))}\n\n"
+            f"**QC / Calibration**\n{bullets(guide.get('qc_calibration'))}\n\n"
+            f"**Escalate If**\n{bullets(guide.get('escalate_if'))}\n\n"
+            f"### 📝 Targeted Checklist for Symptom\n{bullets(checklist.get('steps'))}\n\n"
+            f"**Policy**\n- {policy.get('policy')}\n\n"
+            f"> ⚠️ Education-only. Refer to OEM manuals/policy. No invasive service.\n"
+        )
 
-    caution = "> ⚠️ Education only — no diagnosis or treatment advice."
-    answer = result.final_output or "I couldn’t generate an explanation."
-    await cl.Message(content=f"{quick_ref}{outline_md}---\n{answer}\n\n{caution}").send()
+        answer = result.final_output or "I couldn’t generate a troubleshooting plan."
+        session_reset()
+        await cl.Message(content=f"{quick}\n---\n{answer}").send()
+        await present_device_menu()
+        return