Update app.py

app.py

@@ -1,53 +1,136 @@
-# app.py
+# app.py — Self-contained (no external 'agents' package needed)
 import os, re, json
-from typing import List, Dict
+from typing import Any, Callable, Dict, List, Optional
+from dataclasses import dataclass, field
+
 import chainlit as cl
 from dotenv import load_dotenv
-from pydantic import BaseModel, Field
-
-# === Your agents framework (shim in ./agents) ===
-from agents import (
-    Agent,
-    Runner,
-    AsyncOpenAI,
-    OpenAIChatCompletionsModel,
-    set_tracing_disabled,
-    function_tool,
-)
-from agents.exceptions import InputGuardrailTripwireTriggered
+from openai import AsyncOpenAI as _SDKAsyncOpenAI
 
-# -----------------------------
-# Setup: auto provider (Gemini/OpenAI)
-# -----------------------------
-load_dotenv()
+# ========= Minimal "agents" shim (inline) =========
+def set_tracing_disabled(disabled: bool = True):
+    return disabled
+
+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}
+        if base_url:
+            kwargs["base_url"] = base_url
+        self._client = _SDKAsyncOpenAI(**kwargs)
+    @property
+    def client(self):
+        return self._client
+
+class OpenAIChatCompletionsModel:
+    def __init__(self, model: str, openai_client: AsyncOpenAI):
+        self.model = model
+        self.client = openai_client.client
+
+@dataclass
+class Agent:
+    name: str
+    instructions: str
+    model: OpenAIChatCompletionsModel
+    tools: Optional[List[Callable]] = field(default_factory=list)
 
-GEMINI_API_KEY = os.environ.get("Gem")
-OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+    def tool_specs(self) -> List[Dict[str, Any]]:
+        specs = []
+        for t in (self.tools or []):
+            if getattr(t, "_is_tool", False):
+                argnames = list(t.__code__.co_varnames[:t.__code__.co_argcount])
+                specs.append({
+                    "type": "function",
+                    "function": {
+                        "name": t.__name__,
+                        "description": (t.__doc__ or "")[:512],
+                        "parameters": {
+                            "type": "object",
+                            "properties": {p: {"type": "string"} for p in argnames},
+                            "required": argnames,
+                        },
+                    },
+                })
+        return specs
+
+class Runner:
+    @staticmethod
+    async def run(agent: Agent, user_input: str, context: Optional[Dict[str, Any]] = None):
+        msgs = [
+            {"role": "system", "content": agent.instructions},
+            {"role": "user", "content": user_input},
+        ]
+        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):
+            resp = await agent.model.client.chat.completions.create(
+                model=agent.model.model,
+                messages=msgs,
+                tools=tools or None,
+                tool_choice="auto" if tools else None,
+            )
+            msg = resp.choices[0].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)}
+                    msgs.append({
+                        "role": "tool",
+                        "tool_call_id": call.id,
+                        "name": fn_name,
+                        "content": json.dumps(result),
+                    })
+                continue
+
+            # Final answer
+            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
+
+        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
+
+# ========= Setup: provider auto-detect =========
+load_dotenv()
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 
 if GEMINI_API_KEY:
-    PROVIDER = "gemini"
     API_KEY = GEMINI_API_KEY
     BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/"
     MODEL_ID = "gemini-2.5-flash"
 elif OPENAI_API_KEY:
-    PROVIDER = "openai"
     API_KEY = 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 env/secrets.")
 
-set_tracing_disabled(disabled=True)
+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)
 
-ext_client: AsyncOpenAI = AsyncOpenAI(api_key=API_KEY, base_url=BASE_URL)
-llm_model: OpenAIChatCompletionsModel = OpenAIChatCompletionsModel(
-    model=MODEL_ID,
-    openai_client=ext_client,
-)
-
-# -----------------------------
-# Tools (function calling)
-# -----------------------------
+# ========= Tools =========
 @function_tool
 def infer_modality_from_filename(filename: str) -> dict:
     """
@@ -56,14 +139,14 @@ def infer_modality_from_filename(filename: str) -> dict:
     """
     f = (filename or "").lower()
     mapping = {
-        "xray": "X-ray", "x_ray": "X-ray", "xr": "X-ray", "cXR": "X-ray", "chest": "X-ray",
+        "xray": "X-ray", "x_ray": "X-ray", "xr": "X-ray", "chest": "X-ray",
         "mri": "MRI", "t1": "MRI", "t2": "MRI", "flair": "MRI", "dwi": "MRI", "adc": "MRI", "swi": "MRI",
         "ct": "CT", "cta": "CT",
         "ultrasound": "Ultrasound", "usg": "Ultrasound", "echo": "Ultrasound",
     }
