Chunking Strategies

# script_id: day_025_chunking/chunker
"""atlas/ingest.py, Day 25 additions: split Documents into DocChunks.

This lesson saves it as chunker.py; compare.py and the exercises import it.
Install: uv add tiktoken pydantic
Env:     ORBIT_DATA (optional, defaults to datasets/orbit). Offline, $0.
"""
import os
import re
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path
from typing import Any, Callable, Literal, Sequence

import tiktoken
from pydantic import BaseModel, Field

DATA_DIR = Path(os.environ.get("ORBIT_DATA", "datasets/orbit"))

@lru_cache(maxsize=1)
def _enc() -> tiktoken.Encoding:
    return tiktoken.get_encoding("o200k_base")   # Day 4's stand-in; counts are approximate

def count_tokens(text: str) -> int:              # built on Day 4
    return len(_enc().encode(text, disallowed_special=()))

@dataclass
class Document:                                  # built on Day 24 (compact copy)
    id: str
    title: str
    text: str                                    # normalized Markdown body
    source: str
    metadata: dict[str, str] = field(default_factory=dict)

def load_corpus(root: Path = DATA_DIR) -> list[Document]:
    """Compact Day 24 stand-in: Markdown pages with flat `key: value` front matter."""
    docs = []
    for path in sorted((root / "docs").glob("*.md")):
        _, head, body = path.read_text(encoding="utf-8").split("---", 2)
        meta = {}
        for line in head.strip().splitlines():
            key, _, value = line.partition(":")
            meta[key.strip()] = value.split(" #")[0].strip().strip('"')   # drop inline comments
        docs.append(Document(meta["id"], meta["title"], body.lstrip("\n"), str(path), meta))
    return docs

# script_id: day_025_chunking/chunker
class DocChunk(BaseModel):          # atlas/schemas.py: Day 15 fields, then Day 25 additions
    id: str
    doc_id: str
    section: str
    applies_to: str
    text: str                       # exact slice of Document.text: what you show and cite
    position: int = 0
    n_tokens: int = 0
    header: str = ""                # embedded together with the text, never shown as the text
    metadata: dict[str, Any] = Field(default_factory=dict)

Span = tuple[int, int]              # [start, end) character offsets into Document.text

def _check(max_tokens: int, overlap: int) -> None:
    if not 0 <= overlap < max_tokens:   # step = max_tokens - overlap must be >= 1 to terminate
        raise ValueError(f"need 0 <= overlap < max_tokens, got overlap={overlap}, max_tokens={max_tokens}")

def token_spans(text: str, max_tokens: int, overlap: int = 0, base: int = 0) -> list[Span]:
    """Fixed windows of max_tokens tokens, as offsets into the original string."""
    _check(max_tokens, overlap)
    tokens = _enc().encode(text, disallowed_special=())
    if not tokens:
        return []
    _, starts = _enc().decode_with_offsets(tokens)
    starts.append(len(text))
    step = max_tokens - overlap
    return [(base + starts[i], base + starts[min(i + max_tokens, len(tokens))])
            for i in range(0, max(len(tokens) - overlap, 1), step)]

def fixed_chunker(doc: Document, max_tokens: int, overlap: int) -> list[tuple[str, int, int]]:
    return [(doc.title, s, e) for s, e in token_spans(doc.text, max_tokens, overlap)]

# script_id: day_025_chunking/chunker
SEPARATORS = ("\n\n", "\n", ". ")   # coarse to fine; token windows are the last resort

def _split_keep(text: str, sep: str) -> list[str]:
    """Split on sep but keep it on the left piece, so the pieces re-join to exactly `text`."""
    parts = text.split(sep)
    return [p for p in [p + sep for p in parts[:-1]] + [parts[-1]] if p]

