Fix: strip non-ASCII chars before sending to llama-server

safety_checker.py

@@ -1,78 +1,19 @@
 """
 ClinIQ — Proactive Drug Safety Checker.
 
-Runs automatically after document ingestion — no question needed.
-Uses Qwen2.5-3B-Instruct to:
-  1. Extract all medications + allergies as structured JSON
-  2. Reason over combinations for contraindications and drug interactions
-  3. Return colour-coded alerts: DANGER / WARNING / INFO
-
-This is the signature feature that makes ClinIQ unique:
-a 3B model acting as an on-device pharmacist for small clinics.
+Single-call design: one prompt extracts medications/allergies AND identifies
+safety concerns simultaneously. This is faster and avoids context overflow.
 """
 
 from __future__ import annotations
-
-import json
-import re
+import json, re
 from dataclasses import dataclass, field
 from typing import List, Optional
 
-SAFETY_SYSTEM_PROMPT = """You are a clinical pharmacist assistant with deep knowledge of drug-allergy
-contraindications and drug interactions. You work for small clinics with no pharmacist on staff.
-Your job is to protect patients by flagging dangerous medication combinations BEFORE harm occurs.
-Be precise. Cite the clinical basis for each alert. Output ONLY valid JSON."""
-
-EXTRACTION_PROMPT = """Read the clinical documents below carefully.
-
-Task 1 — Find every medication/drug listed under headings like MEDICATIONS, CURRENT MEDICATIONS, MEDICATIONS AT DISCHARGE, PLAN, or ASSESSMENT.
-Task 2 — Find every allergy listed under headings like ALLERGIES, ALLERGY, or any sentence mentioning "allergic to" or "allergy".
-
-Documents:
-{context}
-
-Output ONLY valid JSON, no markdown, no explanation:
-{{
-  "medications": [
-    {{"name": "drug name", "dose": "dose", "frequency": "frequency"}}
-  ],
-  "allergies": [
-    {{"substance": "substance name", "reaction": "reaction description", "severity": "mild|moderate|severe|unknown"}}
-  ]
-}}"""
-
-SAFETY_PROMPT = """You are a clinical pharmacist. Given the patient's current medications and allergies,
-identify ALL safety concerns.
-
-Medications:
-{medications}
-
-Allergies:
-{allergies}
-
-Analyse for:
-1. Drug-allergy contraindications (same drug class as allergen)
-2. Drug-drug interactions (dangerous combinations)
-3. Missing standard-of-care medications or monitoring
-4. Dosing concerns given the diagnoses mentioned
-
-Output ONLY this JSON (no markdown):
-{{
-  "alerts": [
-    {{
-      "level": "DANGER|WARNING|INFO",
-      "title": "short title",
-      "detail": "clinical explanation with the specific drug names",
-      "recommendation": "what the clinician should do"
-    }}
-  ],
-  "overall_safety": "SAFE|REVIEW_NEEDED|URGENT"
-}}"""
-
 
 @dataclass
 class SafetyAlert:
-    level: str            # DANGER | WARNING | INFO
+    level: str          # DANGER | WARNING | INFO
     title: str
     detail: str
     recommendation: str
@@ -81,8 +22,8 @@ class SafetyAlert:
 @dataclass
 class SafetyReport:
     medications: List[dict] = field(default_factory=list)
-    allergies: List[dict] = field(default_factory=list)
-    alerts: List[SafetyAlert] = field(default_factory=list)
+    allergies:   List[dict] = field(default_factory=list)
+    alerts:      List[SafetyAlert] = field(default_factory=list)
     overall_safety: str = "SAFE"
     error: Optional[str] = None
 
@@ -90,77 +31,78 @@ class SafetyReport:
     def has_danger(self) -> bool:
         return any(a.level == "DANGER" for a in self.alerts)
 
-    @property
-    def alert_count(self) -> dict:
-        counts = {"DANGER": 0, "WARNING": 0, "INFO": 0}
-        for a in self.alerts:
-            counts[a.level] = counts.get(a.level, 0) + 1
-        return counts
 
+# Single combined prompt — extract AND analyse in one call
+_COMBINED_PROMPT = """\
+<|im_start|>system
+You are a clinical pharmacist AI. Analyse the clinical documents and identify drug safety concerns.
+Be thorough — patient safety depends on your accuracy. Output ONLY valid JSON.<|im_end|>
+<|im_start|>user
+Analyse the following clinical documents.
 
-def run_safety_check(context: str, call_model_fn) -> SafetyReport:
-    """
-    Full safety pipeline:
-      Step 1 — extract structured medications + allergies
-      Step 2 — reason over them for safety concerns
+Step 1: Extract every medication and allergy mentioned (look for MEDICATIONS, ALLERGIES, ALLERGY sections and any drug/allergy mentions in the text).
+Step 2: Identify DANGER (drug-allergy contraindications), WARNING (drug interactions, drug-disease contraindications), or INFO alerts.
 