-    for key, mod in mapping.items():
-        if key in f:
-            return {"modality": mod}
+    for k, v in mapping.items():
+        if k in f:
+            return {"modality": v}
     return {"modality": "unknown"}
 
 @function_tool
@@ -77,8 +160,8 @@ def imaging_reference_guide(modality: str) -> dict:
         return {
             "acquisition": [
                 "Projection radiography with ionizing radiation.",
-                "Common views: AP/PA/lateral; adjust kVp/mAs and positioning.",
-                "Grids/collimation reduce scatter and improve contrast."
+                "Views: AP/PA/lateral; tune kVp/mAs and positioning.",
+                "Grids & collimation reduce scatter to improve contrast."
             ],
             "artifacts": [
                 "Motion blur; under/overexposure.",
@@ -91,49 +174,49 @@ def imaging_reference_guide(modality: str) -> dict:
                 "Edge enhancement (unsharp) sparingly to avoid halos."
             ],
             "study_tips": [
-                "Use a systematic pattern (e.g., ABCDE for chest).",
-                "Compare sides; verify markers/labels/devices.",
+                "Use ABCDE (for CXR), check markers/labels/devices.",
+                "Compare sides and prior images.",
                 "Practice with checklists for consistency."
             ],
         }
     if mod in ["mri", "mr"]:
         return {
             "acquisition": [
-                "MR signal via RF pulses in a magnetic field; sequences set contrast.",
+                "MR signal via RF pulses; sequences define contrast.",
                 "Common: T1, T2, FLAIR, DWI/ADC, GRE/SWI.",
                 "TR/TE/flip angle trade off SNR, contrast, scan time."
             ],
             "artifacts": [
                 "Motion/ghosting; susceptibility near metal/air.",
                 "Chemical shift; Gibbs ringing.",
-                "B0/B1 inhomogeneity causing intensity non-uniformity."
+                "B0/B1 inhomogeneity causing intensity bias."
             ],
             "preprocessing": [
                 "Bias-field correction (N4).",
-                "Denoising (NLM); spatial registration/normalization.",
+                "Denoising (NLM); registration/normalization.",
                 "Skull stripping (brain); intensity standardization."
             ],
             "study_tips": [
-                "Learn what each sequence emphasizes (T1 anatomy; T2 fluid; FLAIR edema).",
-                "Always review diffusion for acute ischemia (with ADC).",
-                "Keep window/level consistent for timepoint comparison."
+                "Know sequence emphases (T1 anatomy; T2 fluid; FLAIR edema).",
+                "Review diffusion for acute ischemia (check ADC).",
+                "Keep window/level consistent across timepoints."
             ],
         }
     if mod in ["ct"]:
         return {
             "acquisition": [
-                "Helical CT; HU reflect X-ray attenuation.",
-                "Reconstruction kernels affect sharpness vs noise.",
-                "Contrast timing (arterial/venous) per clinical question."
+                "Helical CT; HU reflect attenuation.",
+                "Recon kernels affect sharpness vs noise.",
+                "Contrast timing (arterial/venous) per question."
             ],
             "artifacts": [
                 "Beam hardening streaks; partial volume; motion.",
-                "Metal artifacts; MAR/iterative recon can help."
+                "Metal artifacts; MAR/iterative recon help."
             ],
             "preprocessing": [
                 "Denoising (bilateral/NLM) with edge preservation.",
                 "Window/level by organ system (lung, mediastinum, bone).",
-                "Metal artifact reduction when available."
+                "Metal artifact reduction if available."
             ],
             "study_tips": [
                 "Use standard planes; scroll systematically.",
@@ -141,23 +224,23 @@ def imaging_reference_guide(modality: str) -> dict:
                 "Compare with priors when teaching cases."
             ],
         }