def recursive_spans(text: str, budget: int, base: int = 0,
                    seps: Sequence[str] = SEPARATORS) -> list[Span]:
    if count_tokens(text) <= budget:
        return [(base, base + len(text))] if text else []
    if not seps:                    # nothing left to split on: token windows, never split("")
        return token_spans(text, budget, 0, base)
    spans, pos = [], base
    for piece in _split_keep(text, seps[0]):
        spans += recursive_spans(piece, budget, pos, seps[1:])
        pos += len(piece)
    return spans

def pack(text: str, spans: list[Span], budget: int) -> list[Span]:
    """Greedily merge neighbouring spans while the merged slice still fits the budget."""
    out: list[Span] = []
    for s, e in spans:
        if out and count_tokens(text[out[-1][0]:e]) <= budget:
            out[-1] = (out[-1][0], e)
        else:
            out.append((s, e))
    return out

def add_overlap(text: str, spans: list[Span], overlap: int) -> list[Span]:
    """Start each span `overlap` tokens early, inside the previous one. Runs AFTER packing."""
    out = spans[:1]
    for (ps, pe), (s, e) in zip(spans, spans[1:]):
        tokens = _enc().encode(text[ps:pe], disallowed_special=())
        if overlap == 0 or not tokens:
            out.append((s, e))
            continue
        _, starts = _enc().decode_with_offsets(tokens)
        start = ps + starts[max(0, len(tokens) - overlap)]
        nl = text.find("\n", start, pe)         # snap to a line start: overlap never begins mid-row
        if nl != -1 and text[nl + 1:pe].strip():
            start = nl + 1
        out.append((min(s, start), e))
    return out

def recursive_chunker(doc: Document, max_tokens: int, overlap: int) -> list[tuple[str, int, int]]:
    budget = max_tokens - overlap   # leave room for the overlap added afterwards
    spans = add_overlap(doc.text, pack(doc.text, recursive_spans(doc.text, budget), budget), overlap)
    return [(doc.title, s, e) for s, e in spans]

# script_id: day_025_chunking/chunker
HEADING = re.compile(r"(#{1,6}) +(.+)")

def md_sections(text: str) -> list[tuple[str, int, int]]:
    """(heading path, start, end) per section. Lines inside ``` fences are never headings."""
    sections, path, start, pos, in_fence = [], [], 0, 0, False
    for line in text.splitlines(keepends=True):
        if line.startswith("```"):
            in_fence = not in_fence
        elif not in_fence and (m := HEADING.match(line)):
            if pos > start:
                sections.append((" › ".join(t for _, t in path), start, pos))
            level = len(m.group(1))
            path = [(lv, t) for lv, t in path if lv < level] + [(level, m.group(2).strip())]
            start = pos
        pos += len(line)
    sections.append((" › ".join(t for _, t in path), start, len(text)))
    return sections

def md_blocks(text: str, start: int, end: int) -> list[Span]:
    """Blocks end at blank lines, except inside a fence; a table has no blank lines, so it stays whole."""
    blocks, block_start, pos, in_fence = [], start, start, False
    for line in text[start:end].splitlines(keepends=True):
        pos += len(line)
        if line.startswith("```"):
            in_fence = not in_fence
        elif not in_fence and not line.strip():
            blocks.append((block_start, pos))
            block_start = pos
    return blocks + ([(block_start, pos)] if pos > block_start else [])

def markdown_chunker(doc: Document, max_tokens: int, overlap: int,
                     min_tokens: int = 40) -> list[tuple[str, int, int]]:
    budget, merged = max_tokens - overlap, []
    for path, s, e in md_sections(doc.text):     # fold a tiny parent ("## General") into its first child
        parent = merged[-1] if merged else None
        if (parent and path.startswith(parent[0] + " › ")
                and count_tokens(doc.text[parent[1]:parent[2]]) < min_tokens):
            merged[-1] = (path, parent[1], e)
        else:
            merged.append((path, s, e))
    out = []
    for path, s, e in merged:
        pieces = [sp for bs, be in md_blocks(doc.text, s, e)   # only an oversized block is cut inside
                  for sp in recursive_spans(doc.text[bs:be], budget, bs)]
        spans = add_overlap(doc.text, pack(doc.text, pieces, budget), overlap)
        out += [(path or doc.title, a, b) for a, b in spans]
    return out

