"""Lab 01 — the knowledge foundation: authorized hybrid retrieval. Two things, and they are not the same thing: * retrieval QUALITY is engineering — chunking, BM25, dense recall, fusion, reranking; * retrieval SAFETY is architecture — namespaces, entitlement as a PRE-filter, citations. The failure this lab is built around has no runtime detection: a shared index with a post-hoc entitlement filter returns tenant B's chunk to tenant A, and the response is a 200 OK with a plausible answer. pytest test_lab.py -v LAB_MODULE=solution pytest test_lab.py -v # the reference, must be green """ from __future__ import annotations import hashlib import math import re from collections import Counter from dataclasses import dataclass, field, replace from enum import Enum from typing import Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple # ====================================================================================== # 1. Documents, chunks and provenance # ====================================================================================== @dataclass(frozen=True) class Document: """``tenant``, ``classification`` and ``barrier`` are not metadata to filter on later — they select the NAMESPACE a chunk is written to.""" doc_id: str tenant: str title: str text: str classification: str = "internal" barrier: Optional[str] = None # information barrier / desk, e.g. "advisory" version: int = 1 updated_tick: int = 0 def __post_init__(self) -> None: # TODO: doc_id and tenant required (ValueError); validate classification via # classification_rank so a typo fails at INGESTION, not at query time. raise NotImplementedError _CLASSIFICATION_ORDER = ("public", "internal", "confidential", "restricted") def classification_rank(name: str) -> int: """Index in _CLASSIFICATION_ORDER. Unknown -> ValueError.""" # TODO raise NotImplementedError @dataclass(frozen=True) class Chunk: chunk_id: str doc_id: str tenant: str text: str section: str start_char: int end_char: int classification: str barrier: Optional[str] doc_version: int updated_tick: int def citation(self) -> str: """``doc_id@vN#section[start:end]`` — enough to point a reviewer at the span.""" # TODO raise NotImplementedError # ====================================================================================== # 2. Chunking # ====================================================================================== _SECTION_RE = re.compile(r"^##\s+(.+)$", re.MULTILINE) def estimate_tokens(text: str) -> int: """ceil(chars / 4); empty -> 0.""" # TODO raise NotImplementedError def chunk_document(document: Document, *, max_tokens: int = 96, overlap_tokens: int = 16) -> List[Chunk]: """Structure-aware chunking: split on ``## headings`` FIRST, then on size. Splitting on structure before size is the highest-leverage quality decision in retrieval. A fixed-size splitter cuts through the middle of a policy clause and produces a chunk that is retrievable but not usable. Validate: ``max_tokens > 0`` and ``0 <= overlap_tokens < max_tokens`` (equal overlap means chunks stop advancing). Chunk ids are ``{doc_id}::{n}``, 1-based, in order. Empty text -> []. """ # TODO: use _split_sections then _split_by_size raise NotImplementedError def _split_sections(text: str) -> List[Tuple[str, str, int]]: """Return ``(section_name, section_text, char_offset)``. No headings at all -> a single ``("body", stripped_text, 0)``. Text before the first heading is its own ``"body"`` section. """ # TODO raise NotImplementedError def _split_by_size(text: str, max_tokens: int, overlap_tokens: int) -> List[Tuple[str, int, int]]: """Split on WORD boundaries into ``(piece, start_char, end_char)``. Never split mid-word. After emitting a piece, step back by roughly ``overlap_tokens`` worth of words so a fact spanning a boundary is still retrievable. """ # TODO raise NotImplementedError def _char_offset(text: str, words: Sequence[str], word_index: int) -> int: """Character offset of ``words[word_index]`` within ``text``, scanning forward so a repeated word resolves to the right occurrence.""" # TODO raise NotImplementedError # ====================================================================================== # 3. Embeddings # ====================================================================================== def hash_embed(text: str, dimensions: int = 96) -> Tuple[float, ...]: """Signed feature hashing, L2-normalized. Per token: blake2b digest; ``index = first 4 bytes mod dimensions``; ``sign = +1 if digest[4] is even else -1``; accumulate. Normalize at the end; an all-zero vector stays all-zero. The SIGN is not decoration: without it, hash collisions always add constructively and every pair of documents looks more similar than it is. """ # TODO raise NotImplementedError def cosine(a: Sequence[float], b: Sequence[float]) -> float: """Dot product — both inputs are normalized.""" # TODO raise NotImplementedError _TOKEN_RE = re.compile(r"[a-z0-9][a-z0-9\-]*") def _tokenize(text: str) -> List[str]: return _TOKEN_RE.findall(text.lower()) # ====================================================================================== # 4. The namespace # ====================================================================================== @dataclass(frozen=True) class Principal: user_id: str tenant: str max_classification: str = "internal" barriers: Tuple[str, ...] = () # barriers this principal is INSIDE def namespace_key(self) -> str: # TODO raise NotImplementedError def namespace_of(chunk: Chunk) -> str: """The partition a chunk is written to. Tenant, and nothing else, decides it. Classification and barrier are enforced WITHIN a namespace; the tenant boundary is enforced BY the namespace. Two mechanisms on purpose: a classification bug leaks within a tenant; a tenant-filter bug leaks across customers. """ # TODO raise NotImplementedError # ====================================================================================== # 5. BM25 # ====================================================================================== @dataclass(frozen=True) class Scored: chunk: Chunk score: float class BM25Index: """Okapi BM25. ``k1`` controls term-frequency saturation; ``b`` controls length normalization (b=1 fully penalizes length, b=0 not at all). """ def __init__(self, *, k1: float = 1.5, b: float = 0.75) -> None: if k1 < 0 or not 0.0 <= b <= 1.0: raise ValueError("k1 must be >= 0 and b in [0, 1]") self.k1 = k1 self.b = b self._by_namespace: Dict[str, List[Chunk]] = {} self._tf: Dict[str, List[Counter]] = {} self._df: Dict[str, Counter] = {} self._lengths: Dict[str, List[int]] = {} def add(self, chunk: Chunk) -> None: """Index into the chunk's namespace: append the chunk, its term Counter, its length, and increment document frequency for each DISTINCT term.""" # TODO raise NotImplementedError def _idf(self, ns: str, term: str) -> float: """``max(0, log((N - df + 0.5) / (df + 0.5) + 1))``. The +1 inside the log and the max(0, ...) both matter: without them a term appearing in most documents contributes a NEGATIVE score, which is nonsense. """ # TODO raise NotImplementedError def search(self, query: str, *, namespace: str, limit: int = 10) -> List[Scored]: """Score every chunk in the namespace: sum over query terms of idf * f * (k1 + 1) / (f + k1 * (1 - b + b * len/avgdl)) Skip terms with f == 0. Drop zero scores. Sort by ``(-score, chunk_id)`` so ties are deterministic. Unknown namespace -> []. """ # TODO raise NotImplementedError # ====================================================================================== # 6. Dense retrieval # ====================================================================================== class VectorIndex: """A namespaced dense index. The silo/pool/bridge topology decision is made HERE, by keying the store on namespace.""" def __init__(self, *, dimensions: int = 96) -> None: self.dimensions = dimensions self._by_namespace: Dict[str, List[Tuple[Chunk, Tuple[float, ...]]]] = {} def add(self, chunk: Chunk) -> None: # TODO: store (chunk, hash_embed(chunk.text, self.dimensions)) in its namespace raise NotImplementedError def search(self, query: str, *, namespace: str, limit: int = 10) -> List[Scored]: """Cosine against every entry in the namespace; drop non-positive scores; sort by ``(-score, chunk_id)``. Unknown namespace -> [].""" # TODO raise NotImplementedError def namespaces(self) -> List[str]: """Sorted, so the test can assert on it.""" # TODO raise NotImplementedError # ====================================================================================== # 7. Fusion and reranking # ====================================================================================== def reciprocal_rank_fusion(rankings: Sequence[Sequence[Scored]], *, k: int = 60, limit: int = 10) -> List[Scored]: """Fused score = sum of ``1 / (k + rank)`` over the rankings a chunk appears in, rank being 1-based. Score-FREE, which is the point: BM25 scores and cosine similarities are incomparable, and normalizing them is a calibration that breaks when you change embedding model. Consequence: consistent agreement beats one strong signal. 1st + 10th = 1/61 + 1/70 = 0.03068; 3rd + 3rd = 2/63 = 0.03175. ``k <= 0`` -> ValueError. Sort by ``(-score, chunk_id)``. """ # TODO raise NotImplementedError Reranker = Callable[[str, Chunk], float] def lexical_overlap_reranker(query: str, chunk: Chunk) -> float: """Deterministic stand-in for a cross-encoder. ``coverage + 0.25 * density`` where coverage is the fraction of DISTINCT query terms present in the chunk and density is (total occurrences of query terms) / (chunk length). Empty query -> 0.0. Faithful in the property that matters: it scores the PAIR. """ # TODO raise NotImplementedError def rerank(query: str, candidates: Sequence[Scored], *, reranker: Reranker, limit: int = 5, min_score: float = 0.0) -> List[Scored]: """Second stage over the top-k ONLY, with a relevance floor. A cross-encoder cannot precompute anything — it runs per (query, document) pair — so it never runs over the index, and it is the first thing shed under latency pressure. ``min_score`` matters more than it looks: first-stage retrieval ALWAYS returns something, so without a floor "nothing is relevant" is indistinguishable from "here are the least-irrelevant chunks". Keep only ``score > min_score``. """ # TODO raise NotImplementedError # ====================================================================================== # 8. The authorized retriever # ====================================================================================== @dataclass(frozen=True) class RetrievalPolicy: max_stale_ticks: Optional[int] = None # freshness contract; None = no contract enable_rerank: bool = True candidate_limit: int = 12 result_limit: int = 5 @dataclass(frozen=True) class RetrievalResult: chunks: Tuple[Scored, ...] degraded: Tuple[str, ...] = () considered: int = 0 excluded_by_entitlement: int = 0 excluded_by_freshness: int = 0 @property def citations(self) -> Tuple[str, ...]: return tuple(s.chunk.citation() for s in self.chunks) class AuthorizedRetriever: """Hybrid retrieval that is authorized by construction — three independent points: 1. the NAMESPACE bounds the search to the caller's tenant (never retrieved, not filtered out); 2. classification and barrier are a PRE-filter inside the namespace, before ranking; 3. the freshness contract drops chunks too stale to rely on. """ def __init__(self, *, bm25: BM25Index, vectors: VectorIndex, reranker: Reranker = lexical_overlap_reranker) -> None: self.bm25 = bm25 self.vectors = vectors self.reranker = reranker self._chunks: Dict[str, Chunk] = {} def ingest(self, document: Document, *, max_tokens: int = 96, overlap_tokens: int = 16) -> List[Chunk]: """Chunk, then add every chunk to BOTH indexes. Return the chunks.""" # TODO raise NotImplementedError def visible(self, principal: Principal, chunk: Chunk) -> bool: """A conjunction: same tenant AND classification within clearance AND (no barrier OR the principal is inside it).""" # TODO raise NotImplementedError def retrieve(self, query: str, principal: Principal, *, now_tick: int = 0, policy: RetrievalPolicy = RetrievalPolicy()) -> RetrievalResult: """The pipeline: 1. search BOTH indexes in ``principal.namespace_key()`` only; 2. count ``considered`` as the union of candidate chunk ids; 3. PRE-FILTER each ranking by ``visible()`` and by the freshness contract, counting exclusions — this happens BEFORE fusion, so an ineligible chunk can neither occupy a slot nor leak through result-set size; 4. fuse the non-empty rankings with RRF; 5. rerank unless disabled (record ``"rerank"`` in ``degraded`` when skipped); 6. return at most ``result_limit``. No rankings survive the filter -> an empty result, not an error. """ # TODO raise NotImplementedError # ====================================================================================== # 9. Grounding # ====================================================================================== @dataclass(frozen=True) class Claim: text: str citation: Optional[str] = None @dataclass(frozen=True) class GroundingReport: supported: Tuple[Claim, ...] unsupported: Tuple[Claim, ...] @property def is_grounded(self) -> bool: return not self.unsupported @property def coverage(self) -> float: total = len(self.supported) + len(self.unsupported) return len(self.supported) / total if total else 1.0 _STOPWORDS = { "the", "a", "an", "is", "are", "was", "were", "be", "been", "of", "to", "in", "on", "for", "and", "or", "it", "this", "that", "with", "as", "at", "by", "from", "has", "have", "had", "will", "would", "can", "could", "should", "may", "might", } def check_grounding(claims: Sequence[Claim], retrieved: Sequence[Scored], *, min_overlap: float = 0.6) -> GroundingReport: """Every claim must map to a retrieved span, or the answer is not defensible. For each claim: take its content words (tokens minus _STOPWORDS). A claim with no content words is trivially supported. Otherwise find the retrieved chunk with the highest overlap fraction; if it meets ``min_overlap``, the claim is supported and carries that chunk's citation. Otherwise it is unsupported. The architecture is the lesson: the check runs over the RETRIEVED SET, so an invented claim has nothing to match. ``min_overlap`` outside (0, 1] -> ValueError. """ # TODO raise NotImplementedError # ====================================================================================== # 10. Context assembly # ====================================================================================== class Segment(str, Enum): """Prompt segments, in emission order. Stable first, volatile last — the first changed byte invalidates provider-side prefix caching (Phase 04).""" INSTRUCTIONS = "instructions" TOOL_SCHEMAS = "tool_schemas" POLICY = "policy" MEMORY = "memory" RETRIEVED = "retrieved" QUESTION = "question" _SEGMENT_ORDER = tuple(Segment) @dataclass(frozen=True) class AssembledContext: text: str tokens: int included_chunks: Tuple[str, ...] dropped_chunks: Tuple[str, ...] stable_prefix_tokens: int @property def cacheable_fraction(self) -> float: return self.stable_prefix_tokens / self.tokens if self.tokens else 0.0 def assemble_context(*, budget_tokens: int, instructions: str, tool_schemas: str = "", policy: str = "", memory: str = "", question: str, retrieved: Sequence[Scored] = ()) -> AssembledContext: """Fill a token budget in cache-friendly order. Render each non-empty segment as ``[name]\\n`` and join blocks with a blank line. Retrieved chunks render as ``<>\\n``, joined the same way. ``stable_prefix_tokens`` counts the rendered instructions, tool schemas and policy blocks only — memory varies per session. Measure the ASSEMBLED text, not the sum of its pieces: segment headers and the separators between blocks are real tokens, and a budget that ignores them over-fills by exactly the amount nobody accounts for. Drop the LOWEST-ranked retrieved chunk and re-measure until it fits. Two rules that are easy to get backwards: 1. the QUESTION is never dropped — if the fixed segments alone exceed the budget, raise; 2. dropping is REPORTED, because "we answered without the third source" is a fact the grounding check and the audit record both need. ``budget_tokens <= 0`` -> ValueError. """ # TODO raise NotImplementedError def _render_chunk(chunk: Chunk) -> str: """``<>\\n``.""" # TODO raise NotImplementedError def main() -> None: print("implement the TODOs, then compare with `python solution.py`") if __name__ == "__main__": main()