-    # Generic fallback
+    # Fallback (generic)
     return {
         "acquisition": [
-            "Acquisition parameters determine contrast, resolution, and noise.",
-            "Positioning and motion control are essential to quality."
+            "Acquisition parameters set contrast, resolution, noise.",
+            "Positioning & motion control drive image quality."
         ],
         "artifacts": [
-            "Motion blur/ghosting; foreign objects and hardware can create streaks/voids.",
+            "Motion blur or ghosting; foreign objects/hardware can streak.",
             "Under/overexposure or parameter misconfiguration."
         ],
         "preprocessing": [
-            "Denoising and contrast normalization can aid teaching clarity.",
-            "Registration and standard planes help consistent comparisons."
+            "Denoising & contrast normalization aid teaching clarity.",
+            "Registration & standard planes for consistent review."
         ],
         "study_tips": [
             "Adopt a checklist; compare bilaterally or across time.",
-            "Understand modality-specific controls (window/level, sequence types)."
+            "Understand modality-specific controls (window/level, sequences)."
         ],
     }
 
@@ -166,17 +249,13 @@ def file_facts(filename: str, size_bytes: str) -> dict:
     """Return simple file facts (name and size)."""
     return {"filename": filename, "size_bytes": size_bytes}
 
-# -----------------------------
-# Guardrails
-# -----------------------------
+# ========= Guardrails =========
 ALLOWED_COMMANDS = ("/help", "/policy")
-
 TOPIC_KEYWORDS = [
     "imaging","image","radiology","biomedical","device","equipment","oem","modality",
-    "acquisition","artifact","artifacts","preprocessing","window","level","sequence","kVp","mAs",
+    "acquisition","artifact","preprocessing","window","level","sequence","kVp","mAs",
     "mri","t1","t2","flair","dwi","adc","swi","ct","xray","x-ray","ultrasound","usg","echo"
 ]
-
 RE_FORBIDDEN_CLINICAL = re.compile(r"\b(diagnos(e|is|tic)|prescrib|medicat|treat(ment|ing)?|dose|drug|therapy)\b", re.I)
 RE_INVASIVE_REPAIR    = re.compile(r"\b(open(ing)?\s+(device|casing|cover)|solder|board[- ]level|reflow|replace\s+(capacitor|ic))\b", re.I)
 RE_ALARM_BYPASS       = re.compile(r"\b(bypass|disable|silence)\s+(alarm|alert|safety|interlock)\b", re.I)
@@ -204,57 +283,10 @@ def local_guard(text: str) -> List[str]:
         issues.append("phi_share_or_collect")
     return issues
 