| CPU-heavy jobs using GIL-holding native code | 900–1800 | 60 |
| Flaky network to the broker | 600 | 30 |

# script_id: day_025_chunking/chunker
Chunker = Callable[[Document, int, int], list[tuple[str, int, int]]]
CHUNKERS: dict[str, Chunker] = {"fixed": fixed_chunker, "recursive": recursive_chunker,
                                "markdown": markdown_chunker}   # "semantic" registers in Go Deeper

def make_header(doc: Document, section: str) -> str:
    trail = section if section.startswith(doc.title) else f"{doc.title} › {section}"
    return f"Orbit docs › {trail} (applies to {doc.metadata.get('applies_to', 'all')})"

def chunk(doc: Document, *, strategy: Literal["fixed", "recursive", "markdown", "semantic"] = "markdown",
          max_tokens: int = 400, overlap: int = 50, add_header: bool = True) -> list[DocChunk]:
    _check(max_tokens, overlap)
    if strategy not in CHUNKERS:
        raise ValueError(f"unknown strategy {strategy!r}; registered: {sorted(CHUNKERS)}")
    params = {"strategy": strategy, "max_tokens": max_tokens, "overlap": overlap}
    applies_to = doc.metadata.get("applies_to", "all")
    return [DocChunk(id=f"{doc.id}#{i}", doc_id=doc.id, section=section, applies_to=applies_to,
                     text=doc.text[s:e], position=i, n_tokens=count_tokens(doc.text[s:e]),
                     header=make_header(doc, section) if add_header else "",
                     metadata={**params, "start": s, "end": e})   # store the recipe with the chunk
            for i, (section, s, e) in enumerate(CHUNKERS[strategy](doc, max_tokens, overlap))]

def embed_text(c: DocChunk) -> str:
    """What the embedding model sees: header + text. Show and cite c.text alone."""
    return f"{c.header}\n\n{c.text.strip()}" if c.header else c.text.strip()

def ranked_doc_ids(chunks: Sequence[DocChunk]) -> list[str]:
    """Best-first chunks -> best-first doc ids, each once (hit@k is scored per page)."""
    return list(dict.fromkeys(c.doc_id for c in chunks))

def rejoin(chunks: Sequence[DocChunk]) -> str:
    """Undo the overlap: each chunk contributes only the text past what came before."""
    out, covered = [], 0
    for c in chunks:
        out.append(c.text[max(0, covered - c.metadata["start"]):])
        covered = c.metadata["end"]
    return "".join(out)

# script_id: day_025_chunking/embedder
"""Interchangeable embedders with one shape: texts -> (n, d) float32 array with unit-length rows.

Install: uv add "openai>=3.46,<4" numpy tiktoken
Env:     OPENAI_API_KEY (openai_embed only), ATLAS_EMBED_MODEL (optional)
"""
import hashlib
import os
import re
import sqlite3
import zlib
from collections.abc import Callable, Sequence
from pathlib import Path

import numpy as np
import tiktoken
from openai import OpenAI

EMBED_MODEL = os.environ.get("ATLAS_EMBED_MODEL", "text-embedding-3-small")   # atlas/config.py (Day 23)
PRICE_PER_1M = 0.02    # USD, text-embedding-3-small; as of 2026-11, verify at https://developers.openai.com/api/docs/pricing
EmbedFn = Callable[[Sequence[str]], np.ndarray]
BILLED = {"tokens": 0}  # filled from resp.usage, the billing truth

def _normalize(m: np.ndarray) -> np.ndarray:
    norms = np.linalg.norm(m, axis=1, keepdims=True)
    return m / np.where(norms == 0, 1, norms)       # a zero row stays zero instead of nan

