Add production condition retrieval smoke test

2026-08-11 14:58:28 +07:00
parent 59e6ad2d0d
commit 7ebbe1f309
38 changed files with 3752 additions and 121 deletions
@@ -0,0 +1,354 @@
+from __future__ import annotations
+
+import json
+from dataclasses import replace
+
+from rag.agent import RagAgent, _patient_generation_query
+from rag.answer import GroundedAnswerService
+from rag.clinical import (
+    CaseContextAction,
+    CandidateStatus,
+    ConditionNormalizer,
+    ConditionQuery,
+    ConditionRelation,
+    PatientContext,
+    RenalContext,
+)
+from rag.models import Evidence, SourceRef
+from rag.understanding import (
+    LlmQueryUnderstander,
+    QueryFrame,
+    _apply_condition_candidate_cue,
+    _apply_named_drug_cues,
+    _apply_reverse_relation_cues,
+    _merge_with_prior_frame,
+)
+
+
+SOURCE = SourceRef(physical_page=10, precision="chunk_page_range", printed_page=11)
+
+
+def _evidence(drug_id: str, section: str = "chi_dinh") -> Evidence:
+    return Evidence(
+        evidence_id=f"{drug_id}__{section}__0",
+        matched_doc_id=f"{drug_id}__{section}__0",
+        kind="prose",
+        text=f"{drug_id} có nội dung {section}.",
+        score=1.0,
+        source_refs=(SOURCE,),
+        hydrated_from_parent=False,
+        requires_visual_check=False,
+        drug_id=drug_id,
+        drug_name=drug_id.upper(),
+        section_key=section,
+        section_title=section,
+    )
+
+
+class _Understander:
+    def __init__(self, frame: QueryFrame) -> None:
+        self.frame = frame
+
+    def understand(self, turn, history=(), budget=None, prior_frame=None):
+        return self.frame
+
+
+class _NoRetrieval:
+    def retrieve_by_indication(self, indication):
+        raise AssertionError(f"retrieval must not run for relation {indication}")
+
+
+def _agent(frame: QueryFrame) -> RagAgent:
+    return RagAgent(
+        _Understander(frame),
+        _NoRetrieval(),
+        GroundedAnswerService(routing=None),
+    )
+
+
+def test_condition_normalizer_handles_professional_aliases_without_drug_mapping():
+    normalizer = ConditionNormalizer()
+    assert normalizer.normalize("THA dùng gì", "THA").normalized_condition == "tăng huyết áp"
+    assert normalizer.normalize("cao huyết áp", "cao huyết áp").normalized_condition == "tăng huyết áp"
+    assert normalizer.normalize("Gout", "gout").normalized_condition == "gút"
+    assert normalizer.normalize("bệnh lạ", "bệnh lạ").normalized_condition == "bệnh lạ"
+
+
+def test_broad_condition_is_clarified_but_specific_subtype_is_not():
+    normalizer = ConditionNormalizer()
+    broad = normalizer.normalize("Viêm gan dùng thuốc gì?", "viêm gan")
+    specific = normalizer.normalize(
+        "Viêm gan B mạn dùng thuốc gì?", "viêm gan B mạn", subtype="B mạn"
+    )
+    assert broad.ambiguous is True
+    assert "A, B, C" in broad.clarify_question
+    assert specific.ambiguous is False
+
+
+def test_bare_broad_question_detector_does_not_overclarify_a_specific_site():
+    normalizer = ConditionNormalizer()
+
+    broad = normalizer.detect_broad_question("Nhiễm trùng dùng thuốc gì?")
+    specific = normalizer.detect_broad_question(
+        "Nhiễm trùng đường tiết niệu dùng thuốc gì?"
+    )
+
+    assert broad is not None
+    assert broad.ambiguous is True
+    assert broad.clarify_question
+    assert specific is None
+
+
+def test_relation_confusion_never_reaches_indication_retrieval():
+    for relation, question in (
+        (ConditionRelation.ADVERSE_EFFECT, "Thuốc nào gây tăng huyết áp?"),
+        (ConditionRelation.CONTRAINDICATION, "Thuốc nào chống chỉ định ở bệnh nhân gout?"),
+    ):
+        reply = _agent(QueryFrame(
+            turn_type="condition_relation",
+            indication="tăng huyết áp",
+            condition_relation=relation,
+        )).handle(question)
+        assert reply.decision == "abstain"
+        assert reply.reason == "unsupported_reverse_relation"
+        assert "không" in reply.answer.lower()
+
+
+def test_explicit_reverse_relation_cue_overrides_a_noisy_llm_clarification():
+    noisy = QueryFrame(
+        turn_type="out_of_scope",
+        needs_clarify=True,
+        clarify_reason="Bạn muốn hỏi thuốc nào?",
+    )
+
+    adverse = _apply_reverse_relation_cues(noisy, "Thuốc nào gây tăng huyết áp?")
+    contraindicated = _apply_reverse_relation_cues(
+        noisy, "Thuốc nào chống chỉ định ở bệnh nhân gout?"
+    )
+
+    assert adverse.turn_type == "condition_relation"
+    assert adverse.condition_relation == ConditionRelation.ADVERSE_EFFECT
+    assert adverse.needs_clarify is False
+    assert contraindicated.condition_relation == ConditionRelation.CONTRAINDICATION
+
+
+def test_patient_candidate_wording_is_not_mistaken_for_reverse_contraindication():
+    normalizer = ConditionNormalizer()
+    noisy = QueryFrame(
+        turn_type="drug_attribute",
+        drugs=("digoxin",),
+        needs_clarify=True,
+        clarify_reason="Bạn muốn hỏi digoxin?",
+    )
+
+    candidate = _apply_condition_candidate_cue(
+        noisy,
+        "BN gout kèm suy thận nặng dùng thuốc nào cần thận trọng hoặc chống chỉ định?",
+        normalizer,
+    )
+    after_relation_guard = _apply_reverse_relation_cues(
+        candidate,
+        "BN gout kèm suy thận nặng dùng thuốc nào cần thận trọng hoặc chống chỉ định?",
+    )
+
+    assert candidate.turn_type == "condition_to_drug"
+    assert candidate.condition is not None
+    assert candidate.condition.normalized_condition == "gút"
+    assert after_relation_guard.turn_type == "condition_to_drug"
+
+
+def test_named_drug_safety_and_purpose_cues_override_noisy_relation_frames():
+    noisy = QueryFrame(
+        turn_type="condition_relation",
+        drugs=("probenecid",),
+        condition_relation=ConditionRelation.CONTRAINDICATION,
+        needs_clarify=True,
+        clarify_reason="Cần làm rõ",
+    )
+
+    safety = _apply_named_drug_cues(
+        noisy, "BN eGFR 25, probenecid có dùng được không?"
+    )
+    purpose = _apply_named_drug_cues(
+        replace(noisy, drugs=("paracetamol_acetaminophen",)),
+        "Paracetamol có tác dụng gì?",
+    )
+
+    assert safety.turn_type == "drug_attribute"
+    assert safety.attribute == "chong_chi_dinh"
+    assert safety.needs_clarify is False
+    assert purpose.turn_type == "drug_to_condition"
+    assert purpose.attribute == "chi_dinh"
+
+
+def test_ambiguous_condition_clarifies_before_retrieval():
+    reply = _agent(QueryFrame(
+        turn_type="condition_to_drug",
+        indication="viêm gan",
+        condition=ConditionQuery(
+            original_query="Viêm gan dùng thuốc gì?",
+            normalized_condition="viêm gan",
+            ambiguous=True,
+            clarify_question="Bạn đang hỏi viêm gan A, B, C hay loại nào?",
+        ),
+    )).handle("Viêm gan dùng thuốc gì?")
+    assert reply.decision == "clarify"
+    assert reply.reason == "ambiguous_condition"
+
+
+def test_patient_context_merges_only_for_explicit_same_case_continuation():
+    prior = QueryFrame(
+        turn_type="condition_to_drug",
+        patient_context=PatientContext(
+            age_text="68 tuổi",
+            comorbidities=("CKD G4",),
+            renal=RenalContext(description="CKD", ckd_stage="G4"),
+        ),
+    )
+    current = QueryFrame(
+        turn_type="condition_to_drug",
+        indication="tăng huyết áp",
+        condition=ConditionQuery("BN bị tăng huyết áp", "tăng huyết áp"),
+        patient_context=PatientContext(primary_condition="tăng huyết áp"),
+        context_action=CaseContextAction.CONTINUE,
+    )
+    merged = _merge_with_prior_frame(current, prior)
+    assert merged.patient_context.age_text == "68 tuổi"
+    assert merged.patient_context.renal.ckd_stage == "G4"
+
+    new_case = _merge_with_prior_frame(
+        QueryFrame(
+            turn_type="condition_to_drug",
+            indication="gút",
+            context_action=CaseContextAction.NEW,
+            patient_context=PatientContext(primary_condition="gút"),
+        ),
+        prior,
+    )
+    assert new_case.patient_context.age_text is None
+    assert new_case.patient_context.renal.present is False
+
+
+def test_patient_generation_query_keeps_task_but_not_user_only_numbers():
+    frame = QueryFrame(
+        turn_type="condition_to_drug",
+        condition=ConditionQuery(
+            "THA", "tăng huyết áp"
+        ),
+        patient_context=PatientContext(
+            age_text="68 tuổi",
+            primary_condition="tăng huyết áp",
+            comorbidities=("CKD G4",),
+            current_medications=("digoxin",),
+            renal=RenalContext(description="eGFR 25", ckd_stage="G4", egfr="25"),
+        ),
+    )
+
+    query = _patient_generation_query(frame)
+
+    assert "tăng huyết áp" in query
+    assert "68" not in query
+    assert "G4" not in query
+    assert "25" not in query
+    assert "digoxin" not in query
+
+
+def test_candidate_status_keeps_indication_separate_from_patient_safety():
+    from rag.clinical import MedicationCandidateAssessment
+
+    assessment = MedicationCandidateAssessment(
+        drug_id="a",
+        drug_name="A",
+        indication_supported=True,
+        indication_evidence=(_evidence("a"),),
+        status=CandidateStatus.INSUFFICIENT_EVIDENCE,
+    )
+    assert assessment.indication_supported is True
+    assert assessment.status == CandidateStatus.INSUFFICIENT_EVIDENCE
+    assert len(assessment.evidence) == 1
+
+
+class _NoCandidates:
+    def resolve(self, query):
+        class Result:
+            status = "not_found"
+            drug_id = None
+            candidate_drug_ids = ()
+        return Result()
+
+    def suggest(self, query, k=3, min_score=0.5):
+        return []
+
+
+class _JsonLlm:
+    def __init__(self, payload: dict) -> None:
+        self.payload = payload
+
+    def generate(self, system, user, schema):
+        return json.dumps(self.payload, ensure_ascii=False)
+
+
+def test_understanding_parses_condition_and_patient_context_without_inventing_fields():
+    payload = {
+        "turn_type": "condition_to_drug",
+        "drugs": [],
+        "unknown_drugs": [],
+        "attribute": None,
+        "population": "suy_than",
+        "weight_kg": None,
+        "age_text": "68 tuổi",
+        "indication": "THA",
+        "condition": {
+            "original_text": "THA",
+            "normalized_condition": "tăng huyết áp",
+            "subtype": None,
+            "qualifiers": [],
+            "ambiguous": False,
+            "clarify_question": None,
+        },
+        "condition_relation": "indication",
+        "patient_context": {
+            "age_text": "68 tuổi",
+            "sex": None,
+            "weight_kg": None,
+            "primary_condition": "tăng huyết áp",
+            "comorbidities": ["CKD G4", "gout"],
+            "allergies": [],
+            "previous_adverse_reactions": [],
+            "current_medications": ["digoxin"],
+            "pregnancy_status": None,
+            "breastfeeding": None,
+            "renal": {
+                "description": "CKD",
+                "ckd_stage": "G4",
+                "egfr": None,
+                "crcl": None,
+                "creatinine": None,
+            },
+            "hepatic": {},
+            "relevant_labs": ["K 5.7"],
+            "treatment_history": [],
+        },
+        "context_action": "none",
+        "route": None,
+        "section_overview": False,
+        "standalone_query": "BN 68 tuổi, THA + CKD G4 + gout, đang dùng digoxin",
+        "depends_on_previous_turn": False,
+        "needs_clarify": False,
+        "clarify_reason": None,
+        "quick_replies": [],
+    }
+    understander = LlmQueryUnderstander(
+        _JsonLlm(payload), {}, _NoCandidates()
+    )
+
+    frame = understander.understand(
+        "BN 68 tuổi, THA + CKD G4 + gout, K 5.7, đang dùng digoxin. Option hạ áp?"
+    )
+
+    assert frame.turn_type == "condition_to_drug"
+    assert frame.condition.normalized_condition == "tăng huyết áp"
+    assert frame.patient_context.comorbidities == ("CKD G4", "gout")
+    assert frame.patient_context.current_medications == ("digoxin",)
+    assert frame.patient_context.renal.ckd_stage == "G4"
+    assert frame.patient_context.hepatic.present is False