-class SafetyCheck(BaseModel):
-    in_scope: bool = Field(description="True if about biomedical imaging education/troubleshooting.")
-    clinical_advice: bool = Field(description="Diagnosis/treatment/dosing/medication/therapy.")
-    invasive_repair: bool = Field(description="Open device, solder, board-level steps.")
-    alarm_bypass: bool = Field(description="Bypass/disable alarms or interlocks.")
-    firmware_tamper: bool = Field(description="Firmware hacks/unlock codes/passwords.")
-    phi_share_or_collect: bool = Field(description="Sharing/collecting personal identifiers.")
-    reasoning: str
-
-guardrail_agent = Agent(
-    name="Safety Classifier",
-    instructions=(
-        "You are a strict safety classifier for a **biomedical imaging education tutor**.\n"
-        "Return ONLY JSON with keys: {in_scope, clinical_advice, invasive_repair, alarm_bypass, "
-        "firmware_tamper, phi_share_or_collect, reasoning}.\n"
-        "- in_scope: true ONLY if the message is about biomedical IMAGING education/troubleshooting (no diagnosis).\n"
-        "- clinical_advice: diagnosis/treatment/dose/medication/therapy.\n"
-        "- invasive_repair: opening casing, soldering, board-level steps.\n"
-        "- alarm_bypass: silencing/disabling alarms/interlocks.\n"
-        "- firmware_tamper: rooting/jailbreaking/unlocking firmware/service modes/passwords.\n"
-        "- phi_share_or_collect: asking to share or store personal identifiers.\n"
-        "Respond with compact JSON only."
-    ),
-    model=llm_model,
-    output_type=SafetyCheck,
-)
-
-def parse_guard_json(s: str) -> Dict[str, bool]:
-    try:
-        m = re.search(r"\{.*\}", s or "", re.S)
-        data = json.loads(m.group(0) if m else (s or "{}"))
-        defaults = {
-            "in_scope": True,
-            "clinical_advice": False,
-            "invasive_repair": False,
-            "alarm_bypass": False,
-            "firmware_tamper": False,
-            "phi_share_or_collect": False
-        }
-        defaults.update({k: bool(v) for k, v in data.items() if k in defaults})
-        return defaults
-    except Exception:
-        return {"in_scope": True, "clinical_advice": False, "invasive_repair": False,
-                "alarm_bypass": False, "firmware_tamper": False, "phi_share_or_collect": False}
-
-# -----------------------------
-# Tutor Agent
-# -----------------------------
+# ========= Tutor Agent =========
 tutor_instructions = (
     "You are a Biomedical Imaging **Education** Tutor. Explain how images are acquired, common artifacts, "
-    "and preprocessing for study/teaching. You do NOT diagnose, identify diseases, or give clinical advice.\n\n"
+    "and preprocessing for study/teaching. Do NOT diagnose or give clinical advice.\n\n"
     "Output a concise, structured answer with sections in this order:\n"
     "1) Acquisition overview\n"
     "2) Common artifacts\n"
@@ -264,7 +296,6 @@ tutor_instructions = (
     "Use tools to infer modality (from filename) and fetch a modality-specific reference guide. "
     "If modality unclear, provide a generic overview and invite the user to specify."
 )
-
 tutor_agent = Agent(
     name="Biomedical Imaging Tutor",
     instructions=tutor_instructions,
@@ -272,9 +303,7 @@ tutor_agent = Agent(
     tools=[infer_modality_from_filename, imaging_reference_guide, file_facts],
 )
 
-# -----------------------------
-# UI strings
-# -----------------------------
+# ========= UI strings =========
 WELCOME = (
     "🎓 **Multimodal Biomedical Imaging Tutor**\n\n"
     "Upload an **MRI/X-ray/CT/Ultrasound** image (PNG/JPG), then ask what you’d like to learn.\n"
@@ -284,10 +313,10 @@ WELCOME = (
 POLICY = (
     "🛡️ **Safety & Scope Policy**\n"
     "- Scope: biomedical **imaging education/troubleshooting** only.\n"
-    "- No clinical advice (no diagnosis/treatment/dosing/medications).\n"
+    "- No clinical advice (diagnosis/treatment/dosing/medications).\n"
     "- No invasive repair steps (opening casing, soldering, board-level).\n"
     "- No alarm bypass or firmware tampering.\n"
-    "- No collecting or sharing personal identifiers.\n"
+    "- No collecting/sharing personal identifiers.\n"
     "- OEM manuals & local policy take priority."
 )
 REFUSAL = (
@@ -295,9 +324,7 @@ REFUSAL = (
     "I can explain **imaging acquisition, artifacts, and preprocessing** for education."
 )
 
-# -----------------------------
-# Chainlit flows
-# -----------------------------
+# ========= Chainlit flow =========
 @cl.on_chat_start
 async def on_chat_start():
     await cl.Message(content=WELCOME).send()
@@ -326,33 +353,16 @@ async def on_message(message: cl.Message):
     if text.lower().startswith("/policy"):
         await cl.Message(content=POLICY).send(); return
 
-    # Topic gate first
+    # Topic & guardrails
     if not on_topic(text):
         await cl.Message(
             content="I only discuss **biomedical imaging education** (acquisition, artifacts, preprocessing). "
                     "Please ask about MRI/X-ray/CT/Ultrasound imaging."
-        ).send()
-        return
+        ).send(); return
 
-    # Local guard (fast)
     issues = local_guard(text)
     if issues:
-        await cl.Message(content=REFUSAL + "\n\n" + POLICY).send()
-        return
-
-    # LLM guard (nuanced)
-    try:
-        verdict = await Runner.run(guardrail_agent, text)
-        flags = parse_guard_json(verdict.final_output)
-        if (not flags.get("in_scope", True)) or any(
-            flags.get(k, False) for k in
-            ["clinical_advice", "invasive_repair", "alarm_bypass", "firmware_tamper", "phi_share_or_collect"]
-        ):
-            await cl.Message(content=REFUSAL + "\n\n" + POLICY).send()
-            return
-    except Exception:
-        # If guard fails, continue (tutor prompt is already safety-constrained)
-        pass
+        await cl.Message(content=REFUSAL + "\n\n" + POLICY).send(); return
 
     # Context from uploaded file
     file_name = cl.user_session.get("last_file_name")
@@ -362,7 +372,7 @@ async def on_message(message: cl.Message):
     if file_size is not None: context_lines.append(f"Size: {file_size} bytes")
     context_block = "\n".join(context_lines)
 
-    # Compose user query for tutor
+    # Compose query for tutor
     user_query = text if not context_block else f"{text}\n\n[Context]\n{context_block}"
 
     # Run tutor