93 lines
3.0 KiB
Python
93 lines
3.0 KiB
Python
"""A free-form question about a resolved drug must not dump the whole monograph
|
|
at the model. When a reranker is configured, the overview is reordered by
|
|
relevance and trimmed to top-k; a bare drug name still returns everything.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from rag.models import EvidenceDecision, RetrievalDocument, SearchHit, SourceRef
|
|
from rag.ports import RerankUnavailable
|
|
from rag.service import EvidencePolicy, RetrievalService
|
|
|
|
|
|
def _hit(i: int) -> SearchHit:
|
|
doc = RetrievalDocument(
|
|
doc_id=f"d{i}",
|
|
drug_id="paracetamol",
|
|
kind="prose",
|
|
text=f"section {i} text",
|
|
section_key=f"sec_{i}",
|
|
source_refs=(
|
|
SourceRef(physical_page=100 + i, precision="exact", printed_page=i),
|
|
),
|
|
)
|
|
return SearchHit(document=doc, score=1.0)
|
|
|
|
|
|
class _Retriever:
|
|
def __init__(self, n: int) -> None:
|
|
self._hits = [_hit(i) for i in range(n)]
|
|
|
|
def find_by_drug(self, drug_id: str) -> list[SearchHit]:
|
|
return list(self._hits)
|
|
|
|
def search(self, query: str, drug_id: str, limit: int) -> list[SearchHit]: # noqa: ARG002
|
|
return []
|
|
|
|
|
|
class _ParentStore:
|
|
def get(self, parent_id: str): # noqa: ARG002
|
|
return None
|
|
|
|
|
|
class _Reranker:
|
|
def __init__(self, order: list[int], fail: bool = False) -> None:
|
|
self._order = order
|
|
self._fail = fail
|
|
self.calls: list[str] = []
|
|
|
|
def rerank(self, query: str, documents: list[str], top_n=None): # noqa: ARG002
|
|
self.calls.append(query)
|
|
if self._fail:
|
|
raise RerankUnavailable("provider down")
|
|
return self._order
|
|
|
|
|
|
def _service(reranker=None, top_k=3) -> RetrievalService:
|
|
return RetrievalService(
|
|
_Retriever(8),
|
|
_ParentStore(),
|
|
EvidencePolicy(rerank_top_k=top_k),
|
|
section_resolver=None,
|
|
reranker=reranker,
|
|
)
|
|
|
|
|
|
def test_a_question_reranks_the_overview_and_keeps_top_k():
|
|
reranker = _Reranker(order=[7, 6, 5, 4, 3, 2, 1, 0])
|
|
result = _service(reranker, top_k=3).retrieve(
|
|
"sốt cao uống được không", "paracetamol"
|
|
)
|
|
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
assert reranker.calls, "reranker should run on a free-form question"
|
|
assert len(result.evidence) == 3
|
|
assert [e.matched_doc_id for e in result.evidence] == ["d7", "d6", "d5"]
|
|
|
|
|
|
def test_a_bare_drug_name_returns_the_whole_monograph_unreranked():
|
|
reranker = _Reranker(order=[0])
|
|
result = _service(reranker).retrieve("PARACETAMOL", "paracetamol")
|
|
|
|
assert reranker.calls == [], "a bare name must not be reranked/trimmed"
|
|
assert len(result.evidence) == 8
|
|
|
|
|
|
def test_rerank_outage_keeps_the_original_order():
|
|
reranker = _Reranker(order=[], fail=True)
|
|
result = _service(reranker).retrieve("sốt cao uống được không", "paracetamol")
|
|
|
|
# Fail-open: the answer survives, in book order, when rerank is unreachable.
|
|
assert result.decision == EvidenceDecision.ANSWERABLE
|
|
assert len(result.evidence) == 8
|
|
assert [e.matched_doc_id for e in result.evidence][:2] == ["d0", "d1"]
|