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,149 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from .models import Evidence, EvidenceDecision, RetrievalResult, SearchHit
+from .ports import ParentStore, QueryEmbeddingUnavailable, Retriever
+from .sections import SectionResolver
+
+
+@dataclass(frozen=True)
+class EvidencePolicy:
+    minimum_score: float = 0.12
+    candidate_limit: int = 5
+    evidence_limit: int = 3
+
+
+class RetrievalService:
+    """Section-filtered retrieval when the question names its attribute.
+
+    Similarity is the fallback, not the default. Measured 2026-08-04, letting
+    similarity choose the section answers "chống chỉ định" correctly 1 time in
+    20, because the largest section (`duoc_ly_va_co_che_tac_dung`) sits close
+    to any question about the drug. When the question says which section it
+    wants, filtering answers it exactly.
+    """
+
+    def __init__(
+        self,
+        retriever: Retriever,
+        parent_store: ParentStore,
+        policy: EvidencePolicy | None = None,
+        section_resolver: SectionResolver | None = None,
+    ) -> None:
+        self._retriever = retriever
+        self._parent_store = parent_store
+        self._policy = policy or EvidencePolicy()
+        self._section_resolver = section_resolver
+
+    def retrieve(self, query: str, drug_id: str) -> RetrievalResult:
+        if not query.strip() or not drug_id.strip():
+            return RetrievalResult(EvidenceDecision.ABSTAIN, "missing_query_or_drug")
+
+        section_hits = self._section_hits(query, drug_id)
+        if section_hits is not None:
+            # No `evidence_limit` here: the whole section is the answer, and a
+            # truncated list of contraindications reads as a complete one.
+            return self._decide(self._hydrate(section_hits, limit=None))
+
+        # Drug resolved but no attribute named ("PARACETAMOL"): show the whole
+        # monograph, in book order, rather than dead-ending on "specify an
+        # attribute". A drug reference answers a drug name with the drug.
+        overview_hits = self._drug_overview(drug_id)
+        if overview_hits is not None:
+            return self._decide(self._hydrate(overview_hits, limit=None))
+
+        try:
+            hits = self._retriever.search(
+                query=query,
+                drug_id=drug_id,
+                limit=self._policy.candidate_limit,
+            )
+        except QueryEmbeddingUnavailable:
+            # Fail closed. The section route needs no embedder, so this only
+            # ever narrows the fallback: the caller is told nothing was found
+            # rather than being shown an error page or, worse, an answer built
+            # from a search that never ran.
+            return RetrievalResult(
+                EvidenceDecision.ABSTAIN, "query_embedding_unavailable"
+            )
+        if not hits or hits[0].score < self._policy.minimum_score:
+            return RetrievalResult(EvidenceDecision.ABSTAIN, "insufficient_retrieval_score")
+        return self._decide(self._hydrate(hits))
+
+    def _drug_overview(self, drug_id: str) -> list[SearchHit] | None:
+        """Every prose section of the drug, or None if the store cannot scroll."""
+        find_by_drug = getattr(self._retriever, "find_by_drug", None)
+        if find_by_drug is None:
+            return None
+        hits = find_by_drug(drug_id)
+        return hits or None
+
+    def _section_hits(self, query: str, drug_id: str) -> list[SearchHit] | None:
+        """Hits for an explicitly named section, or None to fall back.
+
+        Returns None — not an empty list — when this route does not apply, so
+        "no section named" stays distinguishable from "section named but empty".
+        """
+        if self._section_resolver is None:
+            return None
+        find_by_section = getattr(self._retriever, "find_by_section", None)
+        if find_by_section is None:
+            return None
+        match = self._section_resolver.resolve(query)
+        if match is None:
+            return None
+        hits = find_by_section(drug_id, match.section_key)
+        return hits or None
+
+    def _decide(self, evidence: tuple[Evidence, ...]) -> RetrievalResult:
+        if not evidence:
+            return RetrievalResult(EvidenceDecision.ABSTAIN, "parent_hydration_failed")
+        if any(not item.source_refs for item in evidence):
+            return RetrievalResult(EvidenceDecision.ABSTAIN, "missing_provenance")
+        if any(item.requires_visual_check for item in evidence):
+            return RetrievalResult(EvidenceDecision.VERIFY_PDF, "visual_verification_required", evidence)
+        return RetrievalResult(EvidenceDecision.ANSWERABLE, "grounded_evidence_available", evidence)
+
+    def _hydrate(
+        self, hits: list[SearchHit], limit: int | None = -1
+    ) -> tuple[Evidence, ...]:
+        output: list[Evidence] = []
+        seen: set[str] = set()
+        for hit in hits:
+            document = hit.document
+            evidence_id = document.parent_id or document.doc_id
+            if evidence_id in seen:
+                continue
+            seen.add(evidence_id)
+            if document.parent_id:
+                parent = self._parent_store.get(document.parent_id)
+                if parent is None:
+                    continue
+                output.append(Evidence(
+                    evidence_id=parent.parent_id,
+                    matched_doc_id=document.doc_id,
+                    kind=parent.kind,
+                    text=parent.text,
+                    score=hit.score,
+                    source_refs=parent.source_refs,
+                    hydrated_from_parent=True,
+                    requires_visual_check=(
+                        document.requires_visual_check or parent.requires_visual_check
+                    ),
+                ))
+            else:
+                output.append(Evidence(
+                    evidence_id=document.doc_id,
+                    matched_doc_id=document.doc_id,
+                    kind=document.kind,
+                    text=document.text,
+                    score=hit.score,
+                    source_refs=document.source_refs,
+                    hydrated_from_parent=False,
+                    requires_visual_check=document.requires_visual_check,
+                ))
+            cap = self._policy.evidence_limit if limit == -1 else limit
+            if cap is not None and len(output) >= cap:
+                break
+        return tuple(output)