app/rag/retrieval_service.py

"""Hybrid, structure-aware retrieval over canonical academic evidence."""

from __future__ import annotations

import re
from collections import defaultdict
from dataclasses import dataclass
from typing import Sequence

from app.observability.operation import observe_operation, observe_stage
from app.observability.config import get_observability_config
from app.observability.sanitize import sanitize_attribute_value
from app.rag.evidence_store import EvidenceStore, LexicalCandidate
from app.rag.models import (
    ConsumerType,
    EvidenceCitation,
    EvidenceRequest,
    EvidenceType,
    RetrievedEvidence,
    RetrievalDiagnostics,
)
from app.rag.vector_indexes import PaperEvidenceIndex, VectorCandidate
from app.rag.project_memory import ProjectMemoryStore
from app.rag.reranker import EvidenceReranker
from app.rag.pipeline_version import CURRENT_PRODUCT_PIPELINE_VERSION


_RRF_K = 60
_STRUCTURAL_RELATIONS = {"caption_of", "describes", "parent", "continuation", "same_table"}
_SYNTHESIS_CONSUMERS = {
    ConsumerType.WIKI,
    ConsumerType.DRAFT,
    ConsumerType.REPORT,
    ConsumerType.PROJECT_GRAPH,
    ConsumerType.PAPER_GRAPH,
}
_VISUAL_TYPES = {
    EvidenceType.FIGURE,
    EvidenceType.PLOT,
    EvidenceType.DIAGRAM,
    EvidenceType.TABLE,
    EvidenceType.FORMULA,
    EvidenceType.CAPTION,
}
_REFERENCE_SECTION_RE = re.compile(r"^(references|bibliography|works cited)\b", re.IGNORECASE)


@dataclass(frozen=True)
class SourceRetrievalResult:
    evidence: list[RetrievedEvidence]
    citations: list[EvidenceCitation]
    diagnostics: RetrievalDiagnostics


