Fix: targeted retrieval for safety context + sharper extraction prompt

app.py

@@ -79,18 +79,34 @@ def _read_file(path: str) -> str:
 
 
 def _build_full_context() -> str:
-    """Concatenate ALL indexed chunks for safety checking."""
+    """
+    Build safety-check context via targeted retrieval for medications and allergies.
+    This guarantees we get the relevant chunks even if they appear late in documents.
+    """
     if not _retriever.chunks:
         return ""
-    seen = {}
+
+    # Retrieve chunks specifically about meds and allergies
+    med_chunks     = _retriever.retrieve("medications drugs prescribed dose frequency", top_k=6)
+    allergy_chunks = _retriever.retrieve("allergies allergic reactions contraindications", top_k=6)
+
+    seen: set = set()
+    ordered = []
+    for c in med_chunks + allergy_chunks:
+        uid = (c.source, c.index)
+        if uid not in seen:
+            seen.add(uid)
+            ordered.append(c)
+
+    # Also add any remaining chunks not yet included
     for c in _retriever.chunks:
-        if c.source not in seen:
-            seen[c.source] = []
-        seen[c.source].append(c.text)
-    parts = []
-    for src, texts in seen.items():
-        parts.append(f"=== {src} ===\n" + " ".join(texts[:6]))
-    return "\n\n".join(parts)
+        uid = (c.source, c.index)
+        if uid not in seen:
+            seen.add(uid)
+            ordered.append(c)
+
+    parts = [f"[{c.source}]\n{c.text}" for c in ordered]
+    return "\n\n---\n\n".join(parts)
 
 
 # ── HF Hub trace sharing ───────────────────────────────────────────────────────