211 lines
7.9 KiB
Python
211 lines
7.9 KiB
Python
from adapters.qdrant import QdrantRetriever, _source_refs
|
|
|
|
|
|
class _FakePoint:
|
|
def __init__(self, payload: dict) -> None:
|
|
self.payload = payload
|
|
self.score = 1.0
|
|
|
|
|
|
class _FakeScrollClient:
|
|
"""Mimics qdrant-client's `.scroll()` shape closely enough to exercise
|
|
`find_by_indication`'s keyword-matching logic directly — a fake filter
|
|
(not a real one), so it returns every payload handed to it regardless
|
|
of `scroll_filter`; the payloads given in each test already represent
|
|
what a real `section_key=chi_dinh, chunk_kind=prose` filter would have
|
|
returned, which is the part `find_by_indication` cannot get wrong on
|
|
its own (the filter construction itself is a one-line, inspectable
|
|
`Filter(must=[...])` — not worth a second fake layer to prove)."""
|
|
|
|
def __init__(self, payloads: list[dict]) -> None:
|
|
self._payloads = payloads
|
|
|
|
def scroll(self, collection_name, scroll_filter, limit, offset, with_payload): # noqa: ARG002
|
|
return [_FakePoint(p) for p in self._payloads], None
|
|
|
|
|
|
def _chi_dinh_payload(drug_id: str, text: str) -> dict:
|
|
return {
|
|
"chunk_id": f"{drug_id}__chi_dinh__0", "drug_id": drug_id,
|
|
"drug_name": drug_id.upper(), "section_key": "chi_dinh",
|
|
"chunk_kind": "prose", "text": text,
|
|
"heading_physical_page": 100, "printed_page_range": [101, 101],
|
|
}
|
|
|
|
|
|
def test_descriptor_source_ref_comes_from_attachment_not_heading_page():
|
|
refs = _source_refs({
|
|
"chunk_kind": "block_descriptor",
|
|
"heading_physical_page": 208,
|
|
"source_page_range": [209, 209],
|
|
"printed_page_range": [210, 210],
|
|
"attachments": [{
|
|
"block_id": "p209_t0",
|
|
"physical_page": 209,
|
|
"printed_page": 210,
|
|
"bbox": [49.5, 68.1, 289.4, 789.4],
|
|
"source_crop": "crops/p209_t0.png",
|
|
}],
|
|
})
|
|
|
|
assert len(refs) == 1
|
|
assert refs[0].physical_page == 209
|
|
assert refs[0].printed_page == 210
|
|
assert refs[0].block_id == "p209_t0"
|
|
assert refs[0].bbox == (49.5, 68.1, 289.4, 789.4)
|
|
assert refs[0].source_crop == "crops/p209_t0.png"
|
|
assert refs[0].precision == "region"
|
|
|
|
|
|
def test_prose_ref_uses_exact_chunk_range_and_keeps_attachment_region():
|
|
refs = _source_refs({
|
|
"chunk_kind": "prose",
|
|
"heading_physical_page": 100,
|
|
"source_page_range": [104, 105],
|
|
"printed_page_range": [105, 106],
|
|
"attachments": [{
|
|
"block_id": "p105_t0",
|
|
"physical_page": 105,
|
|
"printed_page": 106,
|
|
"bbox": [1.0, 2.0, 3.0, 4.0],
|
|
}],
|
|
})
|
|
|
|
assert refs[0].physical_page == 104
|
|
assert refs[0].page_range == (104, 105)
|
|
assert refs[0].printed_page_range == (105, 106)
|
|
assert refs[1].block_id == "p105_t0"
|
|
assert refs[1].physical_page == 105
|
|
assert refs[1].printed_page == 106
|
|
|
|
|
|
def test_find_by_indication_matches_a_drug_that_names_the_symptom():
|
|
client = _FakeScrollClient([
|
|
_chi_dinh_payload("paracetamol_acetaminophen", "Điều trị sốt, giảm đau nhẹ và vừa."),
|
|
_chi_dinh_payload("amoxicilin", "Điều trị nhiễm khuẩn đường hô hấp."),
|
|
])
|
|
retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
|
|
|
|
hits = retriever.find_by_indication("sốt", limit=8)
|
|
|
|
assert [h.document.drug_id for h in hits] == ["paracetamol_acetaminophen"]
|
|
|
|
|
|
def test_find_by_indication_requires_the_whole_phrase_not_a_scattered_match():
|
|
""""sốt xuất huyết" (dengue) must not match a chunk that only says "sốt"
|
|
— the phrase itself has to appear, not just each of its words somewhere."""
|
|
client = _FakeScrollClient([
|
|
_chi_dinh_payload("paracetamol_acetaminophen", "Điều trị sốt, giảm đau."),
|
|
])
|
|
retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
|
|
|
|
hits = retriever.find_by_indication("sốt xuất huyết", limit=8)
|
|
|
|
assert hits == []
|
|
|
|
|
|
def test_find_by_indication_matches_a_multi_word_phrase_contiguously():
|
|
client = _FakeScrollClient([
|
|
_chi_dinh_payload("paracetamol_acetaminophen", "Điều trị sốt cao ở người lớn."),
|
|
])
|
|
retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
|
|
|
|
hits = retriever.find_by_indication("sốt cao", limit=8)
|
|
|
|
assert [h.document.drug_id for h in hits] == ["paracetamol_acetaminophen"]
|
|
|
|
|
|
def test_find_by_indication_rejects_a_scattered_bag_of_common_words():
|
|
"""Found live 2026-08-07: a token-SUBSET match (every word present
|
|
*somewhere*, any order) let a long nonsense phrase built from common
|
|
filler words false-positive against real chi_dinh text — the words are
|
|
common enough to appear scattered through nearly anything. Phrase
|
|
matching closes it: none of these words are contiguous in the target
|
|
text the way they are in the query."""
|
|
client = _FakeScrollClient([
|
|
_chi_dinh_payload(
|
|
"paracetamol_acetaminophen",
|
|
"Điều trị sốt. Không dùng quá liều khuyến cáo trong sách hướng dẫn.",
|
|
),
|
|
])
|
|
retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
|
|
|
|
hits = retriever.find_by_indication(
|
|
"bệnh chưa từng ghi nhận trong sách abcxyz123", limit=8
|
|
)
|
|
|
|
assert hits == []
|
|
|
|
|
|
def test_find_by_indication_returns_chunk_pool_for_service_level_drug_aggregation():
|
|
client = _FakeScrollClient([
|
|
_chi_dinh_payload("paracetamol_acetaminophen", "Điều trị sốt."),
|
|
{**_chi_dinh_payload("paracetamol_acetaminophen", "Điều trị sốt cao ở trẻ em."),
|
|
"chunk_id": "paracetamol_acetaminophen__chi_dinh__1"},
|
|
])
|
|
retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
|
|
|
|
hits = retriever.find_by_indication("sốt", limit=8)
|
|
|
|
assert len(hits) == 2
|
|
assert {hit.document.drug_id for hit in hits} == {"paracetamol_acetaminophen"}
|
|
|
|
|
|
def test_find_by_indication_respects_the_limit():
|
|
client = _FakeScrollClient([
|
|
_chi_dinh_payload(f"drug_{i}", "Điều trị đau.") for i in range(5)
|
|
])
|
|
retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
|
|
|
|
hits = retriever.find_by_indication("đau", limit=2)
|
|
|
|
assert len(hits) == 2
|
|
|
|
|
|
def _section_payload(drug_id: str, section_key: str, text: str) -> dict:
|
|
return {
|
|
"chunk_id": f"{drug_id}__{section_key}__0", "drug_id": drug_id,
|
|
"drug_name": drug_id.upper(), "section_key": section_key,
|
|
"chunk_kind": "prose", "text": text,
|
|
"heading_physical_page": 100, "printed_page_range": [101, 101],
|
|
}
|
|
|
|
|
|
def test_search_lexical_ranks_by_distinct_matched_term_count():
|
|
client = _FakeScrollClient([
|
|
_section_payload("aspirin", "than_trong", "Thận trọng với suy thận."),
|
|
_section_payload(
|
|
"aspirin", "chong_chi_dinh",
|
|
"Không dùng cho người có loét dạ dày tá tràng đang hoạt động.",
|
|
),
|
|
])
|
|
retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
|
|
|
|
hits = retriever.search_lexical(
|
|
"Thận trọng khi dùng aspirin cho bệnh nhân loét dạ dày", "aspirin", limit=5
|
|
)
|
|
|
|
assert [h.document.section_key for h in hits] == ["chong_chi_dinh", "than_trong"]
|
|
assert hits[0].score > hits[1].score
|
|
|
|
|
|
def test_search_lexical_drops_stopword_only_queries():
|
|
client = _FakeScrollClient([
|
|
_section_payload("aspirin", "than_trong", "Thận trọng với suy thận."),
|
|
])
|
|
retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
|
|
|
|
assert retriever.search_lexical("là gì và của", "aspirin", limit=5) == []
|
|
|
|
|
|
def test_search_lexical_excludes_non_matching_sections():
|
|
client = _FakeScrollClient([
|
|
_section_payload("aspirin", "than_trong", "Thận trọng với suy thận."),
|
|
_section_payload("aspirin", "chi_dinh", "Giảm đau hạ sốt."),
|
|
])
|
|
retriever = QdrantRetriever(client, "duocthu_v1", embedder=None)
|
|
|
|
hits = retriever.search_lexical("loét dạ dày", "aspirin", limit=5)
|
|
|
|
assert hits == []
|