150 lines
6.3 KiB
Python
150 lines
6.3 KiB
Python
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)
|