"""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})" )