from __future__ import annotations from dataclasses import dataclass from .context import pack_evidence from .models import Evidence, EvidenceDecision, RetrievalResult, SearchHit from .ports import ( ParentStore, QueryEmbeddingUnavailable, Reranker, RerankUnavailable, Retriever, ) from .sections import SectionResolver # The sections that introduce a drug: what it is, its class, its main use, its # mechanism — in book order. A bare drug name is answered from these, not from # the dosage-forms table that happens to sit near the top of the monograph. INTRO_SECTIONS = ( "ten_chung_quoc_te", "loai_thuoc", "chi_dinh", "duoc_ly_va_co_che_tac_dung", ) @dataclass(frozen=True) class EvidencePolicy: minimum_score: float = 0.12 candidate_limit: int = 5 evidence_limit: int = 3 # A free-form question about a resolved drug otherwise hands the LLM the # entire monograph; rerank trims it to the sections that actually answer. rerank_top_k: int = 6 # Token budget for the overview/rerank fallback's evidence pool (2026-08-10 # pipeline audit, priority #1). Replaces a flat evidence_limit COUNT: 3 # short chunks wastes budget a real model has, 3 long ones can silently # exceed it. Never applied to the deterministic section route — a whole # section is the answer there, and a truncated contraindication list # reads as a complete one (see `_section_hits`'s own comment). max_context_tokens: int = 6000 # symptom_to_drug: a common symptom can match far more drugs than is # useful to show in one answer. indication_candidate_limit: int = 8 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, reranker: Reranker | None = None, ) -> None: self._retriever = retriever self._parent_store = parent_store self._policy = policy or EvidencePolicy() self._section_resolver = section_resolver self._reranker = reranker def retrieve_framed( self, drug_id: str, section_key: str | None, query: str, is_overview: bool = False, ) -> RetrievalResult: """Retrieve driven by an already-understood frame, not by parsing text. The LLM understanding layer has resolved the drug (against the real catalog) and, when the turn named one, the section. So this skips the fuzzy `CatalogDrugResolver` and the keyword `SectionResolver` entirely: a named `section_key` filters that section whole; without one, this mirrors `retrieve()`'s two remaining cases — `is_overview` (the frame's `turn_type == "drug_overview"`, a bare name) answers from the identity sections only, and a free-form question reranks the full monograph down to `rerank_top_k`. Found live 2026-08-06: without the `is_overview` split, a bare drug name sent the ENTIRE ~29-section monograph as evidence for every generation call (retrieval had no notion of "just the intro"), which is both wrong retrieval and, downstream, an answer so long it intermittently failed generation/entailment outright. `query` is only the rerank signal here, never a resolution input. """ if not drug_id.strip(): return RetrievalResult(EvidenceDecision.ABSTAIN, "missing_query_or_drug") if section_key: find_by_section = getattr(self._retriever, "find_by_section", None) if find_by_section is not None: hits = find_by_section(drug_id, section_key) if hits: hits = hits + self._pooled_neighbour_hits( query, drug_id, section_key, find_by_section ) return self._decide(self._hydrate(hits, limit=None)) overview_hits = self._drug_overview(drug_id) if overview_hits is None: return RetrievalResult(EvidenceDecision.ABSTAIN, "insufficient_retrieval_score") if is_overview: intro = [h for h in overview_hits if h.document.section_key in INTRO_SECTIONS] return self._decide( self._hydrate(intro or overview_hits, limit=None), is_drug_overview=True ) overview_hits = self._rerank(query, overview_hits) # Capped even when rerank is disabled/unavailable and fails open to # the unfiltered list — an ordering aid must never remove the size # bound too, or the same 29-section explosion returns through here. # Token budget (not a flat count): rerank already put the best match # first, so packing in that order keeps as much of it as a real # model's context can hold instead of an arbitrary fixed count. hydrated = self._hydrate(overview_hits, limit=None) packed = pack_evidence(hydrated, max_tokens=self._policy.max_context_tokens) return self._decide(packed.evidence) 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. A *bare* drug name ("PARACETAMOL") # shows the whole monograph in book order. A *free-form question* about # the drug ("sốt cao uống được không?") would otherwise dump all ~29 # sections at the model; rerank keeps only the sections that answer it. overview_hits = self._drug_overview(drug_id) if overview_hits is not None: if self._is_question(query): overview_hits = self._rerank(query, overview_hits) return self._decide(self._hydrate(overview_hits, limit=None)) # A bare drug name is not a question — introduce the drug from its # identity sections (what it is, its class, its main indication), # not the whole monograph starting with the dosage-forms table. intro = [h for h in overview_hits if h.document.section_key in INTRO_SECTIONS] return self._decide( self._hydrate(intro or overview_hits, limit=None), is_drug_overview=True, ) 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(self._rerank(query, hits))) def retrieve_by_indication(self, indication_text: str) -> RetrievalResult: """Reverse lookup: symptom/indication -> candidate drugs. Keyword match first (deterministic, precise — nothing here can fabricate a drug that doesn't genuinely mention the indication). Dense-vector search over `chi_dinh` only is the fallback, tried only when the keyword pass finds nothing, to catch a paraphrase the book's own wording doesn't share. This is the one place in the live path dense search is actually used — see ADR 0008. """ if not indication_text.strip(): return RetrievalResult(EvidenceDecision.ABSTAIN, "missing_indication") find_by_indication = getattr(self._retriever, "find_by_indication", None) hits = ( find_by_indication(indication_text, self._policy.indication_candidate_limit) if find_by_indication is not None else [] ) if not hits: search_indication = getattr(self._retriever, "search_indication", None) if search_indication is not None: try: hits = search_indication( indication_text, self._policy.indication_candidate_limit ) except QueryEmbeddingUnavailable: hits = [] # Dense search always returns its nearest neighbours, even for # an indication the corpus has nothing on — verified live: a # made-up phrase still got 8 unrelated "matches". A weak top # score means those neighbours aren't really about the # question, so don't spend a generation call finding that out # the slow way; abstain here, the same bar `retrieve()`'s own # dense fallback already applies. if hits and hits[0].score < self._policy.minimum_score: hits = [] if not hits: return RetrievalResult(EvidenceDecision.ABSTAIN, "no_indication_match") return self._decide(self._hydrate(hits, limit=None)) @staticmethod def _is_question(query: str) -> bool: """A bare drug name (one or two tokens) wants the whole monograph; more than that is a question whose overview should be reranked to the point.""" return len(query.split()) > 2 def _rerank(self, query: str, hits: list[SearchHit]) -> list[SearchHit]: """Reorder hits by cross-encoder relevance, keep the top-k. Fail-open: no reranker configured, or the provider is unreachable, and the original order is returned unchanged — an ordering aid must never be able to lose an answer. The section route never reaches this. """ if self._reranker is None or len(hits) <= 1: return hits try: order = self._reranker.rerank( query, [hit.document.text for hit in hits], top_n=self._policy.rerank_top_k, ) except RerankUnavailable: return hits reranked = [hits[index] for index in order if 0 <= index < len(hits)] return reranked[: self._policy.rerank_top_k] or 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 # Bounded: a query naming a specific condition can legitimately need one # neighbouring section (than_trong -> chong_chi_dinh was the case found # live); more than this starts pooling tangential sections rather than # the one true answer, which is precision loss dressed as recall. _MAX_LEXICAL_POOLED_SECTIONS = 2 # Measured live 2026-08-10 on the real corpus: a genuine neighbour match # (chong_chi_dinh, the true positive) scored 7 matched terms; the same # query's false-positive attractor scored 6 — close enough that no # threshold alone separates them (see the exclusion below instead). 5 # keeps clearly-incidental overlap (3-4, seen on unrelated sections in # the same measurement) out while still admitting real matches. _LEXICAL_POOL_MIN_SCORE = 5.0 # `duoc_ly_va_co_che_tac_dung` is the corpus's documented false-positive # attractor (see `sections.py`'s own module docstring: it's the largest, # most generic section and "sits close to any question about the drug") # — true for embedding similarity there, and measured true for lexical # overlap here too: it scored a close second (6) right behind the real # answer (7) on the exact query that motivated this pooling mechanism. # Excluded from pooling outright rather than trusting a score margin # that isn't reliably wide enough on its own. _LEXICAL_POOL_EXCLUDED_SECTIONS = frozenset({"duoc_ly_va_co_che_tac_dung"}) # Cross-section pooling solves one measured semantic mismatch: a # "thận trọng" question whose decisive condition is filed under # contraindications. Applying it to every explicit section leaked a # lexically-overlapping interaction section into a dosage answer. Keep # this recall expansion opt-in for the one section with evidence for it; # exact dosage/interaction/contraindication routes remain exact. _LEXICAL_POOL_ENABLED_SECTIONS = frozenset({"than_trong"}) 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) hits = hits + self._pooled_neighbour_hits(query, drug_id, match.section_key, find_by_section) return hits or None def _pooled_neighbour_hits( self, query: str, drug_id: str, resolved_section: str, find_by_section ) -> list[SearchHit]: """Other sections of the SAME drug whose text lexically matches the query strongly enough to suggest the resolved section alone may not answer it — found live 2026-08-10: a "thận trọng" question about a specific condition (loét dạ dày) had its real answer filed only under "chống chỉ định" instead, a category the deterministic keyword route never considers once "thận trọng" itself matched. `search_lexical` generalizes the one hardcoded pairing this started as into a query-driven check across every section, still bounded and still whole-section (never a partial, out-of-context fragment). """ if resolved_section not in self._LEXICAL_POOL_ENABLED_SECTIONS: return [] search_lexical = getattr(self._retriever, "search_lexical", None) if search_lexical is None: return [] lexical_hits = search_lexical(query, drug_id, limit=20) pooled: list[SearchHit] = [] ineligible = {resolved_section} | self._LEXICAL_POOL_EXCLUDED_SECTIONS pooled_sections: set[str] = set() for hit in lexical_hits: section = hit.document.section_key if ( section in ineligible or section in pooled_sections or hit.score < self._LEXICAL_POOL_MIN_SCORE ): continue pooled_sections.add(section) pooled.extend(find_by_section(drug_id, section)) if len(pooled_sections) >= self._MAX_LEXICAL_POOLED_SECTIONS: break return pooled def decide(self, evidence: tuple[Evidence, ...]) -> RetrievalResult: """Public entry point for a caller that assembles its own evidence pool across several `retrieve_framed` calls — e.g. `RagAgent`'s 2-drug interaction path — and needs the same quarantine/provenance policy applied to the combined pool that a single call already gets. Bypassing this (hand-rolling `RetrievalResult(ANSWERABLE, ...)`) is exactly how the interaction path silently dropped a quarantined drug's evidence instead of surfacing VERIFY_PDF for it.""" return self._decide(evidence) def _decide( self, evidence: tuple[Evidence, ...], is_drug_overview: bool = False ) -> 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, is_drug_overview=is_drug_overview, ) 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)