96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
"""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
|