409 lines
13 KiB
Python
409 lines
13 KiB
Python
"""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})"
|
||
)
|
||
|
||
|
||
# -- code -----------------------------------------------------------------
|
||
|
||
|
||
def _code_depth(line: str, depth: int) -> int:
|
||
"""Approximate bracket depth for one line, ignoring strings and comments.
|
||
|
||
This is a heuristic lexer, not a parser: it is here to answer one question —
|
||
"is this a safe line to break after?" — and it errs towards *not* breaking.
|
||
A line with an ambiguous quote or an unterminated bracket simply keeps the
|
||
chunk open, which can only produce a larger chunk, never a broken one.
|
||
"""
|
||
index = 0
|
||
quote = None
|
||
length = len(line)
|
||
while index < length:
|
||
char = line[index]
|
||
if quote:
|
||
if char == "\\":
|
||
index += 2
|
||
continue
|
||
if line.startswith(quote, index):
|
||
index += len(quote)
|
||
quote = None
|
||
continue
|
||
index += 1
|
||
continue
|
||
if line.startswith("//", index) or line.startswith("#", index) and line.lstrip().startswith("#"):
|
||
break
|
||
if char in "\"'`":
|
||
triple = char * 3
|
||
if line.startswith(triple, index):
|
||
quote = triple
|
||
index += 3
|
||
continue
|
||
quote = char
|
||
index += 1
|
||
continue
|
||
if char in "([{":
|
||
depth += 1
|
||
elif char in ")]}":
|
||
depth = max(0, depth - 1)
|
||
index += 1
|
||
return depth
|
||
|
||
|
||
def code_segments(text: str) -> list:
|
||
"""Split code into whole top-level segments. A function is never cut."""
|
||
lines = text.split("\n")
|
||
segments = []
|
||
current = []
|
||
depth = 0
|
||
start = 1
|
||
for number, line in enumerate(lines, start=1):
|
||
if not current:
|
||
start = number
|
||
depth = _code_depth(line, depth)
|
||
current.append(line)
|
||
# A blank line at bracket depth zero is the only automatic boundary.
|
||
if depth == 0 and not line.strip() and len(current) > 1:
|
||
segments.append((start, number, "\n".join(current)))
|
||
current = []
|
||
if current:
|
||
segments.append((start, len(lines), "\n".join(current)))
|
||
return [(start, end, body) for start, end, body in segments if body.strip()]
|
||
|
||
|
||
def chunk_code(text: str) -> list:
|
||
"""Chunk a source file without cutting a function, class or block.
|
||
|
||
Segments are whole top-level units. A unit larger than the ceiling is kept
|
||
as one oversize chunk and flagged, exactly like an oversized Markdown fence:
|
||
cutting it would break the rule this function exists to keep.
|
||
|
||
Overlap repeats the previous unit, never part of one, so an overlapping
|
||
chunk is still syntactically whole.
|
||
"""
|
||
segments = code_segments(text)
|
||
if not segments:
|
||
raise ChunkError("source file produced no segments (empty?)")
|
||
|
||
chunks = []
|
||
current = [] # list of (start, end, body)
|
||
tokens = 0
|
||
|
||
def flush():
|
||
nonlocal current, tokens
|
||
if not current:
|
||
return
|
||
body = "\n".join(segment[2] for segment in current).strip("\n")
|
||
chunks.append(
|
||
Chunk(
|
||
index=len(chunks),
|
||
text=body,
|
||
tokens=estimate_tokens(body),
|
||
line_start=current[0][0],
|
||
line_end=current[-1][1],
|
||
blocks=tuple(
|
||
Block("code", segment[2], estimate_tokens(segment[2]), segment[0], segment[1])
|
||
for segment in current
|
||
),
|
||
oversize=len(current) == 1 and estimate_tokens(current[0][2]) > TOKEN_MAX,
|
||
)
|
||
)
|
||
current = []
|
||
tokens = 0
|
||
|
||
for segment in segments:
|
||
segment_tokens = estimate_tokens(segment[2])
|
||
if current and tokens + segment_tokens > TOKEN_MAX:
|
||
previous = current[-1]
|
||
flush()
|
||
budget = int(round(TOKEN_TARGET * OVERLAP_RATIO))
|
||
if estimate_tokens(previous[2]) <= budget:
|
||
current = [previous]
|
||
tokens = estimate_tokens(previous[2])
|
||
current.append(segment)
|
||
tokens += segment_tokens
|
||
flush()
|
||
return chunks
|