-    call_model_fn(prompt, max_tokens, json_mode) -> str
-    """
-    report = SafetyReport()
+Documents:
+{context}
+
+Output ONLY this JSON (no markdown, no explanation):
+{{
+  "medications": [{{"name": "...", "dose": "...", "frequency": "..."}}],
+  "allergies": [{{"substance": "...", "reaction": "...", "severity": "mild|moderate|severe|unknown"}}],
+  "alerts": [
+    {{
+      "level": "DANGER|WARNING|INFO",
+      "title": "concise title",
+      "detail": "clinical explanation naming the specific drugs involved",
+      "recommendation": "what the clinician should do"
+    }}
+  ],
+  "overall_safety": "SAFE|REVIEW_NEEDED|URGENT"
+}}
+<|im_end|>
+<|im_start|>assistant
+"""
 
-    # ── Step 1: Extract ──────────────────────────────────────────────────────
-    extract_prompt = (
-        "<|im_start|>system\n"
-        "Extract medications and allergies from clinical documents. Output ONLY valid JSON.<|im_end|>\n"
-        "<|im_start|>user\n"
-        + EXTRACTION_PROMPT.format(context=context[:3000])
-        + "\n<|im_end|>\n<|im_start|>assistant\n"
-    )
 
-    try:
-        raw = call_model_fn(extract_prompt, 800, True)
-        match = re.search(r'\{.*\}', raw, re.DOTALL)
-        if match:
-            data = json.loads(match.group())
-            report.medications = data.get("medications", [])
-            report.allergies   = data.get("allergies", [])
-    except Exception as e:
-        report.error = f"Extraction failed: {e}"
+def run_safety_check(context: str, call_model_fn) -> SafetyReport:
+    report = SafetyReport()
+    if not context.strip():
         return report
 
-    if not report.medications and not report.allergies:
-        report.overall_safety = "SAFE"
-        return report
+    # Strip non-ASCII chars (em-dashes etc. corrupt llama-server JSON parsing)
+    import re
+    context = re.sub(r'[^\x00-\x7F]', '-', context)
+    # Keep well under the safe limit (~3000 chars of context)
+    trimmed = context[:3000]
 
-    # ── Step 2: Safety reasoning ──────────────────────────────────────────────
-    safety_prompt = (
-        "<|im_start|>system\n"
-        + SAFETY_SYSTEM_PROMPT
-        + "<|im_end|>\n<|im_start|>user\n"
-        + SAFETY_PROMPT.format(
-            medications=json.dumps(report.medications, indent=2),
-            allergies=json.dumps(report.allergies, indent=2),
-          )
-        + "\n<|im_end|>\n<|im_start|>assistant\n"
-    )
+    prompt = _COMBINED_PROMPT.format(context=trimmed)
 
     try:
-        raw = call_model_fn(safety_prompt, 1000, True)
+        raw = call_model_fn(prompt, 1000, False)
+        # Extract JSON from response
         match = re.search(r'\{.*\}', raw, re.DOTALL)
-        if match:
-            data = json.loads(match.group())
-            report.alerts = [
-                SafetyAlert(
-                    level=a.get("level", "INFO"),
-                    title=a.get("title", ""),
-                    detail=a.get("detail", ""),
-                    recommendation=a.get("recommendation", ""),
-                )
-                for a in data.get("alerts", [])
-            ]
-            report.overall_safety = data.get("overall_safety", "REVIEW_NEEDED")
+        if not match:
+            report.error = "Model did not return JSON"
+            return report
+
+        data = json.loads(match.group())
+        report.medications = data.get("medications", [])
+        report.allergies   = data.get("allergies", [])
+        report.overall_safety = data.get("overall_safety", "REVIEW_NEEDED")
+        report.alerts = [
+            SafetyAlert(
+                level          = a.get("level", "INFO"),
+                title          = a.get("title", ""),
+                detail         = a.get("detail", ""),
+                recommendation = a.get("recommendation", ""),
+            )
+            for a in data.get("alerts", [])
+        ]
+    except json.JSONDecodeError as e:
+        report.error = f"JSON parse error: {e}"
     except Exception as e:
-        report.error = f"Safety analysis failed: {e}"
+        report.error = str(e)
 
     return report
 
@@ -168,15 +110,11 @@ def run_safety_check(context: str, call_model_fn) -> SafetyReport:
 def report_to_dict(report: SafetyReport) -> dict:
     return {
         "overall_safety": report.overall_safety,
-        "medications": report.medications,
-        "allergies": report.allergies,
+        "medications":    report.medications,
+        "allergies":      report.allergies,
         "alerts": [
-            {
-                "level":          a.level,
-                "title":          a.title,
-                "detail":         a.detail,
-                "recommendation": a.recommendation,
-            }
+            {"level": a.level, "title": a.title,
+             "detail": a.detail, "recommendation": a.recommendation}
             for a in report.alerts
         ],
         "error": report.error,