Update app.py

app.py

@@ -56,16 +56,73 @@ print("✅ BioGPT ready!\n")
 
 
 # ── 2. BioGPT LLM — subclasses BaseLLM, bypasses LiteLLM ──────────
+# BioGPT is a causal LM trained on PubMed abstracts — it is NOT
+# instruction-tuned. It works by completing a biomedical sentence/phrase.
+# CrewAI passes a large system+user message block; we extract the clinical
+# case from it and build a tight completion-style prompt BioGPT can handle.
+
+def _extract_case(messages) -> str:
+    """Pull the raw clinical case text out of CrewAI message list."""
+    full = ""
+    if isinstance(messages, str):
+        full = messages
+    else:
+        for m in messages:
+            full += " " + m.get("content", "")
+    # Look for our case marker
+    m = re.search(r"(\d+-year-old .+?\.)", full, re.S)
+    if m:
+        return m.group(1).strip()
+    # Fallback: last non-empty content block
+    if isinstance(messages, list):
+        for m in reversed(messages):
+            c = m.get("content", "").strip()
+            if c:
+                return c[:300]
+    return full.strip()[:300]
+
+
+def _biogpt_complete(prompt: str, max_new: int = 200) -> str:
+    """Run BioGPT completion and return only the generated portion."""
+    raw = _pipe(prompt, max_new_tokens=max_new)[0]["generated_text"]
+    generated = raw[len(prompt):].strip()
+    # Trim to last complete sentence
+    if "." in generated:
+        generated = generated[: generated.rfind(".") + 1]
+    return generated or "Insufficient data generated."
+
+
 class BioGPTLLM(BaseLLM):
     """
     Subclassing BaseLLM routes CrewAI agent calls directly to the local
-    BioGPT HuggingFace pipeline. LiteLLM is never invoked, which
-    eliminates the starlette/streamlit/gradio version conflicts.
+    BioGPT HuggingFace pipeline. LiteLLM is never invoked.
+
+    Key insight: BioGPT is a PubMed completion model, not an instruction
+    follower. We detect which agent role is calling from the system message
+    and craft a tight domain-specific completion prompt accordingly.
     """
 
     def __init__(self):
         super().__init__(model="biogpt-local", temperature=0.65)
 
+    def _detect_role(self, messages) -> str:
+        """Detect which agent is calling by scanning the system message."""
+        system_text = ""
+        if isinstance(messages, list):
+            for m in messages:
+                if m.get("role") == "system":
+                    system_text = m.get("content", "").lower()
+                    break
+        if "diagnos" in system_text:
+            return "diagnosis"
+        if "treatment" in system_text or "pharmacol" in system_text:
+            return "treatment"
+        if "precaution" in system_text or "lifestyle" in system_text:
+            return "precaution"
+        if "coordinat" in system_text or "synthesise" in system_text or "summary" in system_text:
+            return "summary"
+        return "general"
+
     def call(
         self,
         messages: Union[str, List[Dict[str, str]]],
@@ -74,31 +131,42 @@ class BioGPTLLM(BaseLLM):
         available_functions: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> str:
-        # Flatten CrewAI message list → single prompt string
-        if isinstance(messages, str):
-            prompt = messages
+        role = self._detect_role(messages)
+        case = _extract_case(messages)
+
+        # Build tight PubMed-style completion prompts per agent role
+        if role == "diagnosis":
+            prompt = (
+                f"A {case} "
+                f"The differential diagnosis includes:"
+            )
+        elif role == "treatment":
+            prompt = (
+                f"For a patient with {case} "
+                f"The recommended treatment includes:"
+            )
+        elif role == "precaution":
+            prompt = (
+                f"For a patient with {case} "
+                f"Important precautions and follow-up recommendations include:"
+            )
+        elif role == "summary":
+            # For coordinator: re-run all three and concatenate
+            diag = _biogpt_complete(
+                f"A {case} The differential diagnosis includes:", 150)
+            treat = _biogpt_complete(
+                f"For a patient with {case} The recommended treatment includes:", 150)
+            prec = _biogpt_complete(
+                f"For a patient with {case} Important precautions include:", 150)
+            return (
+                f"## Diagnosis\n{diag}\n\n"
+                f"## Treatment\n{treat}\n\n"
+                f"## Precautions\n{prec}"
+            )
         else:
-            parts = []
-            for m in messages:
-                role    = m.get("role", "")
-                content = m.get("content", "")
-                if role == "system":
-                    parts.append(f"[System] {content}")
-                elif role in ("user", "human"):
-                    parts.append(content)
-            prompt = "\n".join(parts).strip()
-
-        if not prompt:
-            return "No prompt received."
-
-        raw       = _pipe(prompt)[0]["generated_text"]
-        generated = raw[len(prompt):].strip()
-
-        # Trim to last complete sentence
-        if "." in generated:
-            generated = generated[: generated.rfind(".") + 1]
+            prompt = f"The clinical findings for {case} suggest:"
 
-        return generated or "Insufficient data. Please consult a licensed physician."
+        return _biogpt_complete(prompt)
 
     def supports_function_calling(self) -> bool:
         return False