Fix: regex extraction on original text + noise filter

safety_checker.py

@@ -1,8 +1,13 @@
 """
 ClinIQ — Proactive Drug Safety Checker.
 
-Single-call design: one prompt extracts medications/allergies AND identifies
-safety concerns simultaneously. This is faster and avoids context overflow.
+Design:
+  Step 1 — Regex extraction: pull MEDICATIONS and ALLERGIES sections directly
+            from the text (fast, 100% reliable, no LLM needed).
+  Step 2 — LLM safety reasoning: feed the clean extracted lists to the model
+            and ask only "are there any dangerous combinations?"
+
+This avoids all JSON key confusion and context overflow issues.
 """
 
 from __future__ import annotations
@@ -11,20 +16,23 @@ from dataclasses import dataclass, field
 from typing import Any, List, Optional
 
 
-def _extract_json(text: str) -> Optional[Any]:
-    """Robustly extract JSON from model output, fixing common issues."""
-    # 1. Try direct parse first
-    try:
-        return json.loads(text.strip())
-    except Exception:
-        pass
+# ── JSON helpers ───────────────────────────────────────────────────────────────
+
+def _parse_json(text: str) -> Optional[Any]:
+    for candidate in [text.strip(), _brace_extract(text), _fix(_brace_extract(text))]:
+        if not candidate:
+            continue
+        try:
+            return json.loads(candidate)
+        except Exception:
+            pass
+    return None
+
 
-    # 2. Extract outermost {...} block
+def _brace_extract(text: str) -> str:
     start = text.find('{')
     if start == -1:
-        return None
-
-    # Walk to find matching closing brace
+        return ""
     depth = 0
     for i, ch in enumerate(text[start:], start):
         if ch == '{':
@@ -32,34 +40,25 @@ def _extract_json(text: str) -> Optional[Any]:
         elif ch == '}':
             depth -= 1
             if depth == 0:
-                candidate = text[start:i+1]
-                try:
-                    return json.loads(candidate)
-                except Exception:
-                    # Try fixing common issues
-                    fixed = _fix_json(candidate)
-                    try:
-                        return json.loads(fixed)
-                    except Exception:
-                        pass
-                break
-    return None
+                return text[start:i + 1]
+    return text[start:]
 
 
-def _fix_json(s: str) -> str:
-    """Fix common JSON errors from LLM output."""
-    # Remove trailing commas before } or ]
+def _fix(s: str) -> str:
+    if not s:
+        return s
     s = re.sub(r',\s*([}\]])', r'\1', s)
-    # Replace single quotes with double quotes (rough)
-    s = re.sub(r"'([^']*)':", r'"\1":', s)
-    # Remove control chars
     s = re.sub(r'[\x00-\x1f\x7f]', ' ', s)
-    return s
+    open_b = s.count('[') - s.count(']')
+    open_c = s.count('{') - s.count('}')
+    return s + ']' * max(0, open_b) + '}' * max(0, open_c)
 
 
+# ── Data classes ───────────────────────────────────────────────────────────────
+
 @dataclass
 class SafetyAlert:
-    level: str          # DANGER | WARNING | INFO
+    level: str
     title: str
     detail: str
     recommendation: str
@@ -67,57 +66,121 @@ 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)
-    overall_safety: str = "SAFE"
-    error: Optional[str] = None
+    medications:    List[str]        = field(default_factory=list)
+    allergies:      List[str]        = field(default_factory=list)
+    alerts:         List[SafetyAlert] = field(default_factory=list)
+    overall_safety: str              = "SAFE"
+    error:          Optional[str]    = None
 
     @property
     def has_danger(self) -> bool:
         return any(a.level == "DANGER" for a in self.alerts)
 
 