def embed_cost(texts: Sequence[str]) -> float:
    """Estimate before you pay: text-embedding-3-* tokenize with cl100k_base."""
    enc = tiktoken.get_encoding("cl100k_base")
    return sum(len(enc.encode(t, disallowed_special=())) for t in texts) * PRICE_PER_1M / 1e6

def local_embed(texts: Sequence[str], dim: int = 4096) -> np.ndarray:
    """Offline stand-in: hashed bag of words. Lexical, NOT semantic. For wiring, tests, $0 runs."""
    m = np.zeros((len(texts), dim), dtype=np.float32)
    for row, text in enumerate(texts):
        for word in re.findall(r"\w[\w.-]*\w|\w", text.lower()):
            m[row, zlib.crc32(word.encode()) % dim] += 1.0   # crc32 is stable across runs; hash() isn't
    return _normalize(np.log1p(m))

# script_id: day_025_chunking/embedder
class EmbeddingCache:   # built on Day 23 (compact): key = sha256(model | dims | text)
    def __init__(self, path: Path = Path(".cache/embeddings.sqlite")):
        path.parent.mkdir(parents=True, exist_ok=True)
        self.db = sqlite3.connect(path)
        self.db.execute("CREATE TABLE IF NOT EXISTS emb (key TEXT PRIMARY KEY, vec BLOB)")

    def key(self, model: str, text: str, dims: int | None = None) -> str:
        return hashlib.sha256(f"{model}|{dims}|{text}".encode()).hexdigest()

    def get(self, key: str) -> np.ndarray | None:
        row = self.db.execute("SELECT vec FROM emb WHERE key = ?", (key,)).fetchone()
        return np.frombuffer(row[0], dtype=np.float32) if row else None

    def put(self, key: str, vec: np.ndarray) -> None:
        self.db.execute("INSERT OR REPLACE INTO emb VALUES (?, ?)", (key, vec.astype(np.float32).tobytes()))
        self.db.commit()

def openai_embed(texts: Sequence[str], model: str = EMBED_MODEL, batch_size: int = 64,
                 cache: EmbeddingCache | None = None) -> np.ndarray:
    """Batched and cached: only texts missing from the cache are sent and billed."""
    cache = cache or EmbeddingCache()
    keys = [cache.key(model, t) for t in texts]
    vecs = {k: v for k in set(keys) if (v := cache.get(k)) is not None}
    todo = [i for i, k in enumerate(keys) if k not in vecs]
    client = OpenAI() if todo else None             # a warm cache needs no key at all
    for b in range(0, len(todo), batch_size):
        idx = todo[b:b + batch_size]
        resp = client.embeddings.create(model=model, input=[texts[i] for i in idx])
        BILLED["tokens"] += resp.usage.prompt_tokens
        for i, d in zip(idx, resp.data):
            vecs[keys[i]] = np.asarray(d.embedding, dtype=np.float32)
            cache.put(keys[i], vecs[keys[i]])
    return _normalize(np.stack([vecs[k] for k in keys]))

def default_embedder() -> tuple[str, EmbedFn]:
    """Real embeddings when a key is set, the offline stand-in otherwise; callers never branch."""
    if os.environ.get("OPENAI_API_KEY"):
        return EMBED_MODEL, openai_embed
    return "local_embed (offline stand-in)", local_embed

# script_id: day_025_chunking/compare
"""Which chunker wins? Doc-level hit@1/@5 on the dev split, chunk counts and embedding cost.

Imports chunker.py and embedder.py from this lesson (same folder).
Env: OPENAI_API_KEY for real embeddings (≈ $0.01, as of 2026-11, verify); without it, $0 offline.
"""
import json
from collections import defaultdict
from dataclasses import dataclass

import numpy as np

from chunker import DATA_DIR, Document, chunk, embed_text, load_corpus, ranked_doc_ids
from embedder import EmbedFn, default_embedder, embed_cost

