Add production condition retrieval smoke test

apps/ai-service/tests/test_retrieval_service.py

@@ -1,6 +1,7 @@
 from pathlib import Path
 
 from rag.artifacts import load_aliases
+from rag.clinical import CandidateStatus, PatientContext, RenalContext
 from rag.evaluation import CaseOrigin, EvaluationCase, EvaluationOutcome, summarize
 from rag.in_memory import InMemoryLexicalRetriever, InMemoryParentStore, _char_ngrams
 from rag.models import (
@@ -519,3 +520,191 @@ def test_retrieve_by_indication_with_blank_text_abstains_without_calling_retriev
 
     assert result.decision == EvidenceDecision.ABSTAIN
     assert result.reason == "missing_indication"
+
+
+def test_indication_candidates_are_ranked_per_drug_not_by_chunk_count():
+    many_weak = [
+        SearchHit(
+            RetrievalDocument(
+                doc_id=f"drug_many__chi_dinh__{index}",
+                drug_id="drug_many",
+                kind="prose",
+                section_key="chi_dinh",
+                text="Điều trị tăng huyết áp.",
+                source_refs=(SOURCE,),
+                part_index=index,
+            ),
+            score=1.0,
+        )
+        for index in range(8)
+    ]
+    one_strong = SearchHit(
+        RetrievalDocument(
+            doc_id="drug_strong__chi_dinh__0",
+            drug_id="drug_strong",
+            kind="prose",
+            section_key="chi_dinh",
+            text="Điều trị tăng huyết áp.",
+            source_refs=(SOURCE,),
+        ),
+        score=2.0,
+    )
+    retriever = _IndicationRetriever(keyword_hits=[*many_weak, one_strong])
+    service = RetrievalService(
+        retriever,
+        InMemoryParentStore([]),
+        EvidencePolicy(indication_candidate_limit=2, indication_evidence_per_drug=2),
+    )
+
+    result = service.retrieve_by_indication("tăng huyết áp")
+
+    assert result.evidence[0].drug_id == "drug_strong"
+    assert [item.drug_id for item in result.evidence].count("drug_many") == 2
+    assert len(result.evidence) == 3
+
+
+class _PatientSafetyRetriever(_IndicationRetriever):
+    def __init__(self) -> None:
+        super().__init__(keyword_hits=[_indication_hit("amlodipin")])
+        self.safety_calls: list[tuple[str, tuple[str, ...]]] = []
+
+    def search_lexical(self, query, drug_id, limit, section_keys=None):
+        self.safety_calls.append((query, section_keys or ()))
+        hits = [
+            SearchHit(
+                RetrievalDocument(
+                    doc_id=f"{drug_id}__than_trong__0",
+                    drug_id=drug_id,
+                    drug_name="AMLODIPIN",
+                    kind="prose",
+                    section_key="than_trong",
+                    section_title="Thận trọng",
+                    text="Thận trọng ở người bệnh suy thận.",
+                    source_refs=(SOURCE,),
+                ),
+                score=3.0,
+            ),
+            SearchHit(
+                RetrievalDocument(
+                    doc_id=f"{drug_id}__tuong_tac_thuoc__0",
+                    drug_id=drug_id,
+                    drug_name="AMLODIPIN",
+                    kind="prose",
+                    section_key="tuong_tac_thuoc",
+                    section_title="Tương tác thuốc",
+                    text="Tương tác được ghi nhận với digoxin.",
+                    source_refs=(SOURCE,),
+                ),
+                score=2.0,
+            ),
+        ]
+        return [
+            hit
+            for hit in hits
+            if (not section_keys or hit.document.section_key in section_keys)
+            and (
+                hit.document.section_key != "tuong_tac_thuoc"
+                or "digoxin" in query.casefold()
+            )
+        ][:limit]
+
+    def find_by_section(self, drug_id, section_key):
+        return []
+
+
+def test_patient_stage_two_targets_renal_and_current_medication_evidence():
+    retriever = _PatientSafetyRetriever()
+    service = RetrievalService(retriever, InMemoryParentStore([]))
+    indication = service.retrieve_by_indication("tăng huyết áp")
+    patient = PatientContext(
+        age_text="68 tuổi",
+        comorbidities=("CKD G4", "gút"),
+        current_medications=("digoxin",),
+        renal=RenalContext(description="CKD", ckd_stage="G4"),
+    )
+
+    result, assessments = service.assess_patient_candidates(indication, patient)
+
+    assert result.decision == EvidenceDecision.ANSWERABLE
+    assert result.reason == "grounded_patient_evidence_available"
+    assert len(assessments) == 1
+    assessment = assessments[0]
+    assert assessment.status == CandidateStatus.SUPPORTED_WITH_CAUTION
+    assert assessment.renal_evidence
+    assert assessment.interaction_evidence
+    searched_sections = {
+        section
+        for _, sections in retriever.safety_calls
+        for section in sections
+    }
+    assert "tuong_tac_thuoc" in searched_sections
+    assert "lieu_luong_va_cach_dung" in searched_sections
+
+
+def test_patient_interaction_requires_current_drug_match_in_interaction_section():
+    retriever = _PatientSafetyRetriever()
+    service = RetrievalService(retriever, InMemoryParentStore([]))
+    indication = service.retrieve_by_indication("tăng huyết áp")
+    patient = PatientContext(
+        current_medications=("warfarin",),
+        renal=RenalContext(description="CKD", ckd_stage="G4"),
+    )
+
+    _, assessments = service.assess_patient_candidates(indication, patient)
+
+    assert len(assessments) == 1
+    assert assessments[0].interaction_evidence == ()
+    interaction_calls = [
+        query
+        for query, sections in retriever.safety_calls
+        if sections == ("tuong_tac_thuoc",)
+    ]
+    assert interaction_calls == ["warfarin"]
+
+
+def test_matching_contraindication_is_retained_without_declaring_patient_status():
+    class ContraindicationRetriever(_PatientSafetyRetriever):
+        def search_lexical(self, query, drug_id, limit, section_keys=None):
+            hits = super().search_lexical(query, drug_id, limit, section_keys)
+            if section_keys and "chong_chi_dinh" in section_keys:
+                hits.append(SearchHit(
+                    RetrievalDocument(
+                        doc_id=f"{drug_id}__chong_chi_dinh__0",
+                        drug_id=drug_id,
+                        drug_name="AMLODIPIN",
+                        kind="prose",
+                        section_key="chong_chi_dinh",
+                        section_title="Chống chỉ định",
+                        text="Chống chỉ định ở người bệnh suy thận nặng.",
+                        source_refs=(SOURCE,),
+                    ),
+                    score=1.0,
+                ))
+            return hits
+
+    retriever = ContraindicationRetriever()
+    service = RetrievalService(retriever, InMemoryParentStore([]))
+    indication = service.retrieve_by_indication("tăng huyết áp")
+
+    _, assessments = service.assess_patient_candidates(
+        indication,
+        PatientContext(renal=RenalContext(description="suy thận nặng")),
+    )
+
+    assert assessments[0].status == CandidateStatus.SUPPORTED_WITH_CAUTION
+    assert assessments[0].contraindication_evidence
+
+
+def test_named_drug_renal_query_adds_targeted_patient_safety_evidence():
+    retriever = _PatientSafetyRetriever()
+    service = RetrievalService(retriever, InMemoryParentStore([]))
+    base = service.retrieve_by_indication("tăng huyết áp")
+    patient = PatientContext(
+        renal=RenalContext(description="suy thận", egfr="25 ml/phút/1,73 m2")
+    )
+
+    result = service.retrieve_patient_drug_context("amlodipin", base, patient)
+
+    assert result.decision == EvidenceDecision.ANSWERABLE
+    assert result.reason == "grounded_patient_evidence_available"
+    assert {item.section_key for item in result.evidence} >= {"chi_dinh", "than_trong"}