app/benchmarks/rag_observability/corpus.py

"""Fixed corpus + labelled query-set loading for the RAG observability benchmark.

Owns exactly one responsibility: interpreting `backend/benchmarks/rag/corpus.v1.json`
and `query_set.v1.json` (plus the corpus text files they reference) into typed
objects the rest of `rag_observability` consumes -- `runner.py` for ingestion,
`quality.py` for scoring. It emits nothing and calls no observability code.

Corpus documents are synthetic short "papers" (Abstract/1. Introduction/
2. Method/3. Results/4. Discussion) authored specifically for this benchmark so
every query's expected answer is verifiable by construction. Because the real
ingestion path (`app.rag.ingestion.ingest_text`) chunks each document with a
plain, section-agnostic splitter (`app.rag.chunker.chunk_text`, 512 chars /
64 overlap), a retrieved candidate never carries a "section" field on its own
-- `section_for_chunk` recovers "which section does this chunk fall under"
after the fact, by replaying the identical chunker against the source text and
locating each chunk's start offset relative to the section headers. This
mirrors real chunk boundaries exactly (same chunker, same parameters) rather
than inventing separate per-section ingestion calls that a real single-file
upload would never produce.
"""

from __future__ import annotations

import hashlib
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Literal

from app.rag.chunker import chunk_text

_RAG_BENCHMARKS_ROOT = Path(__file__).resolve().parents[3] / "benchmarks" / "rag"
_DEFAULT_CORPUS_MANIFEST = _RAG_BENCHMARKS_ROOT / "corpus.v1.json"
_DEFAULT_QUERY_SET = _RAG_BENCHMARKS_ROOT / "query_set.v1.json"

Answerability = Literal["answerable", "unanswerable"]


@dataclass(frozen=True)
class CorpusDocument:
    document_id: str
    title: str
    topic: str
    text: str
    sha256: str


@dataclass(frozen=True)
class Corpus:
    version: str
    section_headers: list[str]
    documents: list[CorpusDocument]

    def by_id(self, document_id: str) -> CorpusDocument | None:
        for doc in self.documents:
            if doc.document_id == document_id:
                return doc
        return None


@dataclass(frozen=True)
class QueryLabel:
    query_id: str
    category: str
    question: str
    expected_documents: list[str]
    expected_sections: list[str]
    answerability: Answerability
    citation_required: bool


@dataclass(frozen=True)
class QuerySet:
    version: str
    queries: list[QueryLabel]


class CorpusIntegrityError(ValueError):
    """A corpus text file's bytes no longer match the manifest's recorded hash."""


def load_corpus(manifest_path: Path = _DEFAULT_CORPUS_MANIFEST) -> Corpus:
    """Load the fixed corpus, verifying every file's bytes against its manifest hash.

    Raises :class:`CorpusIntegrityError` if a corpus text file was edited without
    updating `corpus.v1.json` -- the whole point of hashing source bytes is to
    catch that accidental drift before a benchmark run silently measures against
    a corpus that no longer matches its query labels.
    """
    manifest_path = Path(manifest_path)
    raw = json.loads(manifest_path.read_text(encoding="utf-8"))
    root = manifest_path.parent

    documents: list[CorpusDocument] = []
    for entry in raw["documents"]:
        file_path = root / entry["file"]
        content = file_path.read_bytes()
        actual_hash = hashlib.sha256(content).hexdigest()
        expected_hash = entry["sha256"]
        if actual_hash != expected_hash:
            raise CorpusIntegrityError(
                f"corpus document {entry['document_id']!r} at {file_path} has drifted "
                f"from corpus.v1.json: expected sha256={expected_hash}, got {actual_hash}. "
                "Re-hash and update the manifest if this edit was intentional."
            )
        documents.append(
            CorpusDocument(
                document_id=entry["document_id"],
                title=entry["title"],
                topic=entry["topic"],
                text=content.decode("utf-8"),
                sha256=actual_hash,
            )
        )

    return Corpus(
        version=raw["corpus_version"],
        section_headers=list(raw["section_headers"]),
        documents=documents,
    )


