ai-resume/corpus/chunker.py

"""Code-aware chunking.

The rules come from `docs/phase-4-corpus.md` §4.1:

* 400–800 tokens per chunk, about 12% overlap;
* **never split inside a fenced block, a table or a front-matter block.**

The last rule is structural here, not a post-hoc check: the body is parsed into
*blocks* first and chunks are packed out of whole blocks, so a boundary can only
ever fall between blocks. Fenced code, tables, headings, list items and
paragraphs are all blocks; the large ones cannot be subdivided.

Front-matter is removed by `corpus.loader` before this module sees the body, and
parsing a body that still has front-matter is an error, so it cannot be split
either.

`tokens` are estimated with a word-and-punctuation regex. A real BPE tokenizer
counts 15–30% more for prose and more again for code, so a chunk that measures
800 here is a conservative upper bound.
"""

from __future__ import annotations

import re
from dataclasses import dataclass

# Packing targets (docs/phase-4-corpus.md §4.1).
TOKEN_TARGET = 600
TOKEN_MAX = 800
TOKEN_MIN = 400
OVERLAP_RATIO = 0.12

FENCE_RE = re.compile(r"^ {0,3}(`{3,}|~{3,})")
HEADING_RE = re.compile(r"^ {0,3}#{1,6}\s")
LIST_RE = re.compile(r"^(\s*)([-*+]|\d{1,3}[.)])\s+")
RULE_RE = re.compile(r"^ {0,3}([-*_])(\s*\1){2,}\s*$")
TOKEN_RE = re.compile(r"\w+|[^\w\s]", re.UNICODE)

PROTECTED_KINDS = ("fence", "table")


class ChunkError(RuntimeError):
    """The body could not be chunked without breaking the rules."""


@dataclass(frozen=True)
class Block:
    """An indivisible unit of Markdown. Never split, only packed."""

    kind: str  # fence | table | heading | list | quote | para | rule
    text: str
    tokens: int
    line_start: int  # 1-based, relative to the body this block came from
    line_end: int

    @property
    def protected(self) -> bool:
        return self.kind in PROTECTED_KINDS


@dataclass(frozen=True)
class Chunk:
    index: int
    text: str
    tokens: int
    line_start: int
    line_end: int
    blocks: tuple
    oversize: bool = False

    @property
    def kinds(self) -> tuple:
        return tuple(block.kind for block in self.blocks)


def estimate_tokens(text: str) -> int:
    """Word-plus-punctuation token estimate. Deterministic and cheap."""
    return len(TOKEN_RE.findall(text))


def parse_blocks(body: str) -> list:
    """Split a Markdown body into atomic blocks.

    Raises ChunkError if the body still carries a front-matter block: that is a
    loader bug, and silently chunking it would split something that must not be
    split.
    """
    if body.lstrip().startswith("---"):
        raise ChunkError(
            "body still starts with a front-matter block; strip front-matter "
            "before chunking (corpus.loader does this)"
        )

    lines = body.split("\n")
    blocks = []
    i = 0
    total = len(lines)
    while i < total:
        raw = lines[i]
        stripped = raw.strip()
        if not stripped:
            i += 1
            continue
        start = i
        fence = FENCE_RE.match(raw)
        if fence:
            i = _read_fence(lines, i, fence.group(1))
            kind = "fence"
        elif stripped.startswith("|"):
            while i < total and lines[i].strip().startswith("|"):
                i += 1
            kind = "table"
        elif HEADING_RE.match(raw):
            i += 1
            kind = "heading"
        elif RULE_RE.match(raw):
            i += 1
            kind = "rule"
        elif LIST_RE.match(raw):
            i = _read_list_item(lines, i)
            kind = "list"
        elif stripped.startswith(">"):
            while i < total and lines[i].strip().startswith(">"):
                i += 1
            kind = "quote"
        else:
            i = _read_paragraph(lines, i)
            kind = "para"
        text = "\n".join(lines[start:i])
        blocks.append(
            Block(kind=kind, text=text, tokens=estimate_tokens(text), line_start=start + 1, line_end=i)
        )
    return blocks


def _read_fence(lines: list, i: int, marker: str) -> int:
    """Return the index just past a fenced block. Unterminated runs to EOF."""
    char = marker[0]
    width = len(marker)
    i += 1
    while i < len(lines):
        candidate = FENCE_RE.match(lines[i])
        if candidate and candidate.group(1)[0] == char and len(candidate.group(1)) >= width:
            return i + 1
        i += 1
    return len(lines)


