Fix live multi-turn: pass the resolved drug, stop did-you-mean garbage

2026-08-05 16:54:35 +07:00
parent ef08b4929e
commit 1e8cbdb586
29 changed files with 2013 additions and 83 deletions
@@ -0,0 +1,147 @@
+"""The Converse adapter must translate provider failure into the domain error,
+isolate the JSON envelope the API cannot enforce, and never leak an SDK type.
+
+No network: a stub client stands in for `bedrock-runtime`. The grounding check
+that actually guards correctness lives in `rag.grounding` and is tested there;
+here we prove only the adapter's envelope and failure contract.
+"""
+from __future__ import annotations
+
+import json
+
+import pytest
+
+from adapters.bedrock_converse import (
+    BedrockConverseAnswerGenerator,
+    BedrockCohereReranker,
+    RerankUnavailable,
+    _extract_json,
+)
+from rag.ports import AnswerGenerationUnavailable
+
+
+def _converse_reply(text: str, stop_reason: str = "end_turn") -> dict:
+    return {
+        "stopReason": stop_reason,
+        "output": {"message": {"content": [{"text": text}]}},
+        "usage": {"inputTokens": 10, "outputTokens": 5},
+    }
+
+
+class _Client:
+    def __init__(self, reply: dict) -> None:
+        self._reply = reply
+        self.calls: list[dict] = []
+
+    def converse(self, **kwargs):
+        self.calls.append(kwargs)
+        return self._reply
+
+
+def _payload(answer: str = "Liều 500 mg [1]", sufficient: bool = True) -> str:
+    return json.dumps({"answer": answer, "evidence_sufficient": sufficient}, ensure_ascii=False)
+
+
+def test_returns_the_models_json_on_success():
+    client = _Client(_converse_reply(_payload()))
+    gen = BedrockConverseAnswerGenerator(client=client, model_id="deepseek.v3.2")
+
+    raw = gen.generate("system", "user", {"type": "object"})
+
+    parsed = json.loads(raw)
+    assert parsed["answer"] == "Liều 500 mg [1]"
+    assert parsed["evidence_sufficient"] is True
+    # The prompt carries the schema and a JSON-only directive to the model.
+    sent = client.calls[0]["messages"][0]["content"][0]["text"]
+    assert "JSON" in sent
+
+
+def test_strips_a_markdown_fence_the_api_cannot_forbid():
+    fenced = f"```json\n{_payload()}\n```"
+    gen = BedrockConverseAnswerGenerator(client=_Client(_converse_reply(fenced)))
+
+    parsed = json.loads(gen.generate("s", "u", {}))
+
+    assert parsed["answer"] == "Liều 500 mg [1]"
+
+
+def test_extracts_json_when_the_model_adds_prose_around_it():
+    noisy = f"Đây là câu trả lời: {_payload()} Hết."
+    assert json.loads(_extract_json(noisy))["evidence_sufficient"] is True
+
+
+def test_a_filtered_stop_reason_is_an_outage_not_an_answer():
+    client = _Client(_converse_reply("", stop_reason="content_filtered"))
+    gen = BedrockConverseAnswerGenerator(client=client)
+
+    with pytest.raises(AnswerGenerationUnavailable):
+        gen.generate("s", "u", {})
+
+
+def test_empty_text_raises_rather_than_returning_blank():
+    gen = BedrockConverseAnswerGenerator(client=_Client(_converse_reply("   ")))
+
+    with pytest.raises(AnswerGenerationUnavailable):
+        gen.generate("s", "u", {})
+
+
+def test_real_botocore_client_error_becomes_the_domain_error():
+    botocore_exceptions = pytest.importorskip("botocore.exceptions")
+
+    class _Refusing:
+        def converse(self, **kwargs):
+            raise botocore_exceptions.ClientError(
+                {"Error": {"Code": "AccessDeniedException", "Message": "denied"}},
+                "Converse",
+            )
+
+    gen = BedrockConverseAnswerGenerator(client=_Refusing())
+
+    with pytest.raises(AnswerGenerationUnavailable):
+        gen.generate("s", "u", {})
+
+
+# --- reranker -----------------------------------------------------------------
+
+
+class _RerankBody:
+    def __init__(self, payload: dict) -> None:
+        self._data = json.dumps(payload).encode("utf-8")
+
+    def read(self) -> bytes:
+        return self._data
+
+
+class _RerankClient:
+    def __init__(self, payload: dict) -> None:
+        self._payload = payload
+
+    def invoke_model(self, **kwargs):
+        return {"body": _RerankBody(self._payload)}
+
+
+def test_rerank_returns_indices_most_relevant_first():
+    client = _RerankClient({"results": [{"index": 2}, {"index": 0}, {"index": 1}]})
+    reranker = BedrockCohereReranker(client=client)
+
+    order = reranker.rerank("chống chỉ định", ["a", "b", "c"])
+
+    assert order == [2, 0, 1]
+
+
+def test_rerank_raises_so_the_caller_keeps_original_order():
+    botocore_exceptions = pytest.importorskip("botocore.exceptions")
+
+    class _Refusing:
+        def invoke_model(self, **kwargs):
+            raise botocore_exceptions.ClientError(
+                {"Error": {"Code": "AccessDeniedException", "Message": "denied"}},
+                "InvokeModel",
+            )
+
+    with pytest.raises(RerankUnavailable):
+        BedrockCohereReranker(client=_Refusing()).rerank("q", ["a", "b"])
+
+
+def test_rerank_of_nothing_is_empty():
+    assert BedrockCohereReranker(client=_RerankClient({})).rerank("q", []) == []
@@ -0,0 +1,95 @@
+"""Two answer-UX fixes, pinned:
+
+- citations shown = only the sources the answer cited, not every retrieved chunk;
+- a bare drug name is introduced, not restated section-by-section.
+"""
+from __future__ import annotations
+
+import json
+
+from rag.answer import GroundedAnswerService
+from rag.models import (
+    Evidence,
+    EvidenceDecision,
+    QueryIntent,
+    RetrievalResult,
+    SourceRef,
+    SubjectScope,
+)
+from rag.prompt import build_request
+
+
+def _evidence(i: int, page: int) -> Evidence:
+    return Evidence(
+        evidence_id=f"drug::sec::{i}",
+        matched_doc_id=f"drug::sec::{i}",
+        kind="prose",
+        text=f"đoạn bằng chứng {i}",
+        score=1.0,
+        source_refs=(SourceRef(physical_page=page, precision="exact", printed_page=page),),
+        hydrated_from_parent=False,
+        requires_visual_check=False,
+    )
+
+
+class _Routing:
+    def __init__(self, result: RetrievalResult) -> None:
+        self._result = result
+
+    def retrieve(self, query, subject_scope, intent):  # noqa: ARG002
+        return self._result
+
+
+class _Generator:
+    def __init__(self, payload: dict) -> None:
+        self._payload = payload
+
+    def generate(self, system: str, user: str, schema: dict) -> str:  # noqa: ARG002
+        return json.dumps(self._payload, ensure_ascii=False)
+
+
+def _answerable(*evidence: Evidence, is_overview: bool = False) -> RetrievalResult:
+    return RetrievalResult(
+        EvidenceDecision.ANSWERABLE,
+        "grounded_evidence_available",
+        tuple(evidence),
+        resolved_drug_id="drug",
+        is_drug_overview=is_overview,
+    )
+
+
+def test_only_cited_sources_are_returned():
+    result = _answerable(_evidence(0, 100), _evidence(1, 200), _evidence(2, 300))
+    service = GroundedAnswerService(
+        _Routing(result),
+        _Generator({"answer": "Chỉ dùng đoạn hai [2].", "evidence_sufficient": True}),
+    )
+
+    grounded = service.answer("q", SubjectScope.HUMAN, QueryIntent.FACT_LOOKUP)
+
+    assert grounded.generated is True
+    assert len(grounded.citations) == 1
+    assert grounded.citations[0].printed_page_start == 200
+
+
+def test_answer_citing_nothing_falls_back_to_all_citations():
+    result = _answerable(_evidence(0, 100), _evidence(1, 200))
+    service = GroundedAnswerService(
+        _Routing(result),
+        # no [n] marker at all: rather than show zero provenance, show all.
+        _Generator({"answer": "Không có trích dẫn.", "evidence_sufficient": True}),
+    )
+
+    grounded = service.answer("q", SubjectScope.HUMAN, QueryIntent.FACT_LOOKUP)
+
+    assert len(grounded.citations) == 2
+
+
+def test_bare_name_builds_an_intro_prompt():
+    intro = build_request("PARACETAMOL", ("đoạn A", "đoạn B"), intro=True)
+    assert "GIỚI THIỆU" in intro.user
+    assert "CÂU HỎI:" not in intro.user
+
+    normal = build_request("Liều?", ("đoạn A",), intro=False)
+    assert "CÂU HỎI:" in normal.user
+    assert "GIỚI THIỆU" not in normal.user
@@ -31,8 +31,8 @@ class FakeAnswers:
        self._e = evidence_text
        self.calls = []

