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

2026-08-05 14:33:13 +07:00
parent 834d9e51b0
commit ef08b4929e
127 changed files with 37921 additions and 169 deletions
@@ -0,0 +1,218 @@
+"""Section routing: the fix for hit@1 0.05 on `chong_chi_dinh`.
+
+Measured 2026-08-04 on the real Cohere collection, letting vector similarity
+choose the section answered contraindication questions correctly 1 time in 20.
+These tests pin the two properties that make filtering safe: the longer phrase
+always wins, and an unrecognised question routes nowhere rather than guessing.
+"""
+from __future__ import annotations
+
+from rag.models import (
+    EvidenceDecision,
+    RetrievalDocument,
+    SearchHit,
+    SourceRef,
+)
+from rag.in_memory import InMemoryParentStore
+from rag.sections import SectionResolver
+from rag.service import EvidencePolicy, RetrievalService
+
+SOURCE = SourceRef(physical_page=200, precision="page", printed_page=142)
+
+
+def _doc(doc_id: str, section_key: str, text: str) -> RetrievalDocument:
+    return RetrievalDocument(
+        doc_id=doc_id,
+        parent_id=None,
+        drug_id="aspirin",
+        kind="prose",
+        section_key=section_key,
+        text=text,
+        source_refs=(SOURCE,),
+        requires_visual_check=False,
+    )
+
+
+class SectionAwareRetriever:
+    """Fake that records which route the service actually took."""
+
+    def __init__(self, docs: list[RetrievalDocument]) -> None:
+        self._docs = docs
+        self.search_calls: list[str] = []
+        self.section_calls: list[tuple[str, str]] = []
+
+    def search(
+        self, query: str, drug_id: str, limit: int  # noqa: ARG002 — Retriever protocol
+    ) -> list[SearchHit]:
+        self.search_calls.append(query)
+        # Deliberately wrong on purpose: the whole point is that the section
+        # route must not consult similarity at all.
+        return [SearchHit(self._docs[-1], 0.99)]
+
+    def find_by_section(self, drug_id: str, section_key: str) -> list[SearchHit]:
+        self.section_calls.append((drug_id, section_key))
+        return [
+            SearchHit(doc, 1.0) for doc in self._docs if doc.section_key == section_key
+        ]
+
+
+class SimilarityOnlyRetriever:
+    def __init__(self, docs: list[RetrievalDocument]) -> None:
+        self._docs = docs
+        self.search_calls: list[str] = []
+
+    def search(
+        self, query: str, drug_id: str, limit: int  # noqa: ARG002 — Retriever protocol
+    ) -> list[SearchHit]:
+        self.search_calls.append(query)
+        return [SearchHit(self._docs[0], 0.99)]
+
+
+CONTRA = [
+    _doc("c1", "chong_chi_dinh", "Mẫn cảm với aspirin."),
+    _doc("c2", "chong_chi_dinh", "Loét dạ dày tá tràng đang tiến triển."),
+    _doc("c3", "chong_chi_dinh", "Hen do aspirin."),
+    _doc("c4", "chong_chi_dinh", "Suy gan nặng."),
+    _doc("c5", "chong_chi_dinh", "Trẻ em dưới 16 tuổi có sốt virus."),
+]
+INDICATION = [_doc("i1", "chi_dinh", "Giảm đau, hạ sốt, chống viêm.")]
+PHARMACOLOGY = [_doc("p1", "duoc_ly_va_co_che_tac_dung", "Ức chế cyclooxygenase.")]
+ALL_DOCS = CONTRA + INDICATION + PHARMACOLOGY
+
+
+def _service(retriever, resolver: SectionResolver | None) -> RetrievalService:
+    return RetrievalService(
+        retriever,
+        InMemoryParentStore([]),
+        EvidencePolicy(minimum_score=0.01),
+        section_resolver=resolver,
+    )
+
+
+class TestSectionResolver:
+    def test_contraindication_is_never_read_as_indication(self) -> None:
+        """The one that measured 0.05. "chống chỉ định" contains "chỉ định"."""
+        resolver = SectionResolver()
+        assert resolver.resolve("Chống chỉ định của aspirin là gì?").section_key == (
+            "chong_chi_dinh"
+        )
+        assert resolver.resolve("Chỉ định của aspirin?").section_key == "chi_dinh"
+
+    def test_works_without_diacritics(self) -> None:
+        assert SectionResolver().resolve("aspirin chong chi dinh").section_key == (
+            "chong_chi_dinh"
+        )
+
+    def test_overdose_is_not_read_as_dose(self) -> None:
+        resolver = SectionResolver()
+        assert resolver.resolve("xử trí quá liều metformin").section_key == (
+            "qua_lieu_va_xu_tri"
+        )
+        assert resolver.resolve("liều dùng metformin").section_key == (
+            "lieu_luong_va_cach_dung"
+        )
+
+    def test_adr_management_is_not_read_as_adr_itself(self) -> None:
+        resolver = SectionResolver()
+        assert resolver.resolve("xử trí tác dụng phụ của prednisolon").section_key == (
+            "huong_dan_xu_tri_adr"
+        )
+        assert resolver.resolve("tác dụng phụ của prednisolon").section_key == (
+            "tac_dung_khong_mong_muon"
+        )
+
+    def test_incompatibility_is_not_read_as_interaction(self) -> None:
+        resolver = SectionResolver()
+        assert resolver.resolve("tương kỵ của ceftriaxon").section_key == "tuong_ky"
+        assert resolver.resolve("tương tác của ceftriaxon").section_key == (
+            "tuong_tac_thuoc"
+        )
+
+    def test_bare_lieu_resolves_without_capturing_overdose(self) -> None:
+        """Found by testing on human-written golden questions, not templates.
+
+        4 of 16 said just "Liều Metformin cho người lớn?". Adding bare "liều"
+        is only safe because "quá liều" is longer and is tested first.
+        """
+        resolver = SectionResolver()
+        assert resolver.resolve("Liều Metformin cho người lớn?").section_key == (
+            "lieu_luong_va_cach_dung"
+        )
+        assert resolver.resolve("quá liều paracetamol").section_key == (
+            "qua_lieu_va_xu_tri"
+        )
+
+    def test_colloquial_pregnancy_phrasing(self) -> None:
+        assert SectionResolver().resolve(
+            "Bà bầu dùng Ibuprofen được không?"
+        ).section_key == "thoi_ky_mang_thai"
+
+    def test_unrecognised_question_routes_nowhere(self) -> None:
+        """No match must not become a guess — the caller falls back."""
+        assert SectionResolver().resolve("thuốc này giá bao nhiêu") is None
+        assert SectionResolver().resolve("") is None
+
+    def test_new_section_needs_no_code_change(self) -> None:
+        resolver = SectionResolver({"invented_section": ("một mục hoàn toàn mới",)})
+        assert resolver.resolve("hỏi về một mục hoàn toàn mới").section_key == (
+            "invented_section"
+        )
+
+
+class TestSectionRouting:
+    def test_named_section_bypasses_similarity_entirely(self) -> None:
+        retriever = SectionAwareRetriever(ALL_DOCS)
+        result = _service(retriever, SectionResolver()).retrieve(
+            "Chống chỉ định của aspirin là gì?", "aspirin"
+        )
+        assert retriever.section_calls == [("aspirin", "chong_chi_dinh")]
+        assert retriever.search_calls == []
+        assert result.decision == EvidenceDecision.ANSWERABLE
+        assert {item.evidence_id for item in result.evidence} == {
+            "c1", "c2", "c3", "c4", "c5",
+        }
+
+    def test_whole_section_is_returned_past_the_evidence_limit(self) -> None:
+        """Five contraindications must not arrive as three."""
+        retriever = SectionAwareRetriever(ALL_DOCS)
+        service = RetrievalService(
+            retriever,
+            InMemoryParentStore([]),
+            EvidencePolicy(minimum_score=0.01, evidence_limit=3),
+            section_resolver=SectionResolver(),
+        )
+        result = service.retrieve("chống chỉ định aspirin", "aspirin")
+        assert len(result.evidence) == 5
+
+    def test_unnamed_section_falls_back_to_similarity(self) -> None:
+        retriever = SectionAwareRetriever(ALL_DOCS)
+        result = _service(retriever, SectionResolver()).retrieve(
+            "aspirin dùng cho bệnh nhân này thế nào", "aspirin"
+        )
+        assert retriever.section_calls == []
+        assert retriever.search_calls
+        assert result.decision == EvidenceDecision.ANSWERABLE
+
+    def test_retriever_without_the_capability_still_works(self) -> None:
+        retriever = SimilarityOnlyRetriever(ALL_DOCS)
+        result = _service(retriever, SectionResolver()).retrieve(
+            "Chống chỉ định của aspirin là gì?", "aspirin"
+        )
+        assert retriever.search_calls
+        assert result.decision == EvidenceDecision.ANSWERABLE
+
+    def test_no_resolver_keeps_the_old_behaviour(self) -> None:
+        retriever = SectionAwareRetriever(ALL_DOCS)
+        _service(retriever, None).retrieve("chống chỉ định aspirin", "aspirin")
+        assert retriever.section_calls == []
+        assert retriever.search_calls
+
+    def test_named_but_empty_section_falls_back(self) -> None:
+        """A drug with no such section must not abstain — similarity still tries."""
+        retriever = SectionAwareRetriever(INDICATION + PHARMACOLOGY)
+        result = _service(retriever, SectionResolver()).retrieve(
+            "chống chỉ định aspirin", "aspirin"
+        )
+        assert retriever.section_calls == [("aspirin", "chong_chi_dinh")]
+        assert retriever.search_calls
+        assert result.decision == EvidenceDecision.ANSWERABLE