def _read_list_item(lines: list, i: int) -> int:
    """Return the index just past one top-level list item, continuations included."""
    base_indent = len(LIST_RE.match(lines[i]).group(1))
    i += 1
    while i < len(lines):
        raw = lines[i]
        if not raw.strip():
            # A blank line continues the item only if an indented line follows.
            if i + 1 < len(lines) and lines[i + 1].startswith(" " * (base_indent + 1)):
                i += 1
                continue
            return i
        indent = len(raw) - len(raw.lstrip())
        if indent > base_indent or LIST_RE.match(raw) and len(LIST_RE.match(raw).group(1)) > base_indent:
            i += 1
            continue
        if LIST_RE.match(raw):
            return i
        if indent <= base_indent:
            return i
        i += 1
    return i


def _read_paragraph(lines: list, i: int) -> int:
    i += 1
    while i < len(lines):
        raw = lines[i]
        stripped = raw.strip()
        if not stripped:
            return i
        if FENCE_RE.match(raw) or HEADING_RE.match(raw) or RULE_RE.match(raw):
            return i
        if stripped.startswith("|") or stripped.startswith(">") or LIST_RE.match(raw):
            return i
        i += 1
    return i


def chunk_blocks(blocks: list) -> list:
    """Pack blocks into windows of at most TOKEN_MAX, with ~12% overlap."""
    chunks = []
    total = len(blocks)
    if total == 0:
        return chunks
    i = 0
    while i < total:
        start = i
        tokens = 0
        end = i
        while end < total:
            block = blocks[end]
            if end > start:
                if tokens + block.tokens > TOKEN_MAX:
                    break
                # Prefer breaking in front of a heading.
                if block.kind == "heading" and tokens >= int(TOKEN_TARGET * 0.75):
                    break
            tokens += block.tokens
            end += 1
            if tokens >= TOKEN_TARGET and end < total and blocks[end].kind == "heading":
                break
        window = blocks[start:end]
        chunks.append(_make_chunk(len(chunks), window))
        if end >= total:
            break
        # Back up over trailing blocks until roughly OVERLAP_RATIO of the window
        # repeats, leaving at least one block behind so the walk always advances.
        budget = max(1, int(round(TOKEN_TARGET * OVERLAP_RATIO)))
        carried = 0
        back = end
        while back - 1 > start and carried + blocks[back - 1].tokens <= budget * 2:
            carried += blocks[back - 1].tokens
            back -= 1
            if carried >= budget:
                break
        i = max(back, start + 1)
    return chunks


def _make_chunk(index: int, window: list) -> Chunk:
    text = "\n\n".join(block.text for block in window)
    return Chunk(
        index=index,
        text=text,
        tokens=sum(block.tokens for block in window),
        line_start=window[0].line_start,
        line_end=window[-1].line_end,
        blocks=tuple(window),
        oversize=len(window) == 1 and window[0].tokens > TOKEN_MAX,
    )


def chunk_document(body: str) -> list:
    """Chunk one document body. The pipeline calls this and nothing else."""
    chunks = chunk_blocks(parse_blocks(body))
    verify_no_split(chunks, protected_regions(body))
    return chunks


def protected_regions(body: str) -> list:
    """Line ranges (1-based, inclusive) that a chunk boundary must not cross.

    Recomputed from the body text alone, independently of the packing code, so
    the check is not the implementation grading its own homework.
    """
    regions = []
    lines = body.split("\n")
    i = 0
    while i < len(lines):
        raw = lines[i]
        fence = FENCE_RE.match(raw)
        if fence:
            end = _read_fence(lines, i, fence.group(1))
            regions.append((i + 1, end))
            i = end
            continue
        if raw.strip().startswith("|"):
            start = i
            while i < len(lines) and lines[i].strip().startswith("|"):
                i += 1
            regions.append((start + 1, i))
            continue
        i += 1
    return regions


def verify_no_split(chunks: list, regions: list) -> None:
    """Raise ChunkError if any chunk boundary cuts a protected region."""
    for chunk in chunks:
        for first, last in regions:
            spans_start = chunk.line_start <= first <= chunk.line_end
            spans_end = chunk.line_start <= last <= chunk.line_end
            covers = chunk.line_start <= first and chunk.line_end >= last
            if (spans_start or spans_end) and not covers:
                raise ChunkError(
                    f"chunk {chunk.index} splits protected lines {first}-{last} "
                    f"(chunk covers {chunk.line_start}-{chunk.line_end})"
                )