-# Single combined prompt — extract AND analyse in one call
-_COMBINED_PROMPT = (
+# ── Regex extraction (Step 1 — no LLM) ────────────────────────────────────────
+
+_SECTION_HEADER = re.compile(
+    r'^(PHYSICAL|HISTORY|HOSPITAL|PROCEDURE|ASSESSMENT|PLAN|IMPRESSION|FOLLOW|NOTE|REASON|'
+    r'DISCHARGE|PRIMARY|SECONDARY|VITAL|GENERAL|CARDIO|RESPIRATORY|EXAM|INVEST|PATIENT)',
+    re.IGNORECASE,
+)
+
+
+def _extract_lines(context: str) -> tuple[List[str], List[str]]:
+    """
+    Pull medication and allergy lines directly from document text using regex.
+    Finds ALL MEDICATIONS/ALLERGIES sections per document (not just the first).
+    """
+    meds:   List[str] = []
+    allgs:  List[str] = []
+    seen_m: set = set()
+    seen_a: set = set()
+
+    def _clean(line: str) -> str:
+        return re.sub(r'^[\d\.\-\*\s]+', '', line).strip()
+
+    def _is_noise(line: str) -> bool:
+        """Filter out section headers accidentally captured."""
+        return bool(_SECTION_HEADER.match(line)) or len(line) < 3
+
+    # Find all MEDICATIONS sections (there may be multiple per document)
+    for m in re.finditer(
+        r'MEDICATIONS?[^\n]*\n((?:[ \t]*[-\d*].*\n?){1,20})',
+        context, re.IGNORECASE
+    ):
+        for line in m.group(1).splitlines():
+            line = _clean(line)
+            key = line.lower()[:30]
+            if line and not _is_noise(line) and key not in seen_m:
+                seen_m.add(key)
+                meds.append(line)
+
+    # Find all ALLERGIES sections
+    for m in re.finditer(
+        r'ALLERGIES?[^\n]*\n((?:[ \t]*[-\d*].*\n?){1,10})',
+        context, re.IGNORECASE
+    ):
+        for line in m.group(1).splitlines():
+            line = _clean(line)
+            key = line.lower()[:30]
+            if line and not _is_noise(line) and key not in seen_a:
+                seen_a.add(key)
+                allgs.append(line)
+
+    return meds, allgs
+
+
+# ── LLM safety prompt (Step 2) ─────────────────────────────────────────────────
+
+_SAFETY_PROMPT = (
     "<|im_start|>system\n"
-    "You are a clinical pharmacist AI. Output ONLY valid JSON.<|im_end|>\n"
+    "You are a senior clinical pharmacist. Output ONLY valid JSON.<|im_end|>\n"
     "<|im_start|>user\n"
-    "Analyse documents and find:\n"
-    "1. All medications listed\n"
-    "2. All allergies listed\n"
-    "3. DANGER/WARNING/INFO safety alerts (drug-allergy contraindications, interactions)\n\n"
-    "Documents:\n{context}\n\n"
-    'Output ONLY JSON: {{"medications":[{{"name":"X","dose":"X","frequency":"X"}}],'
-    '"allergies":[{{"substance":"X","reaction":"X","severity":"X"}}],'
-    '"alerts":[{{"level":"DANGER","title":"X","detail":"X","recommendation":"X"}}],'
-    '"overall_safety":"SAFE|REVIEW_NEEDED|URGENT"}}\n'
+    "A patient has these medications and allergies. Find ALL safety concerns.\n\n"
+    "MEDICATIONS:\n{meds}\n\n"
+    "ALLERGIES:\n{allgs}\n\n"
+    "Look for:\n"
+    "- Drug-allergy contraindications (same class as allergen)\n"
+    "- Dangerous drug-drug interactions\n"
+    "- Drug-disease contraindications (e.g. beta-blockers in asthma)\n\n"
+    "Output ONLY this JSON:\n"
+    '{{"alerts":[{{"level":"DANGER","title":"...","detail":"...","recommendation":"..."}}],'
+    '"overall_safety":"SAFE"}}\n'
+    "level must be DANGER, WARNING, or INFO. "
+    "overall_safety must be SAFE, REVIEW_NEEDED, or URGENT.\n"
     "<|im_end|>\n<|im_start|>assistant\n"
 )
 
 
+# ── Main ───────────────────────────────────────────────────────────────────────
+
 def run_safety_check(context: str, call_model_fn) -> SafetyReport:
     report = SafetyReport()
     if not context.strip():
         return report
 
-    # Strip non-ASCII chars (em-dashes etc. corrupt llama-server JSON parsing)
-    import re
+    # Sanitize
     context = re.sub(r'[^\x00-\x7F]', '-', context)
-    # Keep well under the safe limit — 2500 chars of context
-    trimmed = context[:2500]
 
-    prompt = _COMBINED_PROMPT.format(context=trimmed)
+    # Step 1: regex extraction — fast, reliable, no LLM
+    meds, allgs = _extract_lines(context)
+
+    if not meds and not allgs:
+        report.overall_safety = "SAFE"
+        return report
+
+    report.medications = [{"name": m} for m in meds]
+    report.allergies   = [{"substance": a} for a in allgs]
+
+    # Step 2: LLM safety reasoning on clean lists
+    meds_text  = "\n".join(f"- {m}" for m in meds[:15])
+    allgs_text = "\n".join(f"- {a}" for a in allgs[:10])
+
+    prompt = _SAFETY_PROMPT.format(meds=meds_text, allgs=allgs_text)
 
     try:
-        raw = call_model_fn(prompt, 1000, False)
-
-        data = _extract_json(raw)
-        if data is None:
-            report.error = f"Could not parse model output: {raw[:200]}"
-            return report
-        report.medications = data.get("medications", [])
-        report.allergies   = data.get("allergies", [])
+        raw = call_model_fn(prompt, 600, False)
+        data = _parse_json(raw) or {}
         report.overall_safety = data.get("overall_safety", "REVIEW_NEEDED")
         report.alerts = [
             SafetyAlert(
@@ -127,9 +190,10 @@ def run_safety_check(context: str, call_model_fn) -> SafetyReport:
                 recommendation = a.get("recommendation", ""),
             )
             for a in data.get("alerts", [])
+            if isinstance(a, dict) and a.get("title")
         ]
     except Exception as e:
-        report.error = str(e)
+        report.error = f"Safety analysis error: {e}"
 
     return report