duocthu/apps/ai-service/rag/answer.py

from __future__ import annotations

import json
from dataclasses import dataclass, replace

from . import grounding, metrics as metric_names
from .metrics import Metrics, NullMetrics
from .models import EvidenceDecision, QueryIntent, RetrievalResult, SubjectScope
from .ports import AnswerGenerationUnavailable, AnswerGenerator
from .prompt import build_request
from .routing import QueryRoutingService


@dataclass(frozen=True)
class Citation:
    chunk_id: str
    printed_page_start: int
    printed_page_end: int
    physical_page: int
    block_id: str | None = None
    bbox: tuple[float, float, float, float] | None = None
    source_crop: str | None = None
    attachment: str | None = None


@dataclass(frozen=True)
class GroundedAnswer:
    result: RetrievalResult
    answer: str | None
    citations: tuple[Citation, ...] = ()
    generated: bool = False


class GroundedAnswerService:
    """Retrieval decides what is true; generation only decides how it reads.

    When a generator is configured, its output replaces the extractive text
    **only** if `grounding.verify` confirms every figure and citation in it
    traces back to the retrieved evidence. Anything else — an unsupported
    number, a citation to nothing, a provider outage, malformed output — falls
    back to quoting the source verbatim, which is always available because it
    was computed first.
    """

    def __init__(
        self,
        routing: QueryRoutingService,
        generator: AnswerGenerator | None = None,
        metrics: Metrics | None = None,
    ) -> None:
        self._routing = routing
        self._generator = generator
        self._metrics = metrics or NullMetrics()

    def answer(
        self,
        query: str,
        subject_scope: SubjectScope,
        intent: QueryIntent,
    ) -> GroundedAnswer:
        result = self._routing.retrieve(query, subject_scope, intent)
        if result.decision == EvidenceDecision.ABSTAIN:
            self._metrics.increment(metric_names.ABSTENTION, reason=result.reason)
            return GroundedAnswer(result, None)

        citations = self._citations(result)
        if citations is None:
            return GroundedAnswer(
                replace(
                    result,
                    decision=EvidenceDecision.ABSTAIN,
                    reason="missing_printed_page_provenance",
                    evidence=(),
                ),
                None,
            )
        if result.decision == EvidenceDecision.VERIFY_PDF:
            # Never generated over. A quarantined table or formula is exactly
            # the evidence whose numbers were not reliably reconstructed, so
            # rephrasing it is the one case where fluency could invent a dose.
            return GroundedAnswer(
                result,
                "Nguồn có bảng hoặc công thức cần đối chiếu trực tiếp với ảnh PDF; "
                "không tự động trích số liệu.",
                citations,
            )

        evidence_texts = tuple(item.text for item in result.evidence)
        extractive = "\n\n".join(
            f"{text} [{index}]" for index, text in enumerate(evidence_texts, start=1)
        )

        generated = self._generate(query, evidence_texts)
        if generated is None:
            self._metrics.increment(metric_names.ANSWER_EXTRACTIVE)
            return GroundedAnswer(result, extractive, citations)

        self._metrics.increment(metric_names.GENERATION_SERVED)
        return GroundedAnswer(result, generated, citations, generated=True)

    def _generate(self, query: str, evidence_texts: tuple[str, ...]) -> str | None:
        """A verified generation, or None to fall back to the source text."""
        if self._generator is None or not evidence_texts:
            return None

        request = build_request(query, evidence_texts)
        try:
            raw = self._generator.generate(request.system, request.user, request.schema)
        except AnswerGenerationUnavailable:
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="provider_unavailable"
            )
            return None

        try:
            payload = json.loads(raw)
            answer = payload["answer"]
            sufficient = payload["evidence_sufficient"]
        except (ValueError, TypeError, KeyError):
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="malformed_output"
            )
            return None

        if not isinstance(answer, str) or not isinstance(sufficient, bool):
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="malformed_output"
            )
            return None
        if not sufficient:
            # The model says the evidence does not answer the question. Showing
            # the retrieved section verbatim lets the clinician judge that.
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason="evidence_insufficient"
            )
            return None

        report = grounding.verify(answer, evidence_texts)
        if not report.grounded:
            self._metrics.increment(
                metric_names.GENERATION_REJECTED, reason=report.reason
            )
            return None
        return answer

    @staticmethod
    def _citations(result: RetrievalResult) -> tuple[Citation, ...] | None:
        citations = []
        for evidence in result.evidence:
            if not evidence.source_refs:
                return None
            for source in evidence.source_refs:
                printed_range = source.printed_page_range
                if printed_range is not None:
                    start, end = printed_range
                elif source.printed_page is not None:
                    start = end = source.printed_page
                else:
                    return None
                citations.append(Citation(
                    chunk_id=evidence.matched_doc_id,
                    printed_page_start=int(start),
                    printed_page_end=int(end),
                    physical_page=source.physical_page,
                    block_id=source.block_id,
                    bbox=source.bbox,
                    source_crop=source.source_crop,
                    # Backward-compatible compact attachment identifier. A
                    # real crop path wins; otherwise the block id plus the
                    # structured page/bbox fields is enough to render later.
                    attachment=source.source_crop or source.block_id,
                ))
        return tuple(citations)