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), ) 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"