duocthu/apps/ai-service/tests/test_grounded_generation.py

"""The answer layer may rephrase evidence; it may not add to it.

Every test here is a fabrication the generator could plausibly produce, and
the assertion is that the clinician never sees it. The dose figures are taken
from the real METFORMIN and PARACETAMOL sections in `duocthu_v1`.
"""
from __future__ import annotations

import json

import pytest

from rag import grounding
from rag.answer import GroundedAnswerService
from rag.metrics import GENERATION_REJECTED, GENERATION_SERVED, InMemoryMetrics
from rag.models import (
    Evidence,
    EvidenceDecision,
    QueryIntent,
    RetrievalResult,
    SourceRef,
    SubjectScope,
)
from rag.ports import AnswerGenerationUnavailable
from rag.prompt import build_request

SOURCE = SourceRef(
    physical_page=812,
    precision="region",
    printed_page_range=(714, 714),
)

EVIDENCE_TEXT = (
    "Người lớn: uống 500 mg metformin hydroclorid, 2 lần mỗi ngày. "
    "Liều tối đa 2 g mỗi ngày, chia làm nhiều lần."
)


def _result(text: str = EVIDENCE_TEXT) -> RetrievalResult:
    return RetrievalResult(
        EvidenceDecision.ANSWERABLE,
        "grounded_evidence_available",
        (
            Evidence(
                evidence_id="metformin::lieu::0",
                matched_doc_id="metformin::lieu::0",
                kind="prose",
                text=text,
                score=1.0,
                source_refs=(SOURCE,),
                hydrated_from_parent=False,
                requires_visual_check=False,
            ),
        ),
        resolved_drug_id="metformin",
    )


class _FixedRouting:
    def __init__(self, result: RetrievalResult) -> None:
        self._result = result

    def retrieve(self, query, subject_scope, intent):
        return self._result


class _Generator:
    """Returns whatever payload the test wants the model to have produced."""

    def __init__(self, payload) -> None:
        self._payload = payload

    def generate(self, system: str, user: str, schema: dict) -> str:
        if isinstance(self._payload, BaseException):
            raise self._payload
        if isinstance(self._payload, str):
            return self._payload
        return json.dumps(self._payload, ensure_ascii=False)


def _answer(payload, result: RetrievalResult | None = None):
    metrics = InMemoryMetrics()
    service = GroundedAnswerService(
        _FixedRouting(result or _result()), _Generator(payload), metrics
    )
    grounded = service.answer("Liều Metformin?", SubjectScope.HUMAN, QueryIntent.FACT_LOOKUP)
    return grounded, metrics


# --- the guardrail's whole reason to exist ------------------------------------


def test_invented_dose_is_refused_and_never_reaches_the_answer():
    grounded, metrics = _answer(
        {"answer": "Người lớn uống 850 mg, 2 lần mỗi ngày [1].",
         "evidence_sufficient": True}
    )

    assert grounded.generated is False
    assert "850" not in grounded.answer
    assert grounded.answer.startswith(EVIDENCE_TEXT)
    assert metrics.total(GENERATION_REJECTED, reason="ungrounded_number") == 1


def test_a_rounded_figure_counts_as_invented():
    """`2 g` is in the source; `2000 mg` is a conversion, and conversions are
    where unit errors live. The prompt forbids it and the check enforces it."""
    grounded, metrics = _answer(
        {"answer": "Liều tối đa 2000 mg mỗi ngày [1].", "evidence_sufficient": True}
    )

    assert grounded.generated is False
    assert metrics.total(GENERATION_REJECTED, reason="ungrounded_number") == 1


def test_citation_pointing_at_nothing_is_refused():
    grounded, metrics = _answer(
        {"answer": "Người lớn uống 500 mg [3].", "evidence_sufficient": True}
    )

    assert grounded.generated is False
    assert metrics.total(GENERATION_REJECTED, reason="invalid_citation") == 1


def test_faithful_rewrite_is_served():
    grounded, metrics = _answer(
        {"answer": "Người lớn: 500 mg, 2 lần/ngày; tối đa 2 g/ngày [1].",
         "evidence_sufficient": True}
    )

    assert grounded.generated is True
    assert grounded.answer == "Người lớn: 500 mg, 2 lần/ngày; tối đa 2 g/ngày [1]."
    assert metrics.total(GENERATION_SERVED) == 1
    assert metrics.total(GENERATION_REJECTED) == 0


def test_citations_survive_generation():
    """Provenance is the point; a prettier answer must not cost the folio."""
    grounded, _ = _answer(
        {"answer": "Người lớn: 500 mg [1].", "evidence_sufficient": True}
    )

    assert grounded.generated is True
    assert len(grounded.citations) == 1
    assert grounded.citations[0].printed_page_start == 714


# --- degradation is always to the source, never to an error -------------------


@pytest.mark.parametrize(
    "payload, reason",
    [
        (AnswerGenerationUnavailable("revoked"), "provider_unavailable"),
        ("not json at all", "malformed_output"),
        ({"answer": "500 mg [1]"}, "malformed_output"),
        ({"answer": 500, "evidence_sufficient": True}, "malformed_output"),
        ({"answer": "...", "evidence_sufficient": False}, "evidence_insufficient"),
    ],
)
def test_every_generation_failure_falls_back_to_the_source_text(payload, reason):
    grounded, metrics = _answer(payload)

    assert grounded.generated is False
    assert grounded.answer.startswith(EVIDENCE_TEXT)
    assert metrics.total(GENERATION_REJECTED, reason=reason) == 1


def test_no_generator_configured_still_answers():
    service = GroundedAnswerService(_FixedRouting(_result()))

    grounded = service.answer(
        "Liều Metformin?", SubjectScope.HUMAN, QueryIntent.FACT_LOOKUP
    )

    assert grounded.generated is False
    assert grounded.answer.startswith(EVIDENCE_TEXT)


# --- the comparison rule itself ----------------------------------------------


def test_decimal_separators_are_not_interchangeable():
    """`7,5` and `7.5` differ, and so do `7,5` and `75`. Normalising them
    together is how a tenfold dose error scores as a match."""
    source = ("Sơ sinh: 7,5 mg/kg cách 8 giờ/lần.",)

    assert grounding.verify("7,5 mg/kg [1]", source).grounded is True
    assert grounding.verify("7.5 mg/kg [1]", source).grounded is False
    assert grounding.verify("75 mg/kg [1]", source).grounded is False


def test_citation_markers_are_not_read_as_quantities():
    report = grounding.verify("Không dùng cho người suy thận [1].", ("Suy thận.",))

    assert report.grounded is True
    assert report.cited_indices == (1,)


def test_prompt_numbers_evidence_from_one():
    request = build_request("Liều?", ("đoạn A", "đoạn B"))

    assert "[1] đoạn A" in request.user
    assert "[2] đoạn B" in request.user
    assert "CHÉP NGUYÊN VĂN" in request.system


def test_prompt_refuses_to_build_without_evidence():
    with pytest.raises(ValueError):
        build_request("Liều?", ())