482 lines
18 KiB
Python
482 lines
18 KiB
Python
from pathlib import Path
|
|
|
|
from rag.artifacts import load_aliases
|
|
from rag.evaluation import CaseOrigin, EvaluationCase, EvaluationOutcome, summarize
|
|
from rag.in_memory import InMemoryLexicalRetriever, InMemoryParentStore, _char_ngrams
|
|
from rag.models import (
|
|
EvidenceDecision,
|
|
ParentDocument,
|
|
QueryIntent,
|
|
RetrievalDocument,
|
|
SearchHit,
|
|
SourceRef,
|
|
SubjectScope,
|
|
)
|
|
from rag.routing import (
|
|
CatalogDrugResolver,
|
|
DrugResolutionStatus,
|
|
QueryRoutingService,
|
|
)
|
|
from rag.service import EvidencePolicy, RetrievalService
|
|
|
|
SOURCE = SourceRef(
|
|
physical_page=112,
|
|
precision="region",
|
|
block_id="p112_t0",
|
|
bbox=(1, 2, 3, 4),
|
|
source_crop="crops/p112_t0.png",
|
|
)
|
|
|
|
VERIFIED_ENTITIES = (
|
|
Path(__file__).parents[3] / "ingestion/data/verified/drug_entities.json"
|
|
)
|
|
|
|
|
|
def table_service(*, visual: bool = False) -> RetrievalService:
|
|
row = RetrievalDocument(
|
|
doc_id="p112_t0::row::0",
|
|
parent_id="p112_t0",
|
|
drug_id="acetylcystein",
|
|
kind="table_row",
|
|
section_key="lieu_luong_va_cach_dung",
|
|
text="ACETYLCYSTEIN thể trọng 40 đến 49 kg thể tích 34 ml",
|
|
source_refs=(SOURCE,),
|
|
requires_visual_check=visual,
|
|
)
|
|
parent = ParentDocument(
|
|
parent_id="p112_t0",
|
|
kind="table",
|
|
text="| Thể trọng | Thể tích |\n| 40 - 49 kg | 34 ml |",
|
|
source_refs=(SOURCE,),
|
|
)
|
|
return RetrievalService(
|
|
InMemoryLexicalRetriever([row]),
|
|
InMemoryParentStore([parent]),
|
|
EvidencePolicy(minimum_score=0.01),
|
|
)
|
|
|
|
|
|
class _OverviewRetriever:
|
|
"""A fake with `find_by_drug`/`find_by_section` (the Qdrant adapter's
|
|
shape) — `InMemoryLexicalRetriever` doesn't implement either, so
|
|
`retrieve_framed`'s overview path is otherwise untestable.
|
|
|
|
`search_lexical` is scripted per test (`lexical_hits`), not real text
|
|
matching — this file is about `retrieve_framed`'s own wiring, not the
|
|
scorer (see `test_qdrant_adapter.py`/`test_section_routing.py` for
|
|
that)."""
|
|
|
|
def __init__(
|
|
self, documents: list[RetrievalDocument], lexical_hits: list[SearchHit] = ()
|
|
) -> None:
|
|
self._documents = documents
|
|
self._lexical_hits = lexical_hits
|
|
self.lexical_calls: list[tuple[str, str]] = []
|
|
|
|
def find_by_drug(self, drug_id: str) -> list[SearchHit]:
|
|
return [
|
|
SearchHit(document=d, score=1.0)
|
|
for d in self._documents if d.drug_id == drug_id
|
|
]
|
|
|
|
def find_by_section(self, drug_id: str, section_key: str) -> list[SearchHit]:
|
|
return [
|
|
SearchHit(document=d, score=1.0)
|
|
for d in self._documents
|
|
if d.drug_id == drug_id and d.section_key == section_key
|
|
]
|
|
|
|
def search(self, query: str, drug_id: str, limit: int) -> list[SearchHit]:
|
|
return []
|
|
|
|
def search_lexical(self, query: str, drug_id: str, limit: int) -> list[SearchHit]: # noqa: ARG002
|
|
self.lexical_calls.append((query, drug_id))
|
|
return list(self._lexical_hits)
|
|
|
|
|
|
_MONOGRAPH_SECTIONS = (
|
|
"ten_chung_quoc_te", "ma_atc", "loai_thuoc", "dang_thuoc_va_ham_luong",
|
|
"duoc_ly_va_co_che_tac_dung", "chi_dinh", "chong_chi_dinh", "than_trong",
|
|
"tac_dung_khong_mong_muon", "lieu_luong_va_cach_dung", "tuong_tac_thuoc",
|
|
"qua_lieu_va_xu_tri", "do_on_dinh_va_bao_quan", "thong_tin_quy_che",
|
|
)
|
|
|
|
|
|
def _monograph_service(max_context_tokens: int = 6000) -> RetrievalService:
|
|
documents = [
|
|
RetrievalDocument(
|
|
doc_id=f"paracetamol::{section}::0", drug_id="paracetamol",
|
|
kind="prose", section_key=section,
|
|
text=f"Nội dung mục {section}.", source_refs=(SOURCE,),
|
|
)
|
|
for section in _MONOGRAPH_SECTIONS
|
|
]
|
|
return RetrievalService(
|
|
_OverviewRetriever(documents), InMemoryParentStore([]),
|
|
EvidencePolicy(evidence_limit=3, max_context_tokens=max_context_tokens),
|
|
)
|
|
|
|
|
|
def test_retrieve_framed_overview_answers_from_intro_sections_only():
|
|
"""Found live 2026-08-06: a bare drug name sent all 14+ sections of the
|
|
monograph as evidence, producing an answer long enough to intermittently
|
|
fail generation/entailment. `is_overview=True` must narrow this the same
|
|
way `retrieve()`'s bare-name branch always has."""
|
|
result = _monograph_service().retrieve_framed(
|
|
"paracetamol", None, "paracetamol", is_overview=True
|
|
)
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
assert result.is_drug_overview is True
|
|
returned_sections = {e.matched_doc_id.split("::")[1] for e in result.evidence}
|
|
assert returned_sections <= {
|
|
"ten_chung_quoc_te", "loai_thuoc", "chi_dinh", "duoc_ly_va_co_che_tac_dung",
|
|
}
|
|
assert len(result.evidence) < len(_MONOGRAPH_SECTIONS)
|
|
|
|
|
|
def test_retrieve_framed_question_without_section_is_capped_even_without_rerank():
|
|
# No reranker configured: `_rerank` fails open and returns everything
|
|
# unfiltered. Hydration must still bound it by TOKEN budget (2026-08-10:
|
|
# was a flat evidence_limit count, now pack_evidence) — an ordering aid
|
|
# failing open must not also remove the size cap.
|
|
result = _monograph_service(max_context_tokens=50).retrieve_framed(
|
|
"paracetamol", None, "thuốc này có tác dụng phụ gì", is_overview=False
|
|
)
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
assert len(result.evidence) < len(_MONOGRAPH_SECTIONS)
|
|
|
|
|
|
def test_retrieve_framed_packs_overview_by_token_budget_not_flat_count():
|
|
"""2026-08-10 pipeline audit priority #1: a tight token budget can admit
|
|
FEWER than the old flat evidence_limit=3 when blocks are long, and a
|
|
generous one can admit MORE when blocks are short — proving this is
|
|
genuinely token-driven, not a renamed count cap."""
|
|
tiny_budget_result = _monograph_service(max_context_tokens=20).retrieve_framed(
|
|
"paracetamol", None, "thuốc này có tác dụng phụ gì", is_overview=False
|
|
)
|
|
generous_budget_result = _monograph_service(max_context_tokens=6000).retrieve_framed(
|
|
"paracetamol", None, "thuốc này có tác dụng phụ gì", is_overview=False
|
|
)
|
|
assert len(tiny_budget_result.evidence) < 3
|
|
assert len(generous_budget_result.evidence) > 3
|
|
|
|
|
|
def test_retrieve_framed_pools_lexically_strong_neighbour_section():
|
|
"""The LIVE agent path (`RagAgent` -> understanding -> `retrieve_framed`)
|
|
must get the same neighbour-pooling `retrieve()` does — found live
|
|
2026-08-10 that the first version of this fix only wired into
|
|
`retrieve()`, which the real HTTP request path does not call at all;
|
|
`retrieve_framed` has its own separate `if section_key:` branch."""
|
|
documents = [
|
|
RetrievalDocument(
|
|
doc_id=f"aspirin::{section}::0", drug_id="aspirin",
|
|
kind="prose", section_key=section,
|
|
text=f"Nội dung mục {section}.", source_refs=(SOURCE,),
|
|
)
|
|
for section in ("than_trong", "chong_chi_dinh")
|
|
]
|
|
chong_chi_dinh_hit = SearchHit(
|
|
document=next(d for d in documents if d.section_key == "chong_chi_dinh"),
|
|
score=1.0,
|
|
)
|
|
retriever = _OverviewRetriever(documents, lexical_hits=[
|
|
SearchHit(document=chong_chi_dinh_hit.document, score=7.0),
|
|
])
|
|
service = RetrievalService(retriever, InMemoryParentStore([]), EvidencePolicy())
|
|
|
|
result = service.retrieve_framed(
|
|
"aspirin", "than_trong",
|
|
"Thận trọng khi dùng aspirin cho bệnh nhân loét dạ dày là gì?",
|
|
)
|
|
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
returned_sections = {e.matched_doc_id.split("::")[1] for e in result.evidence}
|
|
assert returned_sections == {"than_trong", "chong_chi_dinh"}
|
|
assert retriever.lexical_calls == [
|
|
("Thận trọng khi dùng aspirin cho bệnh nhân loét dạ dày là gì?", "aspirin"),
|
|
]
|
|
|
|
|
|
def test_row_hit_hydrates_complete_parent_and_keeps_citation():
|
|
result = table_service().retrieve("acetylcystein 45 kg bao nhiêu ml", "acetylcystein")
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
assert result.evidence[0].hydrated_from_parent is True
|
|
assert result.evidence[0].text.startswith("| Thể trọng")
|
|
assert result.evidence[0].source_refs == (SOURCE,)
|
|
|
|
|
|
def test_visual_risk_routes_to_pdf_verifier():
|
|
result = table_service(visual=True).retrieve(
|
|
"acetylcystein 45 kg bao nhiêu ml", "acetylcystein",
|
|
)
|
|
assert result.decision == EvidenceDecision.VERIFY_PDF
|
|
assert result.reason == "visual_verification_required"
|
|
|
|
|
|
def test_missing_parent_abstains_instead_of_answering_from_row_fragment():
|
|
row = RetrievalDocument(
|
|
doc_id="row", parent_id="missing", drug_id="drug", kind="table_row",
|
|
section_key="dose", text="drug dose 10 mg", source_refs=(SOURCE,),
|
|
)
|
|
service = RetrievalService(
|
|
InMemoryLexicalRetriever([row]), InMemoryParentStore([]),
|
|
EvidencePolicy(minimum_score=0.01),
|
|
)
|
|
result = service.retrieve("drug dose", "drug")
|
|
assert result.decision == EvidenceDecision.ABSTAIN
|
|
assert result.reason == "parent_hydration_failed"
|
|
|
|
|
|
def test_missing_provenance_abstains():
|
|
document = RetrievalDocument(
|
|
doc_id="prose", drug_id="drug", kind="prose", section_key="dose",
|
|
text="drug dose 10 mg", source_refs=(),
|
|
)
|
|
service = RetrievalService(
|
|
InMemoryLexicalRetriever([document]), InMemoryParentStore([]),
|
|
EvidencePolicy(minimum_score=0.01),
|
|
)
|
|
result = service.retrieve("drug dose", "drug")
|
|
assert result.decision == EvidenceDecision.ABSTAIN
|
|
assert result.reason == "missing_provenance"
|
|
|
|
|
|
class FixedRetriever:
|
|
def __init__(self, hits: list[SearchHit]) -> None:
|
|
self._hits = hits
|
|
|
|
def search(self, query: str, drug_id: str, limit: int) -> list[SearchHit]:
|
|
del query, drug_id
|
|
return self._hits[:limit]
|
|
|
|
|
|
def test_near_tied_different_sources_are_returned_for_evidence_grading():
|
|
first = RetrievalDocument("a", "drug", "prose", "A", "dose", (SOURCE,))
|
|
second = RetrievalDocument("b", "drug", "prose", "B", "dose", (SOURCE,))
|
|
service = RetrievalService(
|
|
FixedRetriever([SearchHit(first, 0.50), SearchHit(second, 0.495)]),
|
|
InMemoryParentStore([]),
|
|
)
|
|
result = service.retrieve("dose", "drug")
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
assert [item.evidence_id for item in result.evidence] == ["a", "b"]
|
|
|
|
|
|
def test_source_derived_cases_do_not_inflate_release_gate_metric():
|
|
outcomes = [
|
|
EvaluationOutcome(
|
|
EvaluationCase(
|
|
"expert-1", "q", "drug", "right", CaseOrigin.EXPERT,
|
|
SubjectScope.HUMAN,
|
|
),
|
|
("wrong",),
|
|
),
|
|
EvaluationOutcome(
|
|
EvaluationCase(
|
|
"generated-1", "q", "drug", "right", CaseOrigin.SOURCE_DERIVED,
|
|
SubjectScope.HUMAN,
|
|
),
|
|
("right",),
|
|
),
|
|
]
|
|
report = summarize(outcomes)
|
|
assert report["expert_release_gate"]["recall_at_1"] == 0.0
|
|
assert report["source_derived_diagnostic"]["recall_at_1"] == 1.0
|
|
assert report["manual_routing_diagnostic"]["cases"] == 0
|
|
|
|
|
|
def test_character_ngrams_preserve_word_order():
|
|
assert _char_ngrams("beta alpha") != _char_ngrams("alpha beta")
|
|
|
|
|
|
def test_drug_resolver_handles_a_typo_without_fixture_drug_id():
|
|
resolver = CatalogDrugResolver({"famciclovir": {"famciclovir"}})
|
|
result = resolver.resolve("famciclovia chỉnh liều khi ClCr 20")
|
|
assert result.status == DrugResolutionStatus.RESOLVED
|
|
assert result.drug_id == "famciclovir"
|
|
|
|
|
|
def test_drug_resolver_does_not_guess_when_query_mentions_two_drugs():
|
|
resolver = CatalogDrugResolver({
|
|
"oresol": {"oresol"},
|
|
"natri_clorid": {"natri clorid"},
|
|
})
|
|
result = resolver.resolve("oresol có bao nhiêu natri clorid")
|
|
assert result.status == DrugResolutionStatus.AMBIGUOUS
|
|
|
|
|
|
def test_verified_aliases_reach_common_parenthesized_drug_names():
|
|
resolver = CatalogDrugResolver(load_aliases(VERIFIED_ENTITIES))
|
|
assert resolver.resolve("Liều paracetamol cho người lớn").drug_id == (
|
|
"paracetamol_acetaminophen"
|
|
)
|
|
assert resolver.resolve("Chống chỉ định aspirin").drug_id == (
|
|
"acid_acetylsalicylic_aspirin"
|
|
)
|
|
assert resolver.resolve("Công thức oresol").drug_id == (
|
|
"thuoc_uong_bu_nuoc_va_ien_giai"
|
|
)
|
|
|
|
|
|
def test_verified_catalog_protects_canonical_substring_traps():
|
|
resolver = CatalogDrugResolver(load_aliases(VERIFIED_ENTITIES))
|
|
traps = {
|
|
"homatropin hydrobromid": "homatropin_hydrobromid",
|
|
"hydroclorothiazid": "hydroclorothiazid",
|
|
"flucloxacilin": "flucloxacilin",
|
|
"pseudoephedrin": "pseudoephedrin",
|
|
"ethinylestradiol": "ethinylestradiol",
|
|
"desloratadin": "desloratadin",
|
|
"ciprofloxacin": "ciprofloxacin",
|
|
"levofloxacin": "levofloxacin",
|
|
"esomeprazol": "esomeprazol",
|
|
"methylprednisolon": "methylprednisolon",
|
|
"medroxyprogesteron acetat": "medroxyprogesteron_acetat",
|
|
"methyltestosteron": "methyltestosteron",
|
|
"oxytetracyclin": "oxytetracyclin",
|
|
}
|
|
for query, expected_id in traps.items():
|
|
result = resolver.resolve(query)
|
|
assert result.status == DrugResolutionStatus.RESOLVED
|
|
assert result.drug_id == expected_id
|
|
|
|
|
|
def test_asymmetric_evidence_resolves_subject_and_component():
|
|
ors = RetrievalDocument(
|
|
doc_id="ors", drug_id="ors", kind="prose", section_key="formula",
|
|
text="Oresol chứa natri clorid", source_refs=(SOURCE,),
|
|
)
|
|
sodium = RetrievalDocument(
|
|
doc_id="sodium", drug_id="sodium", kind="prose", section_key="dose",
|
|
text="Natri clorid dùng đường truyền", source_refs=(SOURCE,),
|
|
)
|
|
routed = QueryRoutingService(
|
|
RetrievalService(
|
|
InMemoryLexicalRetriever([ors, sodium]), InMemoryParentStore([]),
|
|
EvidencePolicy(minimum_score=0.01),
|
|
),
|
|
CatalogDrugResolver({"ors": {"oresol"}, "sodium": {"natri clorid"}}),
|
|
)
|
|
result = routed.retrieve(
|
|
"Oresol có bao nhiêu natri clorid?",
|
|
SubjectScope.HUMAN,
|
|
QueryIntent.FACT_LOOKUP,
|
|
)
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
assert result.resolved_drug_id == "ors"
|
|
|
|
|
|
def test_structured_scope_fails_closed_and_rejects_non_human_subject():
|
|
document = RetrievalDocument(
|
|
doc_id="dose", drug_id="famciclovir", drug_name="FAMCICLOVIR",
|
|
kind="prose", text="Famciclovir liều cho người lớn", section_key="dose",
|
|
source_refs=(SOURCE,),
|
|
)
|
|
routed = QueryRoutingService(
|
|
RetrievalService(
|
|
InMemoryLexicalRetriever([document]), InMemoryParentStore([]),
|
|
EvidencePolicy(minimum_score=0.01),
|
|
),
|
|
CatalogDrugResolver({"famciclovir": {"famciclovir"}}),
|
|
)
|
|
veterinary = routed.retrieve(
|
|
"Liều famciclovir cho mèo", SubjectScope.NON_HUMAN,
|
|
)
|
|
unknown = routed.retrieve("Liều famciclovir")
|
|
adult = routed.retrieve(
|
|
"Liều famciclovir cho người lớn", SubjectScope.HUMAN,
|
|
QueryIntent.FACT_LOOKUP,
|
|
)
|
|
assert veterinary.decision == EvidenceDecision.ABSTAIN
|
|
assert veterinary.reason == "out_of_scope_non_human"
|
|
assert unknown.decision == EvidenceDecision.ABSTAIN
|
|
assert unknown.reason == "subject_scope_unknown"
|
|
assert adult.decision == EvidenceDecision.ANSWERABLE
|
|
assert adult.resolved_drug_id == "famciclovir"
|
|
|
|
|
|
def test_recommendation_intent_is_refused_at_policy_boundary():
|
|
routed = QueryRoutingService(
|
|
table_service(), CatalogDrugResolver({"drug": {"drug"}}),
|
|
)
|
|
result = routed.retrieve(
|
|
"Nên dùng drug nào?", SubjectScope.HUMAN, QueryIntent.RECOMMENDATION,
|
|
)
|
|
assert result.decision == EvidenceDecision.ABSTAIN
|
|
assert result.reason == "recommendation_out_of_scope"
|
|
|
|
|
|
class _IndicationRetriever:
|
|
"""A fake exposing only `find_by_indication`/`search_indication` (the
|
|
Qdrant adapter's shape for the reverse-lookup path), so
|
|
`retrieve_by_indication`'s own orchestration — keyword first, dense
|
|
fallback only when keyword finds nothing — is what's under test here,
|
|
not the matching algorithm itself (that's `test_qdrant_adapter.py`'s job)."""
|
|
|
|
def __init__(
|
|
self,
|
|
keyword_hits: list[SearchHit] | None = None,
|
|
dense_hits: list[SearchHit] | None = None,
|
|
) -> None:
|
|
self._keyword_hits = keyword_hits or []
|
|
self._dense_hits = dense_hits or []
|
|
self.dense_called = False
|
|
|
|
def find_by_indication(self, indication_text, limit): # noqa: ARG002
|
|
return self._keyword_hits
|
|
|
|
def search_indication(self, query, limit): # noqa: ARG002
|
|
self.dense_called = True
|
|
return self._dense_hits
|
|
|
|
|
|
def _indication_hit(drug_id: str) -> SearchHit:
|
|
return SearchHit(
|
|
document=RetrievalDocument(
|
|
doc_id=f"{drug_id}__chi_dinh__0", drug_id=drug_id, kind="prose",
|
|
section_key="chi_dinh", text="Điều trị sốt.", source_refs=(SOURCE,),
|
|
),
|
|
score=1.0,
|
|
)
|
|
|
|
|
|
def test_retrieve_by_indication_uses_keyword_hits_without_trying_dense():
|
|
retriever = _IndicationRetriever(keyword_hits=[_indication_hit("paracetamol_acetaminophen")])
|
|
service = RetrievalService(retriever, InMemoryParentStore([]))
|
|
|
|
result = service.retrieve_by_indication("sốt")
|
|
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
assert len(result.evidence) == 1
|
|
assert retriever.dense_called is False
|
|
|
|
|
|
def test_retrieve_by_indication_falls_back_to_dense_only_when_keyword_is_empty():
|
|
retriever = _IndicationRetriever(dense_hits=[_indication_hit("ibuprofen")])
|
|
service = RetrievalService(retriever, InMemoryParentStore([]))
|
|
|
|
result = service.retrieve_by_indication("thân nhiệt tăng")
|
|
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
assert retriever.dense_called is True
|
|
|
|
|
|
def test_retrieve_by_indication_with_no_match_anywhere_abstains():
|
|
retriever = _IndicationRetriever()
|
|
service = RetrievalService(retriever, InMemoryParentStore([]))
|
|
|
|
result = service.retrieve_by_indication("bệnh chưa từng ghi nhận")
|
|
|
|
assert result.decision == EvidenceDecision.ABSTAIN
|
|
assert result.reason == "no_indication_match"
|
|
|
|
|
|
def test_retrieve_by_indication_with_blank_text_abstains_without_calling_retrieval():
|
|
retriever = _IndicationRetriever(keyword_hits=[_indication_hit("paracetamol_acetaminophen")])
|
|
service = RetrievalService(retriever, InMemoryParentStore([]))
|
|
|
|
result = service.retrieve_by_indication(" ")
|
|
|
|
assert result.decision == EvidenceDecision.ABSTAIN
|
|
assert result.reason == "missing_indication"
|