499 lines
17 KiB
Python
499 lines
17 KiB
Python
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,
|
|
HepaticContext,
|
|
PatientContext,
|
|
RenalContext,
|
|
)
|
|
from rag.models import Evidence, SourceRef
|
|
from rag.sections import SECTION_PHRASES
|
|
from rag.understanding import (
|
|
LlmQueryUnderstander,
|
|
QueryFrame,
|
|
_apply_condition_candidate_cue,
|
|
_apply_contextual_candidate_safety,
|
|
_apply_general_condition_scope,
|
|
_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.detect_known_alias("BN viêm phổi dùng thuốc gì?").normalized_condition == "viêm phổi"
|
|
assert normalizer.normalize("bệnh lạ", "bệnh lạ").normalized_condition == "bệnh lạ"
|
|
|
|
|
|
def test_explicit_indication_relation_is_a_condition_candidate_lookup():
|
|
noisy = QueryFrame(
|
|
turn_type="condition_relation",
|
|
condition=ConditionNormalizer().normalize("bệnh gút", "gút"),
|
|
needs_clarify=True,
|
|
clarify_reason="Hỏi lại sai hướng",
|
|
)
|
|
|
|
frame = _apply_condition_candidate_cue(
|
|
noisy,
|
|
"Thuốc nào có chỉ định liên quan bệnh gút?",
|
|
ConditionNormalizer(),
|
|
)
|
|
|
|
assert frame.turn_type == "condition_to_drug"
|
|
assert frame.condition_relation == ConditionRelation.INDICATION
|
|
assert frame.needs_clarify is False
|
|
|
|
|
|
def test_general_condition_does_not_invent_patient_hepatic_context():
|
|
frame = QueryFrame(
|
|
turn_type="condition_to_drug",
|
|
indication="viêm gan B mạn",
|
|
condition=ConditionQuery(
|
|
original_query="Viêm gan B mạn dùng thuốc gì?",
|
|
normalized_condition="viêm gan B mạn",
|
|
),
|
|
patient_context=PatientContext(
|
|
primary_condition="viêm gan B mạn",
|
|
hepatic=HepaticContext(description="viêm gan B mạn"),
|
|
),
|
|
)
|
|
|
|
cleaned = _apply_general_condition_scope(
|
|
frame, "Viêm gan B mạn dùng thuốc gì?"
|
|
)
|
|
|
|
assert cleaned.patient_context.primary_condition == "viêm gan B mạn"
|
|
assert cleaned.patient_context.requires_safety_review is False
|
|
|
|
|
|
def test_patient_allergy_condition_lookup_keeps_safety_context():
|
|
patient = PatientContext(
|
|
primary_condition="viêm phổi", allergies=("penicillin",)
|
|
)
|
|
frame = QueryFrame(
|
|
turn_type="condition_to_drug",
|
|
condition=ConditionQuery(
|
|
original_query="BN dị ứng penicillin, viêm phổi dùng thuốc gì?",
|
|
normalized_condition="viêm phổi",
|
|
),
|
|
patient_context=patient,
|
|
)
|
|
|
|
kept = _apply_general_condition_scope(
|
|
frame, "BN dị ứng penicillin, viêm phổi dùng thuốc gì?"
|
|
)
|
|
|
|
assert kept.patient_context == patient
|
|
assert kept.patient_context.requires_safety_review is True
|
|
|
|
|
|
def test_candidate_safety_followup_stays_on_prior_condition_lookup():
|
|
prior = QueryFrame(
|
|
turn_type="condition_to_drug",
|
|
indication="tăng huyết áp",
|
|
condition=ConditionQuery(
|
|
original_query="BN bị tăng huyết áp",
|
|
normalized_condition="tăng huyết áp",
|
|
),
|
|
patient_context=PatientContext(
|
|
age_text="68 tuổi",
|
|
renal=RenalContext(description="CKD", ckd_stage="G4"),
|
|
),
|
|
)
|
|
noisy = QueryFrame(
|
|
turn_type="condition_relation",
|
|
condition_relation=ConditionRelation.CONTRAINDICATION,
|
|
depends_on_previous_turn=True,
|
|
patient_context=prior.patient_context,
|
|
needs_clarify=False,
|
|
)
|
|
|
|
corrected = _apply_contextual_candidate_safety(
|
|
noisy,
|
|
"Trong các thuốc trên cái nào cần lưu ý hơn với bệnh thận?",
|
|
prior,
|
|
)
|
|
|
|
assert corrected.turn_type == "condition_to_drug"
|
|
assert corrected.condition == prior.condition
|
|
assert corrected.condition_relation == ConditionRelation.INDICATION
|
|
|
|
|
|
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ì?",
|
|
)
|
|
golden_contraindication = _apply_named_drug_cues(
|
|
replace(noisy, drugs=("paracetamol_acetaminophen",)),
|
|
"Chống chỉ định của Paracetamol là gì?",
|
|
)
|
|
golden_contraindication_with_omitted_llm_drug = _apply_named_drug_cues(
|
|
replace(noisy, drugs=()),
|
|
"Chống chỉ định của Paracetamol là gì?",
|
|
resolved_turn_drug="paracetamol_acetaminophen",
|
|
)
|
|
reverse_contraindication = _apply_named_drug_cues(
|
|
replace(noisy, drugs=("warfarin",)),
|
|
"Thuốc nào chống chỉ định ở bệnh nhân đang dùng warfarin?",
|
|
)
|
|
|
|
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"
|
|
assert golden_contraindication.turn_type == "drug_attribute"
|
|
assert golden_contraindication.attribute == "chong_chi_dinh"
|
|
assert golden_contraindication.condition_relation == ConditionRelation.CONTRAINDICATION
|
|
assert golden_contraindication_with_omitted_llm_drug.turn_type == "drug_attribute"
|
|
assert golden_contraindication_with_omitted_llm_drug.drugs == (
|
|
"paracetamol_acetaminophen",
|
|
)
|
|
assert reverse_contraindication.turn_type == "condition_relation"
|
|
|
|
|
|
def test_owned_section_phrase_routes_every_monograph_section_without_llm_help():
|
|
noisy = QueryFrame(
|
|
turn_type="condition_relation",
|
|
condition_relation=ConditionRelation.CONTRAINDICATION,
|
|
needs_clarify=True,
|
|
clarify_reason="Cần làm rõ",
|
|
)
|
|
|
|
for section_key, phrases in SECTION_PHRASES.items():
|
|
phrase = phrases[0]
|
|
corrected = _apply_named_drug_cues(
|
|
noisy,
|
|
f"{phrase} của Paracetamol là gì?",
|
|
resolved_turn_drug="paracetamol_acetaminophen",
|
|
resolved_section_key=section_key,
|
|
resolved_section_phrase=phrase,
|
|
)
|
|
|
|
assert corrected.turn_type == "drug_attribute", section_key
|
|
assert corrected.drugs == ("paracetamol_acetaminophen",), section_key
|
|
assert corrected.attribute == section_key
|
|
assert corrected.needs_clarify is False
|
|
|
|
|
|
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
|