class PaperEvidenceService:
    """The only product-facing source-evidence retrieval boundary."""

    def __init__(
        self,
        store: EvidenceStore | None = None,
        index: PaperEvidenceIndex | None = None,
        reranker: EvidenceReranker | None = None,
    ) -> None:
        self.store = store or EvidenceStore()
        self.index = index or PaperEvidenceIndex()
        self.reranker = reranker or EvidenceReranker()

    def retrieve(self, request: EvidenceRequest) -> SourceRetrievalResult:
        consumer = request.consumer.value
        diagnostics = RetrievalDiagnostics()
        with observe_operation(
            "rag.retrieve",
            subsystem="retrieval",
            consumer=consumer,
            attributes={
                "query_chars": len(request.query),
                "document_filter_count": len(request.document_ids),
                "anchor_request_count": len(request.anchor_evidence_ids) + len(request.selection_anchors),
                "token_budget": request.token_budget,
                "pipeline.version": CURRENT_PRODUCT_PIPELINE_VERSION,
            },
        ) as op:
            obs_config = get_observability_config()
            if obs_config.capture_content:
                op.event(
                    "local_diagnostic_query",
                    {"query_preview": sanitize_attribute_value(request.query[:240])},
                )
            with observe_stage(op, "anchor_resolution", subsystem="retrieval", consumer=consumer):
                anchor_ids = self._resolve_anchors(request)
            diagnostics.anchor_count = len(anchor_ids)

            lexical: list[LexicalCandidate] = []
            try:
                with observe_stage(op, "lexical_search", subsystem="retrieval", consumer=consumer):
                    lexical = self.store.lexical_search(
                        request.project_id,
                        request.query,
                        limit=40,
                        document_ids=request.document_ids or None,
                    )
            except Exception as exc:
                # SQLite/FTS is one independent retrieval channel. Preserve
                # dense/anchor recovery while making the handled degradation
                # visible on the parent request.
                op.record_exception(
                    exc,
                    escaped=False,
                    stage="lexical_search",
                    category="lexical_search_fallback",
                )
                diagnostics.terminal_state = "degraded"
                op.mark_terminal("degraded")
            diagnostics.lexical_candidate_count = len(lexical)

            dense: list[VectorCandidate] = []
            try:
                # This parent stage deliberately contains both the existing
                # embedding and Chroma child spans; do not add a duplicate
                # query-embedding operation here.
                with observe_stage(op, "dense_search", subsystem="retrieval", consumer=consumer):
                    dense = self.index.search(
                        request.project_id,
                        request.query,
                        limit=40,
                        document_ids=request.document_ids or None,
                        evidence_types=request.modalities or None,
                        consumer=consumer,
                    )
            except Exception as exc:  # lexical/anchor retrieval remains usable
                if not lexical and not anchor_ids:
                    raise
                diagnostics.warnings.append(f"dense_search_fallback:{type(exc).__name__}")
                op.record_exception(
                    exc,
                    escaped=False,
                    stage="dense_search",
                    category="dense_search_fallback",
                )
                diagnostics.terminal_state = "degraded"
                op.mark_terminal("degraded")
            diagnostics.dense_candidate_count = len(dense)

            with observe_stage(op, "rank_fusion", subsystem="retrieval", consumer=consumer):
                fused_ids, scores, reasons = _fuse(anchor_ids, lexical, dense)
            diagnostics.duplicates_removed = max(
                0,
                len(anchor_ids) + len(lexical) + len(dense) - len(fused_ids),
            )
            units = self.store.get_units(request.project_id, evidence_ids=fused_ids)
            unit_map = {unit.evidence_id: unit for unit in units}
            if request.modalities:
                unit_map = {
                    evidence_id: unit
                    for evidence_id, unit in unit_map.items()
                    if unit.element_type in request.modalities or evidence_id in anchor_ids
                }

            with observe_stage(op, "structural_expansion", subsystem="retrieval", consumer=consumer):
                expanded_ids = self._expand_structure(request.project_id, list(unit_map), limit=16)
                missing_ids = [evidence_id for evidence_id in expanded_ids if evidence_id not in unit_map]
                for unit in self.store.get_units(request.project_id, evidence_ids=missing_ids):
                    if request.document_ids and unit.document_id not in request.document_ids:
                        continue
                    unit_map[unit.evidence_id] = unit
                    scores[unit.evidence_id] = max(scores.get(unit.evidence_id, 0.0), 0.015)
                    reasons[unit.evidence_id].append("structural_neighbor")

            ranked = self._rank(request, unit_map, scores, reasons, anchor_ids, lexical, dense)
            with observe_stage(
                op, "reranking", subsystem="retrieval", consumer=consumer
            ) as rerank_stage:
                decision = self.reranker.decide(request, ranked)
                diagnostics.rerank_reason = decision.reason
                if decision.use:
                    try:
                        ranked = self.reranker.rerank(request.query, ranked)
                        diagnostics.rerank_used = True
                    except Exception as exc:  # local model failure must preserve fused retrieval
                        rerank_stage.mark_error(type(exc).__name__)
                        diagnostics.warnings.append(f"rerank_fallback:{type(exc).__name__}")
                        diagnostics.rerank_reason = "model_failure_fallback"
                        op.record_exception(
                            exc,
                            escaped=False,
                            stage="reranking",
                            category="rerank_fallback",
                        )
                        diagnostics.terminal_state = "degraded"
                        op.mark_terminal("degraded")

            with observe_stage(op, "diversity_selection", subsystem="retrieval", consumer=consumer):
                selected, truncated = _select_diverse(ranked, request)
            diagnostics.context_truncated = truncated
            diagnostics.fused_candidate_count = len(ranked)

            with observe_stage(op, "context_packing", subsystem="retrieval", consumer=consumer):
                diagnostics.documents_represented = len({item.evidence.document_id for item in selected})
                diagnostics.sections_represented = len(
                    {(item.evidence.document_id, tuple(item.evidence.section_path)) for item in selected}
                )
                diagnostics.source_context_chars = sum(len(item.evidence.index_text) for item in selected)
            if not selected and diagnostics.terminal_state == "success":
                diagnostics.terminal_state = "success_empty"
                op.mark_terminal("success_empty")

            op.add_count("query_chars", len(request.query))
            op.add_count("anchor_count", diagnostics.anchor_count)
            op.add_count("lexical_candidate_count", diagnostics.lexical_candidate_count)
            op.add_count("dense_candidate_count", diagnostics.dense_candidate_count)
            op.add_count("fused_candidate_count", diagnostics.fused_candidate_count)
            op.add_count("returned_evidence_count", len(selected))
            op.add_count("duplicates_removed", diagnostics.duplicates_removed)
            op.add_count("documents_represented", diagnostics.documents_represented)
            op.add_count("sections_represented", diagnostics.sections_represented)
            op.add_count("context_token_budget", request.token_budget)
            op.add_count("context_truncated", int(diagnostics.context_truncated))
            op.add_count("source_context_chars", diagnostics.source_context_chars)
            op.add_count("empty_result", int(not selected))
            op.set("rerank_used", diagnostics.rerank_used)
            op.set("rerank_reason", diagnostics.rerank_reason)
            for evidence_type in EvidenceType:
                count = sum(1 for item in selected if item.evidence.element_type is evidence_type)
                if count:
                    op.add_count(f"returned_type.{evidence_type.value}", count)
            with observe_stage(op, "citation_materialization", subsystem="retrieval", consumer=consumer):
                citations = self._citations(request.project_id, selected)
            diagnostics.stage_ms = op.stage_durations_ms
            return SourceRetrievalResult(evidence=selected, citations=citations, diagnostics=diagnostics)

    def ground_graph_nodes(
        self,
        project_id: str,
        requests: Sequence[tuple[str, str, Sequence[str]]],
    ) -> tuple[dict[str, list[str]], dict[str, list[str]]]:
        """Ground a graph in one evidence read and one memory read.

        Graph construction can contain 25 nodes. It must not issue 25 FTS
        queries plus selective embedding calls on the request path.
        """

        with observe_operation(
            "rag.ground_graph_nodes",
            subsystem="retrieval",
            consumer="project_graph",
            attributes={"node_request_count": len(requests)},
        ) as op:
            units = self.store.get_units(project_id)
            memories = ProjectMemoryStore(self.index.db).list_project(project_id)
            evidence_result: dict[str, list[str]] = {}
            memory_result: dict[str, list[str]] = {}
            for node_id, query, document_ids in requests:
                allowed = set(document_ids)
                ranked_units = sorted(
                    (
                        (_term_overlap(query, unit.index_text), unit.ordinal, unit.evidence_id)
                        for unit in units
                        if unit.index_text.strip() and (not allowed or unit.document_id in allowed)
                    ),
                    key=lambda item: (-item[0], item[1]),
                )
                evidence_result[node_id] = [
                    evidence_id for score, _, evidence_id in ranked_units if score > 0
                ][:4]
                ranked_memories = sorted(
                    ((_term_overlap(query, item.statement), item.memory_id) for item in memories),
                    key=lambda item: (-item[0], item[1]),
                )
                memory_result[node_id] = [
                    memory_id for score, memory_id in ranked_memories if score > 0
                ][:2]
            op.add_count("source_evidence_count", sum(map(len, evidence_result.values())))
            op.add_count("project_memory_count", sum(map(len, memory_result.values())))
            if not evidence_result:
                op.mark_terminal("success_empty")
            return evidence_result, memory_result

    def project_outlines(self, project_id: str) -> list[dict[str, object]]:
        """Return canonical structural evidence for initial graph generation."""

        with observe_operation(
            "rag.project_outlines",
            subsystem="retrieval",
            consumer="project_graph",
        ) as op:
            outlines: list[dict[str, object]] = []
            structural_types = {
                EvidenceType.TITLE,
                EvidenceType.HEADING,
                EvidenceType.SECTION_CARD,
                EvidenceType.TABLE,
                EvidenceType.FIGURE,
                EvidenceType.PLOT,
                EvidenceType.DIAGRAM,
            }
            for document in self.store.list_documents(project_id):
                units = [
                    unit
                    for unit in self.store.get_units(
                        project_id,
                        document_ids=[document.document_id],
                    )
                    if unit.element_type in structural_types
                ][:80]
                outlines.append({
                    "filename": document.filename,
                    "document_id": document.document_id,
                    "structure": "\n\n".join(
                        f"[{unit.evidence_id} | page {unit.page_start} | {unit.element_type.value}]\n{unit.index_text}"
                        for unit in units
                    )[:16000],
                    "evidence_ids": [unit.evidence_id for unit in units],
                })
            op.add_count("documents_represented", len(outlines))
            op.add_count(
                "source_evidence_count",
                sum(len(item["evidence_ids"]) for item in outlines),
            )
            if not outlines:
                op.mark_terminal("success_empty")
            return outlines

    def project_citation_material(self, project_id: str) -> list[dict[str, object]]:
        """Return front matter and bibliography from canonical evidence only."""

        with observe_operation(
            "rag.project_citation_material",
            subsystem="retrieval",
            consumer="paper_graph",
        ) as op:
            payloads: list[dict[str, object]] = []
            for document in self.store.list_documents(project_id):
                units = self.store.get_units(
                    project_id,
                    document_ids=[document.document_id],
                )
                front_units = [
                    unit for unit in units
                    if unit.page_start == 1
                    and unit.element_type in {
                        EvidenceType.TITLE,
                        EvidenceType.HEADING,
                        EvidenceType.PARAGRAPH,
                    }
                ][:20]
                reference_units = [
                    unit for unit in units
                    if (
                        unit.element_type is EvidenceType.REFERENCE
                        or self._is_reference_section_unit(unit.section_path)
                    )
                    and unit.index_text.strip()
                ]
                reference_units.sort(key=lambda unit: unit.ordinal)
                references = [
                    unit.raw_text.strip() or unit.index_text
                    for unit in reference_units
                ]
                # ``index_text`` includes an internal retrieval envelope
                # (Paper/Section/Element labels).  Citation identity must use
                # the source text, never that implementation detail.
                front_text = "\n".join(
                    unit.raw_text.strip() or unit.index_text
                    for unit in front_units
                )
                body = "\n".join(
                    part for part in (
                        document.title,
                        document.abstract,
                        front_text,
                        "References\n" + "\n".join(references) if references else "",
                    ) if part
                )
                payloads.append({
                    "file_id": document.document_id,
                    "filename": document.filename,
                    "text": body,
                    "document_metadata": {
                        "Title": document.title,
                        "Author": "; ".join(document.authors),
                    },
                    "front_matter_text": front_text,
                    "reference_evidence_ids": [
                        unit.evidence_id for unit in reference_units
                    ],
                })
            op.add_count("documents_represented", len(payloads))
            op.add_count(
                "source_evidence_count",
                sum(len(item["reference_evidence_ids"]) for item in payloads),
            )
            if not payloads:
                op.mark_terminal("success_empty")
            return payloads

    @staticmethod
    def _is_reference_section_unit(section_path: Sequence[str]) -> bool:
        """Support evidence written before bibliography headings were normalized."""

        if not section_path:
            return False
        heading = re.sub(r"[*_`]+", "", section_path[-1] or "").strip()
        return bool(_REFERENCE_SECTION_RE.match(heading))

    def _resolve_anchors(self, request: EvidenceRequest) -> list[str]:
        ids = list(request.anchor_evidence_ids)
        for anchor in request.selection_anchors:
            if anchor.evidence_id:
                ids.append(anchor.evidence_id)
            if anchor.region_id:
                ids.extend(
                    self.store.resolve_region(
                        request.project_id,
                        anchor.document_id,
                        anchor.region_id,
                    )
                )
            ids.extend(
                self.store.resolve_selection(
                    request.project_id,
                    anchor.document_id,
                    anchor.page_number,
                    anchor.boxes,
                    anchor.text,
                )
            )
        existing = self.store.get_units(request.project_id, evidence_ids=list(dict.fromkeys(ids)))
        allowed_documents = set(request.document_ids)
        return [
            unit.evidence_id
            for unit in existing
            if not allowed_documents or unit.document_id in allowed_documents
        ]

    def _expand_structure(self, project_id: str, evidence_ids: Sequence[str], limit: int) -> list[str]:
        relations = self.store.get_relations(evidence_ids)
        expanded: list[str] = []
        for relation in relations:
            if str(relation.relation_type) not in _STRUCTURAL_RELATIONS:
                continue
            if relation.source_evidence_id in evidence_ids:
                expanded.append(relation.target_evidence_id)
            if relation.target_evidence_id in evidence_ids:
                expanded.append(relation.source_evidence_id)
        return list(dict.fromkeys(expanded))[:limit]

    def _rank(
        self,
        request: EvidenceRequest,
        units: dict,
        fused_scores: dict[str, float],
        reasons: dict[str, list[str]],
        anchor_ids: Sequence[str],
        lexical: Sequence[LexicalCandidate],
        dense: Sequence[VectorCandidate],
    ) -> list[RetrievedEvidence]:
        lexical_scores = {candidate.evidence_id: candidate.score for candidate in lexical}
        dense_scores = {candidate.item_id: candidate.score for candidate in dense}
        anchors = set(anchor_ids)
        ranked: list[RetrievedEvidence] = []
        for evidence_id, unit in units.items():
            structural = 0.0
            if evidence_id in anchors:
                structural += 1.0
            if unit.element_type in {EvidenceType.SECTION_CARD, EvidenceType.LOCAL_WINDOW}:
                structural += 0.05 if request.consumer in _SYNTHESIS_CONSUMERS else -0.01
            if unit.element_type in _VISUAL_TYPES and request.consumer in {
                ConsumerType.VISUALIZATION,
                ConsumerType.CHAT,
                ConsumerType.WIKI,
            }:
                structural += 0.04
            if unit.quality_flags:
                structural -= 0.01 * min(3, len(unit.quality_flags))
            ranked.append(
                RetrievedEvidence(
                    evidence=unit,
                    fused_score=fused_scores.get(evidence_id, 0.0) + structural,
                    dense_score=dense_scores.get(evidence_id),
                    lexical_score=lexical_scores.get(evidence_id),
                    structural_score=structural,
                    retrieval_reasons=list(dict.fromkeys(reasons.get(evidence_id, []))),
                )
            )
        ranked.sort(key=lambda item: (-item.fused_score, item.evidence.ordinal))
        return ranked

    def _citations(
        self,
        project_id: str,
        evidence: Sequence[RetrievedEvidence],
    ) -> list[EvidenceCitation]:
        documents = {doc.document_id: doc for doc in self.store.list_documents(project_id)}
        return [
            EvidenceCitation(
                evidence_id=item.evidence.evidence_id,
                document_id=item.evidence.document_id,
                filename=documents.get(item.evidence.document_id).filename
                if item.evidence.document_id in documents
                else "",
                page_start=item.evidence.page_start,
                page_end=item.evidence.page_end,
                bbox_norm=item.evidence.bbox_norm,
                section_path=item.evidence.section_path,
            )
            for item in evidence
        ]


