from __future__ import annotations from dataclasses import dataclass from .clinical import ( CandidateStatus, MedicationCandidateAssessment, PatientContext, ) from .context import pack_evidence from .models import Evidence, EvidenceDecision, RetrievalResult, SearchHit from .ports import ( ParentStore, QueryEmbeddingUnavailable, Reranker, RerankUnavailable, Retriever, ) from .sections import SectionResolver from .text import normalize_name # 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 # Retrieve a wider chunk pool before grouping/ranking at drug level. The # final candidate cap above is applied only after entity aggregation. indication_retrieval_limit: int = 40 indication_evidence_per_drug: int = 2 # Patient-specific stage 2 is intentionally narrower than a general list. patient_candidate_limit: int = 2 safety_hits_per_section: int = 1 safety_sections_per_candidate: int = 4 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_retrieval_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_retrieval_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") groups = self._rank_indication_drugs(indication_text, hits) selected_hits = [ hit for group in groups[: self._policy.indication_candidate_limit] for hit in group[: self._policy.indication_evidence_per_drug] ] return self._decide(self._hydrate(selected_hits, limit=None)) def assess_patient_candidates( self, indication_result: RetrievalResult, patient: PatientContext, ) -> tuple[RetrievalResult, tuple[MedicationCandidateAssessment, ...]]: """Targeted stage-2 safety evidence for already-indicated candidates. This never creates candidates. It groups the stage-1 indication evidence, keeps a bounded number of drugs, then searches only safety facets relevant to facts actually present in ``patient``. A lexical hit selects a chunk; pregnancy/breastfeeding sections are direct metadata routes because their relation is explicit in the section key. Absence of a hit is recorded as insufficient evidence, never "safe". """ if not patient.present or not indication_result.evidence: return indication_result, () indication_by_drug: dict[str, list[Evidence]] = {} for evidence in indication_result.evidence: drug_id = _evidence_drug_id(evidence) if drug_id: indication_by_drug.setdefault(drug_id, []).append(evidence) assessments: list[MedicationCandidateAssessment] = [] combined: list[Evidence] = [] for drug_id, indication_evidence in list(indication_by_drug.items())[ : self._policy.patient_candidate_limit ]: selected = self._patient_safety_evidence(drug_id, patient) safety_evidence = tuple( evidence for values in selected.values() for evidence in values ) status = ( CandidateStatus.REQUIRES_ADDITIONAL_INFORMATION if any(item.requires_visual_check for item in safety_evidence) else CandidateStatus.SUPPORTED_WITH_CAUTION if safety_evidence else CandidateStatus.INSUFFICIENT_EVIDENCE ) name = next( (item.drug_name for item in indication_evidence if item.drug_name), None, ) or drug_id.replace("_", " ").title() assessment = MedicationCandidateAssessment( drug_id=drug_id, drug_name=name, indication_supported=True, indication_evidence=tuple(indication_evidence), contraindication_evidence=tuple(selected.get("chong_chi_dinh", ())), precaution_evidence=tuple(selected.get("than_trong", ())), interaction_evidence=tuple(selected.get("tuong_tac_thuoc", ())), renal_evidence=_facet_evidence( selected, patient.renal.present, ("chong_chi_dinh", "than_trong", "lieu_luong_va_cach_dung"), ), hepatic_evidence=_facet_evidence( selected, patient.hepatic.present, ("chong_chi_dinh", "than_trong", "lieu_luong_va_cach_dung"), ), pregnancy_evidence=tuple(selected.get("thoi_ky_mang_thai", ())), breastfeeding_evidence=tuple(selected.get("thoi_ky_cho_con_bu", ())), age_evidence=_facet_evidence( selected, bool(patient.age_text), ("than_trong", "lieu_luong_va_cach_dung"), ), dose_evidence=tuple(selected.get("lieu_luong_va_cach_dung", ())), status=status, ) assessments.append(assessment) # Quarantined tables/formulas remain visible in the structured # assessment/status but never enter generation. Applying the # single-drug global VERIFY_PDF rule to a multi-candidate list # would suppress every otherwise verified prose candidate merely # because one candidate has one visual-only renal table. combined.extend( item for item in assessment.evidence if not item.requires_visual_check ) result = self._decide(tuple(combined)) if result.decision == EvidenceDecision.ANSWERABLE: result = RetrievalResult( result.decision, "grounded_patient_evidence_available", result.evidence, ) return result, tuple(assessments) def retrieve_patient_drug_context( self, drug_id: str, base_result: RetrievalResult, patient: PatientContext, ) -> RetrievalResult: """Add bounded patient-relevant facets to a named-drug lookup.""" if not patient.requires_safety_review: return base_result selected = self._patient_safety_evidence(drug_id, patient) evidence = list(base_result.evidence) seen = {item.evidence_id for item in evidence} for values in selected.values(): for item in values: if item.requires_visual_check or item.evidence_id in seen: continue seen.add(item.evidence_id) evidence.append(item) result = self._decide(tuple(evidence)) if result.decision == EvidenceDecision.ANSWERABLE: return RetrievalResult( result.decision, "grounded_patient_evidence_available", result.evidence, resolved_drug_id=base_result.resolved_drug_id, is_drug_overview=base_result.is_drug_overview, ) return result def _patient_safety_evidence( self, drug_id: str, patient: PatientContext ) -> dict[str, list[Evidence]]: selected: dict[str, list[Evidence]] = {} search_lexical = getattr(self._retriever, "search_lexical", None) find_by_section = getattr(self._retriever, "find_by_section", None) def lexical_facets( query: str, section_keys: tuple[str, ...], *, require_context_match: bool = False, ) -> dict[str, list[SearchHit]]: """Return at most the configured hits per requested relation. Each clinical facet gets its own query. In particular, a current medicine may select an interaction chunk only when that medicine matches inside the interaction section; CKD/age terms from another facet cannot make an unrelated interaction look supported. """ if not query or search_lexical is None: return {} hits = search_lexical( query, drug_id, max(20, self._policy.safety_hits_per_section * len(section_keys)), section_keys=section_keys, ) per_section: dict[str, list[SearchHit]] = {} for hit in hits: section = hit.document.section_key if section not in section_keys: continue if require_context_match and not _patient_context_matches( hit.document.text, patient ): continue bucket = per_section.setdefault(section, []) if len(bucket) < self._policy.safety_hits_per_section: bucket.append(hit) return per_section if search_lexical is not None: # These three searches deliberately keep their relations separate. # Absence of an exact lexical hit means "not evidenced in the # retrieved Dược thư text", never "no interaction/contraindication". interaction = lexical_facets( patient.interaction_query(), ("tuong_tac_thuoc",) ) warnings = lexical_facets( patient.warning_query(), ("chong_chi_dinh", "than_trong"), require_context_match=True, ) dosage = lexical_facets( patient.dosage_context_query(), ("lieu_luong_va_cach_dung",), require_context_match=True, ) candidates: list[tuple[str, list[SearchHit]]] = [] if "tuong_tac_thuoc" in interaction: candidates.append(("tuong_tac_thuoc", interaction["tuong_tac_thuoc"])) warning_sections = sorted( warnings, key=lambda section: ( section != "chong_chi_dinh", -warnings[section][0].score, section, ), ) for section in warning_sections: candidates.append((section, warnings[section])) if "lieu_luong_va_cach_dung" in dosage: candidates.append( ("lieu_luong_va_cach_dung", dosage["lieu_luong_va_cach_dung"]) ) for section, hits in candidates[ : self._policy.safety_sections_per_candidate ]: selected[section] = list(self._hydrate(hits, limit=None)) # These sections encode the patient relation themselves; no lexical # coincidence is needed to decide they are relevant. direct_sections = [] if patient.pregnancy_status: direct_sections.append("thoi_ky_mang_thai") if patient.breastfeeding is True: direct_sections.append("thoi_ky_cho_con_bu") if find_by_section is not None: for section in direct_sections: if section in selected: continue hits = find_by_section(drug_id, section) if hits: selected[section] = list(self._hydrate(hits, limit=None)) return selected def _rank_indication_drugs( self, query: str, hits: list[SearchHit] ) -> list[list[SearchHit]]: """Group and rank entities without rewarding duplicate chunks.""" by_drug: dict[str, list[SearchHit]] = {} for hit in hits: by_drug.setdefault(hit.document.drug_id, []).append(hit) groups = [ sorted(group, key=lambda hit: (-hit.score, hit.document.doc_id)) for group in by_drug.values() ] groups.sort(key=lambda group: (-group[0].score, group[0].document.drug_id)) if self._reranker is None or len(groups) <= 1: return groups documents = [ f"{group[0].document.drug_name or group[0].document.drug_id}\n" + "\n".join(hit.document.text for hit in group[:2]) for group in groups ] try: order = self._reranker.rerank( query, documents, top_n=self._policy.indication_candidate_limit, ) except RerankUnavailable: return groups ranked = [groups[index] for index in order if 0 <= index < len(groups)] return ranked or groups @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 ), drug_id=document.drug_id, drug_name=document.drug_name, section_key=document.section_key, section_title=document.section_title, source_document=document.source_document, )) 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, drug_id=document.drug_id, drug_name=document.drug_name, section_key=document.section_key, section_title=document.section_title, source_document=document.source_document, )) cap = self._policy.evidence_limit if limit == -1 else limit if cap is not None and len(output) >= cap: break return tuple(output) def _evidence_drug_id(evidence: Evidence) -> str | None: if evidence.drug_id: return evidence.drug_id if "__" in evidence.matched_doc_id: return evidence.matched_doc_id.split("__", 1)[0] return None def _facet_evidence( selected: dict[str, list[Evidence]], enabled: bool, section_keys: tuple[str, ...], ) -> tuple[Evidence, ...]: if not enabled: return () output: list[Evidence] = [] seen: set[str] = set() for section in section_keys: for evidence in selected.get(section, ()): if evidence.evidence_id in seen: continue seen.add(evidence.evidence_id) output.append(evidence) return tuple(output) def _patient_context_matches(text: str, patient: PatientContext) -> bool: """Require a clinical anchor, not overlap on generic words like 'chức năng'.""" haystack = normalize_name(text) supplied = ( *patient.comorbidities, *patient.allergies, *patient.previous_adverse_reactions, *patient.relevant_labs, ) raw_terms = [normalize_name(term) for term in supplied if term.strip()] if any(term in haystack for term in raw_terms if len(term) >= 3): return True if patient.renal.present and any( term in haystack for term in ( "suy than", "chuc nang than", "than nang", "creatinin", "crcl", "egfr", "loc cau than", "do thanh thai", ) ): return True if patient.hepatic.present and any( term in haystack for term in ( "suy gan", "chuc nang gan", "benh gan", "xo gan", "child pugh", "ast", "alt", "bilirubin", ) ): return True if patient.age_text and any( term in haystack for term in ("nguoi cao tuoi", "cao tuoi", "tre em", "tre so sinh") ): return True return False