363 lines
16 KiB
Python
363 lines
16 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 LexicalAwareRetriever(SectionAwareRetriever):
|
|
"""Adds `search_lexical`, scripted per test rather than doing real
|
|
text matching — this suite is about `_pooled_neighbour_hits`' bounding
|
|
and threshold logic, not the lexical scorer itself (see
|
|
`adapters/qdrant.py`'s own coverage for that)."""
|
|
|
|
def __init__(self, docs: list[RetrievalDocument], lexical_hits: list[SearchHit]) -> None:
|
|
super().__init__(docs)
|
|
self._lexical_hits = lexical_hits
|
|
self.lexical_calls: list[tuple[str, str]] = []
|
|
|
|
def search_lexical(self, query: str, drug_id: str, limit: int) -> list[SearchHit]: # noqa: ARG002
|
|
self.lexical_calls.append((query, drug_id))
|
|
return self._lexical_hits
|
|
|
|
|
|
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.")]
|
|
PRECAUTION = [_doc("t1", "than_trong", "Thận trọng với người suy thận.")]
|
|
INTERACTION = [_doc("x1", "tuong_tac_thuoc", "Tương tác với thuốc chống đông máu.")]
|
|
ALL_DOCS = CONTRA + INDICATION + PHARMACOLOGY + PRECAUTION + INTERACTION
|
|
|
|
|
|
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 TestSectionResolverResolveAll:
|
|
"""`resolve_all` — the fix for `abstain/incomplete_answer` reproduced
|
|
live 2026-08-14 on "Chỉ định và chống chỉ định của Aspirin là gì?":
|
|
`resolve()` silently picked one section, retrieval only fetched that
|
|
one, and the still-broad question failed the completeness check against
|
|
it. These pin that a genuine two-section question reports both, while a
|
|
substring collision (the exact case `resolve()` itself guards against)
|
|
still reports only one.
|
|
"""
|
|
|
|
def test_two_genuinely_named_sections_both_reported(self) -> None:
|
|
resolver = SectionResolver()
|
|
matches = resolver.resolve_all(
|
|
"Chỉ định và chống chỉ định của Aspirin là gì?"
|
|
)
|
|
assert {m.section_key for m in matches} == {"chi_dinh", "chong_chi_dinh"}
|
|
|
|
def test_substring_collision_is_not_double_counted(self) -> None:
|
|
""""chỉ định" is a literal substring of "chống chỉ định" — this must
|
|
stay a single match, exactly like `resolve()` already guarantees."""
|
|
resolver = SectionResolver()
|
|
matches = resolver.resolve_all("Chống chỉ định của aspirin là gì?")
|
|
assert [m.section_key for m in matches] == ["chong_chi_dinh"]
|
|
|
|
def test_three_sections_named_at_once(self) -> None:
|
|
resolver = SectionResolver()
|
|
matches = resolver.resolve_all(
|
|
"Liều dùng, chống chỉ định và tương tác thuốc của Metformin?"
|
|
)
|
|
assert {m.section_key for m in matches} == {
|
|
"lieu_luong_va_cach_dung", "chong_chi_dinh", "tuong_tac_thuoc",
|
|
}
|
|
|
|
def test_single_section_question_still_returns_one(self) -> None:
|
|
resolver = SectionResolver()
|
|
matches = resolver.resolve_all("Liều dùng metformin?")
|
|
assert [m.section_key for m in matches] == ["lieu_luong_va_cach_dung"]
|
|
|
|
def test_unrecognised_question_returns_empty(self) -> None:
|
|
assert SectionResolver().resolve_all("thuốc này giá bao nhiêu") == ()
|
|
assert SectionResolver().resolve_all("") == ()
|
|
|
|
|
|
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_lexically_strong_neighbour_section_is_pooled_in(self) -> None:
|
|
"""A precaution some drug's own "thận trọng" text never mentions can
|
|
still be filed under "chống chỉ định" (found live: Aspirin + loét dạ
|
|
dày). A neighbour section with real term overlap (score above
|
|
threshold) gets pooled in whole, not just the matching fragment."""
|
|
retriever = LexicalAwareRetriever(
|
|
ALL_DOCS, lexical_hits=[SearchHit(CONTRA[1], 6.0)], # c2: "chong_chi_dinh"
|
|
)
|
|
result = _service(retriever, SectionResolver()).retrieve(
|
|
"Thận trọng khi dùng aspirin cho bệnh nhân loét dạ dày là gì?", "aspirin"
|
|
)
|
|
assert retriever.section_calls == [
|
|
("aspirin", "than_trong"), ("aspirin", "chong_chi_dinh"),
|
|
]
|
|
assert retriever.lexical_calls == [
|
|
("Thận trọng khi dùng aspirin cho bệnh nhân loét dạ dày là gì?", "aspirin"),
|
|
]
|
|
# Whole section pooled (c1-c5), not just the one matching chunk (c2).
|
|
assert {item.evidence_id for item in result.evidence} == {
|
|
"t1", "c1", "c2", "c3", "c4", "c5",
|
|
}
|
|
|
|
def test_weak_lexical_match_is_not_pooled(self) -> None:
|
|
"""One incidental token overlap must not drag in a whole tangential
|
|
section — only real term overlap (>= threshold) counts."""
|
|
retriever = LexicalAwareRetriever(
|
|
ALL_DOCS, lexical_hits=[SearchHit(CONTRA[1], 4.0)],
|
|
)
|
|
result = _service(retriever, SectionResolver()).retrieve(
|
|
"Thận trọng khi dùng aspirin là gì?", "aspirin"
|
|
)
|
|
assert retriever.section_calls == [("aspirin", "than_trong")]
|
|
assert {item.evidence_id for item in result.evidence} == {"t1"}
|
|
|
|
def test_lexical_pooling_is_bounded(self) -> None:
|
|
"""At most 2 neighbour sections pool in, even if more score above
|
|
threshold — unbounded pooling is precision loss dressed as recall."""
|
|
retriever = LexicalAwareRetriever(
|
|
ALL_DOCS,
|
|
lexical_hits=[
|
|
SearchHit(CONTRA[0], 7.0), # chong_chi_dinh
|
|
SearchHit(INDICATION[0], 6.0), # chi_dinh
|
|
SearchHit(INTERACTION[0], 5.0), # tuong_tac_thuoc — 3rd, over bound
|
|
],
|
|
)
|
|
result = _service(retriever, SectionResolver()).retrieve(
|
|
"Thận trọng khi dùng aspirin là gì?", "aspirin"
|
|
)
|
|
assert retriever.section_calls == [
|
|
("aspirin", "than_trong"),
|
|
("aspirin", "chong_chi_dinh"),
|
|
("aspirin", "chi_dinh"),
|
|
]
|
|
# tuong_tac_thuoc (3rd-ranked lexical hit) never pooled — the bound
|
|
# (2 neighbours) already used by chong_chi_dinh + chi_dinh.
|
|
assert "x1" not in {item.evidence_id for item in result.evidence}
|
|
|
|
def test_generic_pharmacology_section_is_never_pooled(self) -> None:
|
|
"""`duoc_ly_va_co_che_tac_dung` is the corpus's documented
|
|
false-positive attractor (see `sections.py`) — excluded outright,
|
|
even with a lexical score that would otherwise clear the threshold
|
|
(measured live: it scored a close second right behind the real
|
|
answer on the exact query that motivated this whole mechanism)."""
|
|
retriever = LexicalAwareRetriever(
|
|
ALL_DOCS, lexical_hits=[SearchHit(PHARMACOLOGY[0], 9.0)],
|
|
)
|
|
result = _service(retriever, SectionResolver()).retrieve(
|
|
"Thận trọng khi dùng aspirin là gì?", "aspirin"
|
|
)
|
|
assert retriever.section_calls == [("aspirin", "than_trong")]
|
|
assert "p1" not in {item.evidence_id for item in result.evidence}
|
|
|
|
def test_retriever_without_lexical_search_still_works(self) -> None:
|
|
"""No `search_lexical` on the retriever (interface segregation, same
|
|
pattern as `find_by_section`/`find_by_drug`) — pooling is skipped,
|
|
not an error."""
|
|
retriever = SectionAwareRetriever(ALL_DOCS)
|
|
result = _service(retriever, SectionResolver()).retrieve(
|
|
"Thận trọng khi dùng aspirin là gì?", "aspirin"
|
|
)
|
|
assert retriever.section_calls == [("aspirin", "than_trong")]
|
|
assert {item.evidence_id for item in result.evidence} == {"t1"}
|
|
|
|
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
|