def load_query_set(path: Path = _DEFAULT_QUERY_SET) -> QuerySet:
    raw = json.loads(Path(path).read_text(encoding="utf-8"))
    queries = [
        QueryLabel(
            query_id=q["query_id"],
            category=q["category"],
            question=q["question"],
            expected_documents=list(q["expected_documents"]),
            expected_sections=list(q["expected_sections"]),
            answerability=q["answerability"],
            citation_required=bool(q["citation_required"]),
        )
        for q in raw["queries"]
    ]
    return QuerySet(version=raw["query_set_version"], queries=queries)


def corpus_manifest_fingerprint(manifest_path: Path = _DEFAULT_CORPUS_MANIFEST) -> str:
    """Sha256 of the manifest file itself (distinct from any one document's hash) --
    changes whenever a document is added/removed/retitled even if no existing
    document's bytes changed."""
    return hashlib.sha256(Path(manifest_path).read_bytes()).hexdigest()


def query_set_fingerprint(path: Path = _DEFAULT_QUERY_SET) -> str:
    return hashlib.sha256(Path(path).read_bytes()).hexdigest()


# --- Section recovery (chunk_index -> section header) -----------------------

_CHUNK_OVERLAP = 64  # must match app.rag.chunker.chunk_text's default


def _section_offsets(text: str, headers: list[str]) -> list[tuple[int, str]]:
    offsets = [(text.find(h), h) for h in headers]
    return sorted((off, h) for off, h in offsets if off != -1)


def _chunk_start_offsets(text: str, chunks: list[str]) -> list[int]:
    """Locate each chunk's start offset in `text`, in order.

    `chunk_text` (RecursiveCharacterTextSplitter) does not return offsets, so
    this replays a forward-only search: each chunk is searched for starting at
    (or after) the previous chunk's own start offset, advancing the cursor by
    at least `len(chunk) - overlap` so a short, possibly-repeated chunk prefix
    can't match an earlier occurrence.
    """
    cursor = 0
    offsets: list[int] = []
    for chunk in chunks:
        start = text.find(chunk, cursor)
        if start == -1:
            start = text.find(chunk)  # fallback: search from the beginning
        offsets.append(start)
        if start != -1:
            cursor = start + max(1, len(chunk) - _CHUNK_OVERLAP)
    return offsets


def build_section_map(document: CorpusDocument, headers: list[str]) -> list[str | None]:
    """Return, for each chunk `chunk_text(document.text)` produces (in order),
    the section header whose text most closely precedes that chunk's start
    offset -- or `None` for a chunk that starts before any header (the title/
    preamble chunk)."""
    chunks = chunk_text(document.text)
    header_offsets = _section_offsets(document.text, headers)
    chunk_offsets = _chunk_start_offsets(document.text, chunks)

    sections: list[str | None] = []
    for start in chunk_offsets:
        if start == -1:
            sections.append(None)
            continue
        section: str | None = None
        for off, header in header_offsets:
            if off <= start:
                section = header
            else:
                break
        sections.append(section)
    return sections


class SectionIndex:
    """Cached `document_id -> [section per chunk_index]` lookup for a `Corpus`."""

    def __init__(self, corpus: Corpus) -> None:
        self._corpus = corpus
        self._cache: dict[str, list[str | None]] = {}

    def section_for_chunk(self, document_id: str, chunk_index: int | None) -> str | None:
        if chunk_index is None:
            return None
        if document_id not in self._cache:
            doc = self._corpus.by_id(document_id)
            if doc is None:
                self._cache[document_id] = []
            else:
                self._cache[document_id] = build_section_map(doc, self._corpus.section_headers)
        sections = self._cache[document_id]
        if 0 <= chunk_index < len(sections):
            return sections[chunk_index]
        return None