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

apps/ai-service/tests/test_citation_and_intro.py

@@ -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