-    def answer(self, query, subject_scope, intent):
-        self.calls.append(query)
+    def answer(self, query, subject_scope, intent, drug_id=None):
+        self.calls.append((query, drug_id))
        return _grounded(self._a, self._e)


@@ -66,7 +66,7 @@ def test_medical_turn_returns_grounded_answer():
    assert out.grounded is not None


-def test_followup_inherits_drug_and_names_it_and_rewrites_query():
+def test_followup_inherits_drug_and_passes_it_resolved():
    answers = FakeAnswers(
        "Ở trẻ em điều chỉnh theo cân nặng.",
        "Ở trẻ em, liều metformin điều chỉnh theo cân nặng.",
@@ -76,12 +76,40 @@ def test_followup_inherits_drug_and_names_it_and_rewrites_query():
    out = svc.answer("c3", "còn trẻ em thì sao?", SubjectScope.HUMAN, QueryIntent.FACT_LOOKUP)
    assert out.inherited_drug == "metformin"
    assert out.answer.startswith("Về metformin:")
-    # The follow-up was rewritten self-contained before hitting the engine.
-    assert "metformin" in answers.calls[-1]
+    # The inherited drug is passed already-resolved (not re-resolved from the
+    # rewritten turn text), and the raw follow-up drives section routing.
+    last_query, last_drug_id = answers.calls[-1]
+    assert last_drug_id == "metformin"
+    assert last_query == "còn trẻ em thì sao?"
    # State carried the drug forward.
    assert svc._store.load("c3").focus.drug_id == "metformin"


+def test_confirmation_is_not_fuzzy_matched_to_a_drug():
+    """'đúng' must not be fuzzy-matched to terbinafin/tretinoin (the did-you-mean
+    loop the reviewer hit); it asks which drug instead."""
+    answers = FakeAnswers("x", "x")
+    svc = _service(answers)
+    out = svc.answer("cc", "đúng", SubjectScope.HUMAN, QueryIntent.FACT_LOOKUP)
+    assert out.clarification is not None
+    assert out.clarification.reason == "confirm_without_context"
+    assert answers.calls == []
+
+
+def test_a_long_sentence_that_names_no_drug_is_not_offered_did_you_mean():
+    """A full question ('EPO điều trị thiếu máu...') that resolves no drug is
+    answered honestly, not with garbage suggestions from fuzzing the sentence."""
+    answers = FakeAnswers("x", "x")
+    svc = _service(answers)  # catalog holds only metformin
+    out = svc.answer(
+        "cl", "EPO điều trị thiếu máu do hóa trị ung thư liều khởi đầu bao nhiêu",
+        SubjectScope.HUMAN, QueryIntent.FACT_LOOKUP,
+    )
+    assert out.clarification is not None
+    assert out.clarification.reason == "drug_not_supported"
+    assert answers.calls == []
+
+
 def test_no_close_drug_reports_not_supported():
    answers = FakeAnswers("x", "x")
    svc = _service(answers)  # catalog holds only metformin
@@ -0,0 +1,92 @@
+"""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"]