394 lines
18 KiB
Python
394 lines
18 KiB
Python
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"})
|
|
|
|
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).
|
|
"""
|
|
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)
|