def _fuse(
    anchor_ids: Sequence[str],
    lexical: Sequence[LexicalCandidate],
    dense: Sequence[VectorCandidate],
) -> tuple[list[str], dict[str, float], dict[str, list[str]]]:
    scores: dict[str, float] = defaultdict(float)
    reasons: dict[str, list[str]] = defaultdict(list)
    order: list[str] = []
    for anchor_rank, evidence_id in enumerate(anchor_ids):
        scores[evidence_id] += max(1.25, 2.0 - anchor_rank * 0.08)
        reasons[evidence_id].append("explicit_anchor")
        order.append(evidence_id)
    for rank, candidate in enumerate(lexical, 1):
        scores[candidate.evidence_id] += 1.0 / (_RRF_K + rank)
        reasons[candidate.evidence_id].append("lexical")
        order.append(candidate.evidence_id)
    for rank, candidate in enumerate(dense, 1):
        scores[candidate.item_id] += 1.0 / (_RRF_K + rank)
        reasons[candidate.item_id].append("dense")
        order.append(candidate.item_id)
    return list(dict.fromkeys(order)), scores, reasons


def _term_overlap(query: str, text: str) -> int:
    terms = {
        term
        for term in re.findall(r"[a-zA-Z][a-zA-Z0-9_-]{2,}", query.casefold())
        if term not in {"about", "from", "paper", "that", "this", "with"}
    }
    lowered = text.casefold()
    return sum(1 for term in terms if term in lowered)


def _select_diverse(
    ranked: Sequence[RetrievedEvidence],
    request: EvidenceRequest,
) -> tuple[list[RetrievedEvidence], bool]:
    char_budget = max(1024, request.token_budget * 4)
    per_document = 8 if request.consumer in _SYNTHESIS_CONSUMERS else 5
    per_section = 3
    selected: list[RetrievedEvidence] = []
    document_counts: dict[str, int] = defaultdict(int)
    section_counts: dict[tuple[str, tuple[str, ...]], int] = defaultdict(int)
    used_chars = 0
    truncated = False
    for item in ranked:
        unit = item.evidence
        section_key = (unit.document_id, tuple(unit.section_path))
        anchored = "explicit_anchor" in item.retrieval_reasons
        if not anchored and (
            document_counts[unit.document_id] >= per_document or section_counts[section_key] >= per_section
        ):
            continue
        size = max(1, len(unit.index_text))
        if selected and used_chars + size > char_budget:
            truncated = True
            continue
        selected.append(item)
        document_counts[unit.document_id] += 1
        section_counts[section_key] += 1
        used_chars += size
    return selected, truncated