219 lines
8.8 KiB
Python
219 lines
8.8 KiB
Python
"""Section routing: the fix for hit@1 0.05 on `chong_chi_dinh`.
|
|
|
|
Measured 2026-08-04 on the real Cohere collection, letting vector similarity
|
|
choose the section answered contraindication questions correctly 1 time in 20.
|
|
These tests pin the two properties that make filtering safe: the longer phrase
|
|
always wins, and an unrecognised question routes nowhere rather than guessing.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from rag.models import (
|
|
EvidenceDecision,
|
|
RetrievalDocument,
|
|
SearchHit,
|
|
SourceRef,
|
|
)
|
|
from rag.in_memory import InMemoryParentStore
|
|
from rag.sections import SectionResolver
|
|
from rag.service import EvidencePolicy, RetrievalService
|
|
|
|
SOURCE = SourceRef(physical_page=200, precision="page", printed_page=142)
|
|
|
|
|
|
def _doc(doc_id: str, section_key: str, text: str) -> RetrievalDocument:
|
|
return RetrievalDocument(
|
|
doc_id=doc_id,
|
|
parent_id=None,
|
|
drug_id="aspirin",
|
|
kind="prose",
|
|
section_key=section_key,
|
|
text=text,
|
|
source_refs=(SOURCE,),
|
|
requires_visual_check=False,
|
|
)
|
|
|
|
|
|
class SectionAwareRetriever:
|
|
"""Fake that records which route the service actually took."""
|
|
|
|
def __init__(self, docs: list[RetrievalDocument]) -> None:
|
|
self._docs = docs
|
|
self.search_calls: list[str] = []
|
|
self.section_calls: list[tuple[str, str]] = []
|
|
|
|
def search(
|
|
self, query: str, drug_id: str, limit: int # noqa: ARG002 — Retriever protocol
|
|
) -> list[SearchHit]:
|
|
self.search_calls.append(query)
|
|
# Deliberately wrong on purpose: the whole point is that the section
|
|
# route must not consult similarity at all.
|
|
return [SearchHit(self._docs[-1], 0.99)]
|
|
|
|
def find_by_section(self, drug_id: str, section_key: str) -> list[SearchHit]:
|
|
self.section_calls.append((drug_id, section_key))
|
|
return [
|
|
SearchHit(doc, 1.0) for doc in self._docs if doc.section_key == section_key
|
|
]
|
|
|
|
|
|
class SimilarityOnlyRetriever:
|
|
def __init__(self, docs: list[RetrievalDocument]) -> None:
|
|
self._docs = docs
|
|
self.search_calls: list[str] = []
|
|
|
|
def search(
|
|
self, query: str, drug_id: str, limit: int # noqa: ARG002 — Retriever protocol
|
|
) -> list[SearchHit]:
|
|
self.search_calls.append(query)
|
|
return [SearchHit(self._docs[0], 0.99)]
|
|
|
|
|
|
CONTRA = [
|
|
_doc("c1", "chong_chi_dinh", "Mẫn cảm với aspirin."),
|
|
_doc("c2", "chong_chi_dinh", "Loét dạ dày tá tràng đang tiến triển."),
|
|
_doc("c3", "chong_chi_dinh", "Hen do aspirin."),
|
|
_doc("c4", "chong_chi_dinh", "Suy gan nặng."),
|
|
_doc("c5", "chong_chi_dinh", "Trẻ em dưới 16 tuổi có sốt virus."),
|
|
]
|
|
INDICATION = [_doc("i1", "chi_dinh", "Giảm đau, hạ sốt, chống viêm.")]
|
|
PHARMACOLOGY = [_doc("p1", "duoc_ly_va_co_che_tac_dung", "Ức chế cyclooxygenase.")]
|
|
ALL_DOCS = CONTRA + INDICATION + PHARMACOLOGY
|
|
|
|
|
|
def _service(retriever, resolver: SectionResolver | None) -> RetrievalService:
|
|
return RetrievalService(
|
|
retriever,
|
|
InMemoryParentStore([]),
|
|
EvidencePolicy(minimum_score=0.01),
|
|
section_resolver=resolver,
|
|
)
|
|
|
|
|
|
class TestSectionResolver:
|
|
def test_contraindication_is_never_read_as_indication(self) -> None:
|
|
"""The one that measured 0.05. "chống chỉ định" contains "chỉ định"."""
|
|
resolver = SectionResolver()
|
|
assert resolver.resolve("Chống chỉ định của aspirin là gì?").section_key == (
|
|
"chong_chi_dinh"
|
|
)
|
|
assert resolver.resolve("Chỉ định của aspirin?").section_key == "chi_dinh"
|
|
|
|
def test_works_without_diacritics(self) -> None:
|
|
assert SectionResolver().resolve("aspirin chong chi dinh").section_key == (
|
|
"chong_chi_dinh"
|
|
)
|
|
|
|
def test_overdose_is_not_read_as_dose(self) -> None:
|
|
resolver = SectionResolver()
|
|
assert resolver.resolve("xử trí quá liều metformin").section_key == (
|
|
"qua_lieu_va_xu_tri"
|
|
)
|
|
assert resolver.resolve("liều dùng metformin").section_key == (
|
|
"lieu_luong_va_cach_dung"
|
|
)
|
|
|
|
def test_adr_management_is_not_read_as_adr_itself(self) -> None:
|
|
resolver = SectionResolver()
|
|
assert resolver.resolve("xử trí tác dụng phụ của prednisolon").section_key == (
|
|
"huong_dan_xu_tri_adr"
|
|
)
|
|
assert resolver.resolve("tác dụng phụ của prednisolon").section_key == (
|
|
"tac_dung_khong_mong_muon"
|
|
)
|
|
|
|
def test_incompatibility_is_not_read_as_interaction(self) -> None:
|
|
resolver = SectionResolver()
|
|
assert resolver.resolve("tương kỵ của ceftriaxon").section_key == "tuong_ky"
|
|
assert resolver.resolve("tương tác của ceftriaxon").section_key == (
|
|
"tuong_tac_thuoc"
|
|
)
|
|
|
|
def test_bare_lieu_resolves_without_capturing_overdose(self) -> None:
|
|
"""Found by testing on human-written golden questions, not templates.
|
|
|
|
4 of 16 said just "Liều Metformin cho người lớn?". Adding bare "liều"
|
|
is only safe because "quá liều" is longer and is tested first.
|
|
"""
|
|
resolver = SectionResolver()
|
|
assert resolver.resolve("Liều Metformin cho người lớn?").section_key == (
|
|
"lieu_luong_va_cach_dung"
|
|
)
|
|
assert resolver.resolve("quá liều paracetamol").section_key == (
|
|
"qua_lieu_va_xu_tri"
|
|
)
|
|
|
|
def test_colloquial_pregnancy_phrasing(self) -> None:
|
|
assert SectionResolver().resolve(
|
|
"Bà bầu dùng Ibuprofen được không?"
|
|
).section_key == "thoi_ky_mang_thai"
|
|
|
|
def test_unrecognised_question_routes_nowhere(self) -> None:
|
|
"""No match must not become a guess — the caller falls back."""
|
|
assert SectionResolver().resolve("thuốc này giá bao nhiêu") is None
|
|
assert SectionResolver().resolve("") is None
|
|
|
|
def test_new_section_needs_no_code_change(self) -> None:
|
|
resolver = SectionResolver({"invented_section": ("một mục hoàn toàn mới",)})
|
|
assert resolver.resolve("hỏi về một mục hoàn toàn mới").section_key == (
|
|
"invented_section"
|
|
)
|
|
|
|
|
|
class TestSectionRouting:
|
|
def test_named_section_bypasses_similarity_entirely(self) -> None:
|
|
retriever = SectionAwareRetriever(ALL_DOCS)
|
|
result = _service(retriever, SectionResolver()).retrieve(
|
|
"Chống chỉ định của aspirin là gì?", "aspirin"
|
|
)
|
|
assert retriever.section_calls == [("aspirin", "chong_chi_dinh")]
|
|
assert retriever.search_calls == []
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
assert {item.evidence_id for item in result.evidence} == {
|
|
"c1", "c2", "c3", "c4", "c5",
|
|
}
|
|
|
|
def test_whole_section_is_returned_past_the_evidence_limit(self) -> None:
|
|
"""Five contraindications must not arrive as three."""
|
|
retriever = SectionAwareRetriever(ALL_DOCS)
|
|
service = RetrievalService(
|
|
retriever,
|
|
InMemoryParentStore([]),
|
|
EvidencePolicy(minimum_score=0.01, evidence_limit=3),
|
|
section_resolver=SectionResolver(),
|
|
)
|
|
result = service.retrieve("chống chỉ định aspirin", "aspirin")
|
|
assert len(result.evidence) == 5
|
|
|
|
def test_unnamed_section_falls_back_to_similarity(self) -> None:
|
|
retriever = SectionAwareRetriever(ALL_DOCS)
|
|
result = _service(retriever, SectionResolver()).retrieve(
|
|
"aspirin dùng cho bệnh nhân này thế nào", "aspirin"
|
|
)
|
|
assert retriever.section_calls == []
|
|
assert retriever.search_calls
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
|
|
def test_retriever_without_the_capability_still_works(self) -> None:
|
|
retriever = SimilarityOnlyRetriever(ALL_DOCS)
|
|
result = _service(retriever, SectionResolver()).retrieve(
|
|
"Chống chỉ định của aspirin là gì?", "aspirin"
|
|
)
|
|
assert retriever.search_calls
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
|
|
def test_no_resolver_keeps_the_old_behaviour(self) -> None:
|
|
retriever = SectionAwareRetriever(ALL_DOCS)
|
|
_service(retriever, None).retrieve("chống chỉ định aspirin", "aspirin")
|
|
assert retriever.section_calls == []
|
|
assert retriever.search_calls
|
|
|
|
def test_named_but_empty_section_falls_back(self) -> None:
|
|
"""A drug with no such section must not abstain — similarity still tries."""
|
|
retriever = SectionAwareRetriever(INDICATION + PHARMACOLOGY)
|
|
result = _service(retriever, SectionResolver()).retrieve(
|
|
"chống chỉ định aspirin", "aspirin"
|
|
)
|
|
assert retriever.section_calls == [("aspirin", "chong_chi_dinh")]
|
|
assert retriever.search_calls
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|