"""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