Fix: flexible medication/allergy extraction for real-world document formats

safety_checker.py

@@ -89,7 +89,8 @@ _SECTION_HEADER = re.compile(
 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).
+    Covers many real-world clinical document formats, not just standard headers.
+    Finds ALL matching sections per document (not just the first).
     """
     meds:   List[str] = []
     allgs:  List[str] = []
@@ -97,35 +98,90 @@ def _extract_lines(context: str) -> tuple[List[str], List[str]]:
     seen_a: set = set()
 
     def _clean(line: str) -> str:
-        return re.sub(r'^[\d\.\-\*\s]+', '', line).strip()
+        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)
+    # ── Medication section headers — covers common real-world formats ──────────
+    # Examples: "MEDICATIONS:", "Current Medications:", "Active Medications:",
+    #           "Medication List:", "Rx:", "Prescriptions:", "Drug List:",
+    #           "Home Medications:", "Discharge Medications:", "Outpatient Meds:"
+    MED_HEADER = re.compile(
+        r'(?:current\s+|active\s+|home\s+|discharge\s+|outpatient\s+|admission\s+|'
+        r'inpatient\s+|chronic\s+|long.?term\s+)?'
+        r'(?:medications?|meds?|medicines?|drugs?|prescriptions?|rx)'
+        r'(?:\s+list)?'
+        r'[^\n]{0,30}\n',
+        re.IGNORECASE,
+    )
+
+    # ── Allergy section headers — covers common real-world formats ─────────────
+    # Examples: "ALLERGIES:", "Drug Allergies:", "Known Allergies:",
+    #           "Allergy List:", "Sensitivities:", "Adverse Reactions:",
+    #           "Drug Intolerances:", "NKA" (no known allergies)
+    ALLERGY_HEADER = re.compile(
+        r'(?:drug\s+|food\s+|environmental\s+|known\s+|reported\s+)?'
+        r'(?:allergies?|allergy|sensitivities|sensitivity|intolerances?|'
+        r'adverse\s+reactions?|hypersensitivities?)'
+        r'(?:\s+list)?'
+        r'[^\n]{0,30}\n',
+        re.IGNORECASE,
+    )
+
+    # ── Also catch inline bullet/colon lists after a header on the same line ───
+    # e.g. "Allergies: Penicillin (rash), Sulfa"
+    ALLERGY_INLINE = re.compile(
+        r'(?:drug\s+|known\s+)?(?:allergies?|allergy|sensitivities?|adverse\s+reactions?)'
+        r'\s*:\s*([^\n]{3,200})',
+        re.IGNORECASE,
+    )
+    MED_INLINE = re.compile(
+        r'(?:current\s+|active\s+|home\s+)?(?:medications?|meds?|rx)'
+        r'\s*:\s*([^\n]{3,200})',
+        re.IGNORECASE,
+    )
+
+    def _harvest(pattern, context, seen, out, max_lines=20):
+        """Find all sections matching pattern and extract bullet/numbered lines."""
+        for match in re.finditer(pattern, context):
+            # Grab the block of lines after the header
+            start = match.end()
+            block = context[start:start + 1200]
+            count = 0
+            for line in block.splitlines():
+                if count >= max_lines:
+                    break
+                # Stop if we hit another section header (blank-ish line or ALL CAPS heading)
+                stripped = line.strip()
+                if not stripped:
+                    continue
+                if re.match(r'^[A-Z][A-Z\s]{4,}:', stripped) and not re.match(r'^[A-Z][a-z]', stripped):
+                    break
+                cleaned = _clean(stripped)
+                key = cleaned.lower()[:30]
+                if cleaned and not _is_noise(cleaned) and key not in seen:
+                    seen.add(key)
+                    out.append(cleaned)
+                    count += 1
+
+    def _harvest_inline(pattern, context, seen, out):
+        """Catch single-line format: 'Allergies: Penicillin, Sulfa'"""
+        for match in re.finditer(pattern, context):
+            raw = match.group(1)
+            # Split on commas or semicolons
+            for part in re.split(r'[,;]', raw):
+                cleaned = _clean(part.strip())
+                key = cleaned.lower()[:30]
+                if cleaned and not _is_noise(cleaned) and key not in seen:
+                    seen.add(key)
+                    out.append(cleaned)
+
+    _harvest(MED_HEADER,     context, seen_m, meds, max_lines=20)
+    _harvest(ALLERGY_HEADER, context, seen_a, allgs, max_lines=10)
+    _harvest_inline(MED_INLINE,     context, seen_m, meds)
+    _harvest_inline(ALLERGY_INLINE, context, seen_a, allgs)
 
     return meds, allgs