Spaces:
Sleeping
Sleeping
| """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"] | |
| class CorpusDocument: | |
| document_id: str | |
| title: str | |
| topic: str | |
| text: str | |
| sha256: str | |
| 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 | |
| class QueryLabel: | |
| query_id: str | |
| category: str | |
| question: str | |
| expected_documents: list[str] | |
| expected_sections: list[str] | |
| answerability: Answerability | |
| citation_required: bool | |
| 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 | |