Add production condition retrieval smoke test
This commit is contained in:
@@ -2,6 +2,11 @@ from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from .clinical import (
|
||||
CandidateStatus,
|
||||
MedicationCandidateAssessment,
|
||||
PatientContext,
|
||||
)
|
||||
from .context import pack_evidence
|
||||
from .models import Evidence, EvidenceDecision, RetrievalResult, SearchHit
|
||||
from .ports import (
|
||||
@@ -12,6 +17,7 @@ from .ports import (
|
||||
Retriever,
|
||||
)
|
||||
from .sections import SectionResolver
|
||||
from .text import normalize_name
|
||||
|
||||
|
||||
# The sections that introduce a drug: what it is, its class, its main use, its
|
||||
@@ -43,6 +49,14 @@ class EvidencePolicy:
|
||||
# symptom_to_drug: a common symptom can match far more drugs than is
|
||||
# useful to show in one answer.
|
||||
indication_candidate_limit: int = 8
|
||||
# Retrieve a wider chunk pool before grouping/ranking at drug level. The
|
||||
# final candidate cap above is applied only after entity aggregation.
|
||||
indication_retrieval_limit: int = 40
|
||||
indication_evidence_per_drug: int = 2
|
||||
# Patient-specific stage 2 is intentionally narrower than a general list.
|
||||
patient_candidate_limit: int = 2
|
||||
safety_hits_per_section: int = 1
|
||||
safety_sections_per_candidate: int = 4
|
||||
|
||||
|
||||
class RetrievalService:
|
||||
@@ -189,7 +203,7 @@ class RetrievalService:
|
||||
|
||||
find_by_indication = getattr(self._retriever, "find_by_indication", None)
|
||||
hits = (
|
||||
find_by_indication(indication_text, self._policy.indication_candidate_limit)
|
||||
find_by_indication(indication_text, self._policy.indication_retrieval_limit)
|
||||
if find_by_indication is not None
|
||||
else []
|
||||
)
|
||||
@@ -198,7 +212,7 @@ class RetrievalService:
|
||||
if search_indication is not None:
|
||||
try:
|
||||
hits = search_indication(
|
||||
indication_text, self._policy.indication_candidate_limit
|
||||
indication_text, self._policy.indication_retrieval_limit
|
||||
)
|
||||
except QueryEmbeddingUnavailable:
|
||||
hits = []
|
||||
@@ -213,7 +227,261 @@ class RetrievalService:
|
||||
hits = []
|
||||
if not hits:
|
||||
return RetrievalResult(EvidenceDecision.ABSTAIN, "no_indication_match")
|
||||
return self._decide(self._hydrate(hits, limit=None))
|
||||
groups = self._rank_indication_drugs(indication_text, hits)
|
||||
selected_hits = [
|
||||
hit
|
||||
for group in groups[: self._policy.indication_candidate_limit]
|
||||
for hit in group[: self._policy.indication_evidence_per_drug]
|
||||
]
|
||||
return self._decide(self._hydrate(selected_hits, limit=None))
|
||||
|
||||
def assess_patient_candidates(
|
||||
self,
|
||||
indication_result: RetrievalResult,
|
||||
patient: PatientContext,
|
||||
) -> tuple[RetrievalResult, tuple[MedicationCandidateAssessment, ...]]:
|
||||
"""Targeted stage-2 safety evidence for already-indicated candidates.
|
||||
|
||||
This never creates candidates. It groups the stage-1 indication
|
||||
evidence, keeps a bounded number of drugs, then searches only safety
|
||||
facets relevant to facts actually present in ``patient``. A lexical
|
||||
hit selects a chunk; pregnancy/breastfeeding sections are direct
|
||||
metadata routes because their relation is explicit in the section key.
|
||||
Absence of a hit is recorded as insufficient evidence, never "safe".
|
||||
"""
|
||||
if not patient.present or not indication_result.evidence:
|
||||
return indication_result, ()
|
||||
|
||||
indication_by_drug: dict[str, list[Evidence]] = {}
|
||||
for evidence in indication_result.evidence:
|
||||
drug_id = _evidence_drug_id(evidence)
|
||||
if drug_id:
|
||||
indication_by_drug.setdefault(drug_id, []).append(evidence)
|
||||
|
||||
assessments: list[MedicationCandidateAssessment] = []
|
||||
combined: list[Evidence] = []
|
||||
|
||||
for drug_id, indication_evidence in list(indication_by_drug.items())[
|
||||
: self._policy.patient_candidate_limit
|
||||
]:
|
||||
selected = self._patient_safety_evidence(drug_id, patient)
|
||||
|
||||
safety_evidence = tuple(
|
||||
evidence for values in selected.values() for evidence in values
|
||||
)
|
||||
status = (
|
||||
CandidateStatus.REQUIRES_ADDITIONAL_INFORMATION
|
||||
if any(item.requires_visual_check for item in safety_evidence)
|
||||
else CandidateStatus.SUPPORTED_WITH_CAUTION
|
||||
if safety_evidence
|
||||
else CandidateStatus.INSUFFICIENT_EVIDENCE
|
||||
)
|
||||
name = next(
|
||||
(item.drug_name for item in indication_evidence if item.drug_name),
|
||||
None,
|
||||
) or drug_id.replace("_", " ").title()
|
||||
assessment = MedicationCandidateAssessment(
|
||||
drug_id=drug_id,
|
||||
drug_name=name,
|
||||
indication_supported=True,
|
||||
indication_evidence=tuple(indication_evidence),
|
||||
contraindication_evidence=tuple(selected.get("chong_chi_dinh", ())),
|
||||
precaution_evidence=tuple(selected.get("than_trong", ())),
|
||||
interaction_evidence=tuple(selected.get("tuong_tac_thuoc", ())),
|
||||
renal_evidence=_facet_evidence(
|
||||
selected, patient.renal.present,
|
||||
("chong_chi_dinh", "than_trong", "lieu_luong_va_cach_dung"),
|
||||
),
|
||||
hepatic_evidence=_facet_evidence(
|
||||
selected, patient.hepatic.present,
|
||||
("chong_chi_dinh", "than_trong", "lieu_luong_va_cach_dung"),
|
||||
),
|
||||
pregnancy_evidence=tuple(selected.get("thoi_ky_mang_thai", ())),
|
||||
breastfeeding_evidence=tuple(selected.get("thoi_ky_cho_con_bu", ())),
|
||||
age_evidence=_facet_evidence(
|
||||
selected, bool(patient.age_text),
|
||||
("than_trong", "lieu_luong_va_cach_dung"),
|
||||
),
|
||||
dose_evidence=tuple(selected.get("lieu_luong_va_cach_dung", ())),
|
||||
status=status,
|
||||
)
|
||||
assessments.append(assessment)
|
||||
# Quarantined tables/formulas remain visible in the structured
|
||||
# assessment/status but never enter generation. Applying the
|
||||
# single-drug global VERIFY_PDF rule to a multi-candidate list
|
||||
# would suppress every otherwise verified prose candidate merely
|
||||
# because one candidate has one visual-only renal table.
|
||||
combined.extend(
|
||||
item for item in assessment.evidence
|
||||
if not item.requires_visual_check
|
||||
)
|
||||
|
||||
result = self._decide(tuple(combined))
|
||||
if result.decision == EvidenceDecision.ANSWERABLE:
|
||||
result = RetrievalResult(
|
||||
result.decision,
|
||||
"grounded_patient_evidence_available",
|
||||
result.evidence,
|
||||
)
|
||||
return result, tuple(assessments)
|
||||
|
||||
def retrieve_patient_drug_context(
|
||||
self,
|
||||
drug_id: str,
|
||||
base_result: RetrievalResult,
|
||||
patient: PatientContext,
|
||||
) -> RetrievalResult:
|
||||
"""Add bounded patient-relevant facets to a named-drug lookup."""
|
||||
if not patient.requires_safety_review:
|
||||
return base_result
|
||||
selected = self._patient_safety_evidence(drug_id, patient)
|
||||
evidence = list(base_result.evidence)
|
||||
seen = {item.evidence_id for item in evidence}
|
||||
for values in selected.values():
|
||||
for item in values:
|
||||
if item.requires_visual_check or item.evidence_id in seen:
|
||||
continue
|
||||
seen.add(item.evidence_id)
|
||||
evidence.append(item)
|
||||
result = self._decide(tuple(evidence))
|
||||
if result.decision == EvidenceDecision.ANSWERABLE:
|
||||
return RetrievalResult(
|
||||
result.decision,
|
||||
"grounded_patient_evidence_available",
|
||||
result.evidence,
|
||||
resolved_drug_id=base_result.resolved_drug_id,
|
||||
is_drug_overview=base_result.is_drug_overview,
|
||||
)
|
||||
return result
|
||||
|
||||
def _patient_safety_evidence(
|
||||
self, drug_id: str, patient: PatientContext
|
||||
) -> dict[str, list[Evidence]]:
|
||||
selected: dict[str, list[Evidence]] = {}
|
||||
search_lexical = getattr(self._retriever, "search_lexical", None)
|
||||
find_by_section = getattr(self._retriever, "find_by_section", None)
|
||||
|
||||
def lexical_facets(
|
||||
query: str,
|
||||
section_keys: tuple[str, ...],
|
||||
*,
|
||||
require_context_match: bool = False,
|
||||
) -> dict[str, list[SearchHit]]:
|
||||
"""Return at most the configured hits per requested relation.
|
||||
|
||||
Each clinical facet gets its own query. In particular, a current
|
||||
medicine may select an interaction chunk only when that medicine
|
||||
matches inside the interaction section; CKD/age terms from another
|
||||
facet cannot make an unrelated interaction look supported.
|
||||
"""
|
||||
if not query or search_lexical is None:
|
||||
return {}
|
||||
hits = search_lexical(
|
||||
query,
|
||||
drug_id,
|
||||
max(20, self._policy.safety_hits_per_section * len(section_keys)),
|
||||
section_keys=section_keys,
|
||||
)
|
||||
per_section: dict[str, list[SearchHit]] = {}
|
||||
for hit in hits:
|
||||
section = hit.document.section_key
|
||||
if section not in section_keys:
|
||||
continue
|
||||
if require_context_match and not _patient_context_matches(
|
||||
hit.document.text, patient
|
||||
):
|
||||
continue
|
||||
bucket = per_section.setdefault(section, [])
|
||||
if len(bucket) < self._policy.safety_hits_per_section:
|
||||
bucket.append(hit)
|
||||
return per_section
|
||||
|
||||
if search_lexical is not None:
|
||||
# These three searches deliberately keep their relations separate.
|
||||
# Absence of an exact lexical hit means "not evidenced in the
|
||||
# retrieved Dược thư text", never "no interaction/contraindication".
|
||||
interaction = lexical_facets(
|
||||
patient.interaction_query(), ("tuong_tac_thuoc",)
|
||||
)
|
||||
warnings = lexical_facets(
|
||||
patient.warning_query(),
|
||||
("chong_chi_dinh", "than_trong"),
|
||||
require_context_match=True,
|
||||
)
|
||||
dosage = lexical_facets(
|
||||
patient.dosage_context_query(),
|
||||
("lieu_luong_va_cach_dung",),
|
||||
require_context_match=True,
|
||||
)
|
||||
|
||||
candidates: list[tuple[str, list[SearchHit]]] = []
|
||||
if "tuong_tac_thuoc" in interaction:
|
||||
candidates.append(("tuong_tac_thuoc", interaction["tuong_tac_thuoc"]))
|
||||
warning_sections = sorted(
|
||||
warnings,
|
||||
key=lambda section: (
|
||||
section != "chong_chi_dinh",
|
||||
-warnings[section][0].score,
|
||||
section,
|
||||
),
|
||||
)
|
||||
for section in warning_sections:
|
||||
candidates.append((section, warnings[section]))
|
||||
if "lieu_luong_va_cach_dung" in dosage:
|
||||
candidates.append(
|
||||
("lieu_luong_va_cach_dung", dosage["lieu_luong_va_cach_dung"])
|
||||
)
|
||||
|
||||
for section, hits in candidates[
|
||||
: self._policy.safety_sections_per_candidate
|
||||
]:
|
||||
selected[section] = list(self._hydrate(hits, limit=None))
|
||||
|
||||
# These sections encode the patient relation themselves; no lexical
|
||||
# coincidence is needed to decide they are relevant.
|
||||
direct_sections = []
|
||||
if patient.pregnancy_status:
|
||||
direct_sections.append("thoi_ky_mang_thai")
|
||||
if patient.breastfeeding is True:
|
||||
direct_sections.append("thoi_ky_cho_con_bu")
|
||||
if find_by_section is not None:
|
||||
for section in direct_sections:
|
||||
if section in selected:
|
||||
continue
|
||||
hits = find_by_section(drug_id, section)
|
||||
if hits:
|
||||
selected[section] = list(self._hydrate(hits, limit=None))
|
||||
return selected
|
||||
|
||||
def _rank_indication_drugs(
|
||||
self, query: str, hits: list[SearchHit]
|
||||
) -> list[list[SearchHit]]:
|
||||
"""Group and rank entities without rewarding duplicate chunks."""
|
||||
by_drug: dict[str, list[SearchHit]] = {}
|
||||
for hit in hits:
|
||||
by_drug.setdefault(hit.document.drug_id, []).append(hit)
|
||||
groups = [
|
||||
sorted(group, key=lambda hit: (-hit.score, hit.document.doc_id))
|
||||
for group in by_drug.values()
|
||||
]
|
||||
groups.sort(key=lambda group: (-group[0].score, group[0].document.drug_id))
|
||||
if self._reranker is None or len(groups) <= 1:
|
||||
return groups
|
||||
documents = [
|
||||
f"{group[0].document.drug_name or group[0].document.drug_id}\n"
|
||||
+ "\n".join(hit.document.text for hit in group[:2])
|
||||
for group in groups
|
||||
]
|
||||
try:
|
||||
order = self._reranker.rerank(
|
||||
query,
|
||||
documents,
|
||||
top_n=self._policy.indication_candidate_limit,
|
||||
)
|
||||
except RerankUnavailable:
|
||||
return groups
|
||||
ranked = [groups[index] for index in order if 0 <= index < len(groups)]
|
||||
return ranked or groups
|
||||
|
||||
@staticmethod
|
||||
def _is_question(query: str) -> bool:
|
||||
@@ -384,6 +652,11 @@ class RetrievalService:
|
||||
requires_visual_check=(
|
||||
document.requires_visual_check or parent.requires_visual_check
|
||||
),
|
||||
drug_id=document.drug_id,
|
||||
drug_name=document.drug_name,
|
||||
section_key=document.section_key,
|
||||
section_title=document.section_title,
|
||||
source_document=document.source_document,
|
||||
))
|
||||
else:
|
||||
output.append(Evidence(
|
||||
@@ -395,8 +668,74 @@ class RetrievalService:
|
||||
source_refs=document.source_refs,
|
||||
hydrated_from_parent=False,
|
||||
requires_visual_check=document.requires_visual_check,
|
||||
drug_id=document.drug_id,
|
||||
drug_name=document.drug_name,
|
||||
section_key=document.section_key,
|
||||
section_title=document.section_title,
|
||||
source_document=document.source_document,
|
||||
))
|
||||
cap = self._policy.evidence_limit if limit == -1 else limit
|
||||
if cap is not None and len(output) >= cap:
|
||||
break
|
||||
return tuple(output)
|
||||
|
||||
|
||||
def _evidence_drug_id(evidence: Evidence) -> str | None:
|
||||
if evidence.drug_id:
|
||||
return evidence.drug_id
|
||||
if "__" in evidence.matched_doc_id:
|
||||
return evidence.matched_doc_id.split("__", 1)[0]
|
||||
return None
|
||||
|
||||
|
||||
def _facet_evidence(
|
||||
selected: dict[str, list[Evidence]],
|
||||
enabled: bool,
|
||||
section_keys: tuple[str, ...],
|
||||
) -> tuple[Evidence, ...]:
|
||||
if not enabled:
|
||||
return ()
|
||||
output: list[Evidence] = []
|
||||
seen: set[str] = set()
|
||||
for section in section_keys:
|
||||
for evidence in selected.get(section, ()):
|
||||
if evidence.evidence_id in seen:
|
||||
continue
|
||||
seen.add(evidence.evidence_id)
|
||||
output.append(evidence)
|
||||
return tuple(output)
|
||||
|
||||
|
||||
def _patient_context_matches(text: str, patient: PatientContext) -> bool:
|
||||
"""Require a clinical anchor, not overlap on generic words like 'chức năng'."""
|
||||
haystack = normalize_name(text)
|
||||
supplied = (
|
||||
*patient.comorbidities,
|
||||
*patient.allergies,
|
||||
*patient.previous_adverse_reactions,
|
||||
*patient.relevant_labs,
|
||||
)
|
||||
raw_terms = [normalize_name(term) for term in supplied if term.strip()]
|
||||
if any(term in haystack for term in raw_terms if len(term) >= 3):
|
||||
return True
|
||||
if patient.renal.present and any(
|
||||
term in haystack
|
||||
for term in (
|
||||
"suy than", "chuc nang than", "than nang", "creatinin", "crcl",
|
||||
"egfr", "loc cau than", "do thanh thai",
|
||||
)
|
||||
):
|
||||
return True
|
||||
if patient.hepatic.present and any(
|
||||
term in haystack
|
||||
for term in (
|
||||
"suy gan", "chuc nang gan", "benh gan", "xo gan", "child pugh",
|
||||
"ast", "alt", "bilirubin",
|
||||
)
|
||||
):
|
||||
return True
|
||||
if patient.age_text and any(
|
||||
term in haystack for term in ("nguoi cao tuoi", "cao tuoi", "tre em", "tre so sinh")
|
||||
):
|
||||
return True
|
||||
return False
|
||||
|
||||
Reference in New Issue
Block a user