@dataclass
class RetrievalCase:                 # built on Day 23 (compact)
    id: str
    query: str
    relevant_doc_ids: list[str]
    kind: str                        # exact_term | semantic | multi_hop

def load_retrieval_cases(split: str = "dev") -> list[RetrievalCase]:
    """Day 23's deterministic split: every other case by id within each kind."""
    path = DATA_DIR / "eval" / "retrieval_queries.jsonl"
    by_kind: dict[str, list[RetrievalCase]] = defaultdict(list)
    for r in sorted(map(json.loads, path.read_text(encoding="utf-8").splitlines()), key=lambda r: r["id"]):
        by_kind[r["kind"]].append(RetrievalCase(r["id"], r["query"], r["relevant_doc_ids"], r["kind"]))
    if split == "all":
        return [c for cs in by_kind.values() for c in cs]
    return [c for cs in by_kind.values() for i, c in enumerate(cs) if i % 2 == (split == "test")]

# script_id: day_025_chunking/compare
def compare_chunkers(docs: list[Document], cases: list[RetrievalCase], strategies: dict[str, dict],
                     k: int = 5, embed_fn: EmbedFn | None = None) -> list[dict]:
    """One row per strategy: chunk, embed, search every case, score doc-level hit@1 and hit@k."""
    embed_fn = embed_fn or default_embedder()[1]
    queries = embed_fn([c.query for c in cases])
    rows = []
    for label, kwargs in strategies.items():
        chunks = [c for d in docs for c in chunk(d, **kwargs)]
        texts = [embed_text(c) for c in chunks]
        sims = queries @ embed_fn(texts).T           # unit-length rows, so dot product = cosine
        hits = dict.fromkeys((1, k), 0)
        for case, row in zip(cases, sims):
            ranked = ranked_doc_ids([chunks[i] for i in np.argsort(-row)])
            for n in hits:
                hits[n] += bool(set(ranked[:n]) & set(case.relevant_doc_ids))
        rows.append({"strategy": label, "chunks": len(chunks),
                     "avg_tokens": round(sum(c.n_tokens for c in chunks) / len(chunks)),
                     "hit@1": hits[1] / len(cases), f"hit@{k}": hits[k] / len(cases),
                     "embed_usd": embed_cost(texts)})
    return rows

STRATEGIES = {
    "page (baseline)":   {"strategy": "fixed", "max_tokens": 8000, "overlap": 0, "add_header": False},
    "fixed 400/50":      {"strategy": "fixed", "add_header": False},
    "recursive 400/50":  {"strategy": "recursive", "add_header": False},
    "markdown 400/50":   {"strategy": "markdown", "add_header": False},
    "markdown + header": {"strategy": "markdown"},
}

if __name__ == "__main__":
    name, embed_fn = default_embedder()
    cases = load_retrieval_cases("dev")
    rows = compare_chunkers(load_corpus(), cases, STRATEGIES, embed_fn=embed_fn)
    print(f"embedder: {name} · {len(cases)} dev queries")
    print(f"{'strategy':18} {'chunks':>6} {'avg tok':>7} {'hit@1':>6} {'hit@5':>6} {'embed $':>8}")
    for r in rows:
        print(f"{r['strategy']:18} {r['chunks']:6} {r['avg_tokens']:7} {r['hit@1']:6.2f} "
              f"{r['hit@5']:6.2f} {r['embed_usd']:8.4f}")

embedder: local_embed (offline stand-in) · 20 dev queries
strategy           chunks avg tok  hit@1  hit@5  embed $
page (baseline)        45    1415   0.70   0.95   0.0013
fixed 400/50          199     359   0.60   1.00   0.0014
recursive 400/50      221     318   0.70   1.00   0.0014
markdown 400/50       423     152   0.60   0.95   0.0013
markdown + header     423     152   0.75   0.95   0.0015