Chunking is a retrieval decision, not a preprocessing step

from dataclasses import dataclass, field
from typing import Iterator


@dataclass(slots=True)
class Chunk:
    """One retrievable unit, plus everything needed to rank and cite it."""

    doc_id: str
    text: str
    heading_path: list[str] = field(default_factory=list)
    kind: str = "prose"          # prose | table | code | definition
    start: int = 0
    end: int = 0

    @property
    def embedding_text(self) -> str:
        """What the embedder sees: the chunk under its own headings."""
        trail = " > ".join(self.heading_path)
        return f"{trail}\n\n{self.text}" if trail else self.text


def split_markdown(doc_id: str, source: str) -> Iterator[Chunk]:
    """Split on structure first, size second.

    Headings open a new chunk. Tables and fenced code are never split, however
    long they run — a half table retrieves for nothing.
    """
    heading_path: list[str] = []
    buffer: list[str] = []
    kind = "prose"
    start = 0

    def flush(end: int) -> Iterator[Chunk]:
        body = "\n".join(buffer).strip()
        if body:
            yield Chunk(doc_id, body, list(heading_path), kind, start, end)

    for offset, line in enumerate(source.splitlines()):
        if line.startswith("#"):
            yield from flush(offset)
            buffer, kind, start = [], "prose", offset
            level = len(line) - len(line.lstrip("#"))
            heading_path = heading_path[: level - 1] + [line.lstrip("# ").strip()]
            continue

        if line.startswith("```"):
            kind = "prose" if kind == "code" else "code"

        if line.startswith("|") and kind == "prose":
            kind = "table"

        buffer.append(line)

    yield from flush(len(source.splitlines()))

from dataclasses import dataclass
from statistics import mean
from typing import Callable, Sequence

Retriever = Callable[[str, int], Sequence[str]]


@dataclass(frozen=True, slots=True)
class LabelledQuery:
    """One real question, and the chunk ids that genuinely answer it."""

    question: str
    gold_chunk_ids: frozenset[str]


def recall_at_k(retriever: Retriever, queries: Sequence[LabelledQuery], k: int) -> float:
    """Share of questions whose answer appears anywhere in the top k."""
    hits = 0
    for query in queries:
        retrieved = set(retriever(query.question, k))
        if retrieved & query.gold_chunk_ids:
            hits += 1
    return hits / len(queries)


def mrr(retriever: Retriever, queries: Sequence[LabelledQuery], k: int) -> float:
    """How far down the list the first correct chunk sits, averaged."""
    scores = []
    for query in queries:
        rank = next(
            (i for i, chunk_id in enumerate(retriever(query.question, k), start=1)
             if chunk_id in query.gold_chunk_ids),
            None,
        )
        scores.append(1 / rank if rank else 0.0)
    return mean(scores)


def compare(strategies: dict[str, Retriever], queries: Sequence[LabelledQuery], k: int = 10) -> None:
    width = max(len(name) for name in strategies)
    for name, retriever in strategies.items():
        print(
            f"{name:<{width}}  "
            f"recall@{k} {recall_at_k(retriever, queries, k):.2f}  "
            f"mrr@{k} {mrr(retriever, queries, k):.2f}"
        )

question class      strategy            recall@10   mrr@10
prose               fixed-512           ....        ....
prose               structural          ....        ....
identifier          fixed-512           ....        ....
identifier          structural          ....        ....

import { z } from "zod";

export const ChunkSchema = z.object({
  docId: z.string(),
  text: z.string().min(1),
  headingPath: z.array(z.string()),
  kind: z.enum(["prose", "table", "code", "definition"]),
  start: z.number().int().nonnegative(),
  end: z.number().int().nonnegative(),
});

export type Chunk = z.infer<typeof ChunkSchema>;

/** Never let a table or a fenced block be cut in half by a size limit. */
export function enforceBoundaries(chunks: Chunk[], maxChars: number): Chunk[] {
  return chunks.flatMap((chunk) => {
    if (chunk.text.length <= maxChars || chunk.kind !== "prose") {
      return [chunk];
    }

    const paragraphs = chunk.text.split(/\n{2,}/);
    const packed: Chunk[] = [];

    for (const paragraph of paragraphs) {
      const last = packed.at(-1);

      if (last && last.text.length + paragraph.length <= maxChars) {
        last.text = `${last.text}\n\n${paragraph}`;
        continue;
      }

      packed.push({ ...chunk, text: paragraph });
    }

    return packed;
  });
}