Wire the guarded conversational RAG answer layer end-to-end

apps/ai-service/rag/artifacts.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+from .models import ParentDocument, RetrievalDocument, SourceRef
+
+
+def _read_jsonl(path: Path) -> list[dict]:
+    with path.open(encoding="utf-8") as handle:
+        return [json.loads(line) for line in handle if line.strip()]
+
+
+def _source_ref(raw: dict) -> SourceRef:
+    bbox = raw.get("bbox")
+    page_range = raw.get("page_range")
+    printed_page_range = raw.get("printed_page_range")
+    return SourceRef(
+        physical_page=int(raw["physical_page"]),
+        precision=raw["precision"],
+        block_id=raw.get("block_id"),
+        bbox=tuple(bbox) if bbox else None,
+        source_crop=raw.get("source_crop"),
+        page_range=tuple(page_range) if page_range else None,
+        printed_page=raw.get("printed_page"),
+        printed_page_range=(
+            tuple(printed_page_range) if printed_page_range else None
+        ),
+    )
+
+
+def load_documents(path: Path) -> list[RetrievalDocument]:
+    documents = []
+    for raw in _read_jsonl(path):
+        documents.append(RetrievalDocument(
+            doc_id=raw["doc_id"],
+            drug_id=raw["drug_id"],
+            kind=raw["kind"],
+            text=raw["text"],
+            section_key=raw["section_key"],
+            source_refs=tuple(_source_ref(item) for item in raw.get("source_refs", [])),
+            parent_id=raw.get("parent_id"),
+            requires_visual_check=raw.get("requires_visual_check", False),
+            drug_name=raw.get("drug_name"),
+        ))
+    return documents
+
+
+def load_parents(path: Path) -> list[ParentDocument]:
+    parents = []
+    for raw in _read_jsonl(path):
+        parents.append(ParentDocument(
+            parent_id=raw["logical_table_id"],
+            kind=raw["kind"],
+            text=raw["markdown"],
+            source_refs=tuple(_source_ref(item) for item in raw.get("source_refs", [])),
+            requires_visual_check=raw.get("requires_visual_check", False),
+        ))
+    return parents
+
+
+def load_aliases(path: Path | None) -> dict[str, set[str]]:
+    if path is None:
+        return {}
+    raw = json.loads(path.read_text(encoding="utf-8"))
+    if isinstance(raw, dict) and "entities" in raw:
+        return {
+            entity["drug_id"]: set(entity["aliases"])
+            for entity in raw["entities"]
+        }
+    return {drug_id: set(aliases) for drug_id, aliases in raw.items()}
+
+
+def build_drug_catalog(
+    documents: list[RetrievalDocument],
+    extra_aliases: dict[str, set[str]] | None = None,
+) -> dict[str, set[str]]:
+    catalog: dict[str, set[str]] = {}
+    for document in documents:
+        aliases = catalog.setdefault(document.drug_id, set())
+        aliases.add(document.drug_id.replace("_", " "))
+        if document.drug_name:
+            aliases.add(document.drug_name)
+    for drug_id, aliases in (extra_aliases or {}).items():
+        catalog.setdefault(drug_id, set()).update(aliases)
+    return catalog