nanobot: 2026-09-10 12:33:37
This commit is contained in:
lachtan
260
skills/wiki/scripts/wiki_chunker.py
Normal file
260
skills/wiki/scripts/wiki_chunker.py
Normal file
@@ -0,0 +1,260 @@
|
||||
#!/usr/bin/env python3
|
||||
# /// script
|
||||
# requires-python = ">=3.11"
|
||||
# dependencies = ["markdown-it-py", "pyyaml"]
|
||||
# ///
|
||||
"""Structure-aware markdown chunker for the wiki skill.
|
||||
|
||||
Splits a document along its H1-H3 heading hierarchy, then merges small siblings and
|
||||
splits oversized sections along block boundaries. Every chunk carries a breadcrumb
|
||||
(`path > title > section > subsection`) which goes into the embedded text as well as
|
||||
the metadata, so a vector represents a passage in context rather than in isolation.
|
||||
|
||||
markdown-it-py supplies the AST (it knows the CommonMark edge cases); the chunking
|
||||
policy below is ours. Block boundaries come from token line maps, so the emitted text
|
||||
is the original markdown — code fences and tables stay byte-for-byte intact.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
import yaml
|
||||
from markdown_it import MarkdownIt # ty: ignore[unresolved-import]
|
||||
|
||||
CHUNKER_VERSION = "1"
|
||||
|
||||
MERGE_BELOW = 200
|
||||
SPLIT_ABOVE = 800
|
||||
OVERLAP = 64
|
||||
CHARS_PER_TOKEN = 4
|
||||
|
||||
SECTION_LEVELS = (1, 2, 3)
|
||||
BREADCRUMB_SEP = " > "
|
||||
|
||||
_FRONTMATTER_RE = re.compile(r"\A---[ \t]*\r?\n(.*?)\r?\n---[ \t]*(?:\r?\n|\Z)", re.DOTALL)
|
||||
|
||||
_md = MarkdownIt("commonmark")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Chunk:
|
||||
breadcrumb: str
|
||||
text: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ParsedFile:
|
||||
title: str | None = None
|
||||
tags: list[str] = field(default_factory=list)
|
||||
headings: list[str] = field(default_factory=list)
|
||||
chunks: list[Chunk] = field(default_factory=list)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Heading:
|
||||
start: int
|
||||
end: int
|
||||
level: int
|
||||
title: str
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Section:
|
||||
breadcrumb: tuple[str, ...]
|
||||
text: str
|
||||
|
||||
@property
|
||||
def parent(self) -> tuple[str, ...]:
|
||||
return self.breadcrumb[:-1]
|
||||
|
||||
|
||||
def token_estimate(text: str) -> int:
|
||||
"""Approximate token count. Sizing does not need a real tokenizer (plan: 4 chars/token)."""
|
||||
return max(1, len(text) // CHARS_PER_TOKEN)
|
||||
|
||||
|
||||
def parse_markdown(text: str, path: str) -> ParsedFile:
|
||||
"""Chunk one markdown document. `path` is the source-relative path, the breadcrumb root."""
|
||||
frontmatter, body = _split_frontmatter(text)
|
||||
frontmatter_title = _frontmatter_title(frontmatter)
|
||||
tags = _frontmatter_tags(frontmatter)
|
||||
|
||||
root: tuple[str, ...] = (path,) if frontmatter_title is None else (path, frontmatter_title)
|
||||
sections, headings = _split_sections(body, root)
|
||||
merged = _merge_small_siblings(sections)
|
||||
chunks = [chunk for section in merged for chunk in _split_oversized(section)]
|
||||
|
||||
# The catalog title falls back to the first heading, because notes carry their title as
|
||||
# `# H1` far more often than as frontmatter. It deliberately does NOT feed the breadcrumb
|
||||
# root: the H1 already reaches the breadcrumb through the heading stack, and changing the
|
||||
# root would rewrite every chunk's text and force a full reindex.
|
||||
title = frontmatter_title or (headings[0] if headings else None)
|
||||
return ParsedFile(title=title, tags=tags, headings=headings, chunks=chunks)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# frontmatter
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _split_frontmatter(text: str) -> tuple[dict, str]:
|
||||
"""Peel off a leading YAML frontmatter block. Malformed frontmatter stays body text."""
|
||||
match = _FRONTMATTER_RE.match(text)
|
||||
if not match:
|
||||
return {}, text
|
||||
try:
|
||||
data = yaml.safe_load(match.group(1))
|
||||
except yaml.YAMLError:
|
||||
return {}, text
|
||||
if not isinstance(data, dict):
|
||||
return {}, text
|
||||
return data, text[match.end() :]
|
||||
|
||||
|
||||
def _frontmatter_title(frontmatter: dict) -> str | None:
|
||||
title = frontmatter.get("title")
|
||||
if isinstance(title, str) and title.strip():
|
||||
return title.strip()
|
||||
return None
|
||||
|
||||
|
||||
def _frontmatter_tags(frontmatter: dict) -> list[str]:
|
||||
raw = frontmatter.get("tags")
|
||||
if isinstance(raw, str):
|
||||
values = raw.split(",")
|
||||
elif isinstance(raw, list):
|
||||
values = [str(item) for item in raw]
|
||||
else:
|
||||
return []
|
||||
return [tag.strip() for tag in values if tag.strip()]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# sectioning
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _split_sections(body: str, root: tuple[str, ...]) -> tuple[list[_Section], list[str]]:
|
||||
"""Cut the body at H1-H3 headings. H4+ stay inside their parent section."""
|
||||
lines = body.split("\n")
|
||||
tokens = _md.parse(body)
|
||||
|
||||
starts: list[_Heading] = []
|
||||
for index, token in enumerate(tokens):
|
||||
if token.type != "heading_open" or token.map is None:
|
||||
continue
|
||||
level = int(token.tag[1:])
|
||||
if level not in SECTION_LEVELS:
|
||||
continue
|
||||
inline = tokens[index + 1] if index + 1 < len(tokens) else None
|
||||
title = inline.content.strip() if inline is not None else ""
|
||||
starts.append(_Heading(start=token.map[0], end=token.map[1], level=level, title=title))
|
||||
|
||||
headings = [heading.title for heading in starts]
|
||||
boundaries = [heading.start for heading in starts] + [len(lines)]
|
||||
|
||||
sections: list[_Section] = []
|
||||
preamble = "\n".join(lines[: boundaries[0]]).strip()
|
||||
if preamble:
|
||||
sections.append(_Section(root, preamble))
|
||||
|
||||
stack: list[tuple[int, str]] = []
|
||||
for position, heading in enumerate(starts):
|
||||
while stack and stack[-1][0] >= heading.level:
|
||||
stack.pop()
|
||||
stack.append((heading.level, heading.title))
|
||||
section_end = boundaries[position + 1]
|
||||
# A heading with no prose of its own would embed as a bare title; the heading
|
||||
# still reaches the index through its children's breadcrumbs, so drop it.
|
||||
if not "\n".join(lines[heading.end : section_end]).strip():
|
||||
continue
|
||||
text = "\n".join(lines[heading.start : section_end]).strip()
|
||||
sections.append(_Section(_dedupe(root + tuple(title for _, title in stack)), text))
|
||||
return sections, headings
|
||||
|
||||
|
||||
def _dedupe(parts: tuple[str, ...]) -> tuple[str, ...]:
|
||||
"""Drop consecutive repeats so a frontmatter title matching the H1 shows up once."""
|
||||
out: list[str] = []
|
||||
for part in parts:
|
||||
if not out or out[-1] != part:
|
||||
out.append(part)
|
||||
return tuple(out)
|
||||
|
||||
|
||||
def _merge_small_siblings(sections: list[_Section]) -> list[_Section]:
|
||||
"""Glue a too-small section onto the following sibling under the same parent."""
|
||||
merged: list[_Section] = []
|
||||
for section in sections:
|
||||
if not merged:
|
||||
merged.append(section)
|
||||
continue
|
||||
previous = merged[-1]
|
||||
combined = f"{previous.text}\n\n{section.text}"
|
||||
if (
|
||||
token_estimate(previous.text) < MERGE_BELOW
|
||||
and previous.parent == section.parent
|
||||
and token_estimate(combined) <= SPLIT_ABOVE
|
||||
):
|
||||
merged[-1] = _Section(previous.breadcrumb, combined)
|
||||
else:
|
||||
merged.append(section)
|
||||
return merged
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# splitting
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _split_oversized(section: _Section) -> list[Chunk]:
|
||||
"""Break a section over SPLIT_ABOVE into block-aligned pieces with OVERLAP carry-over."""
|
||||
breadcrumb = BREADCRUMB_SEP.join(section.breadcrumb)
|
||||
if token_estimate(section.text) <= SPLIT_ABOVE:
|
||||
return [Chunk(breadcrumb, _embed_text(breadcrumb, section.text))]
|
||||
|
||||
blocks = _top_level_blocks(section.text)
|
||||
chunks: list[Chunk] = []
|
||||
current: list[str] = []
|
||||
for block in blocks:
|
||||
candidate = [*current, block]
|
||||
if current and token_estimate("\n\n".join(candidate)) > SPLIT_ABOVE:
|
||||
chunks.append(Chunk(breadcrumb, _embed_text(breadcrumb, "\n\n".join(current))))
|
||||
current = [*_overlap_tail(current), block]
|
||||
else:
|
||||
current = candidate
|
||||
if current:
|
||||
chunks.append(Chunk(breadcrumb, _embed_text(breadcrumb, "\n\n".join(current))))
|
||||
return chunks
|
||||
|
||||
|
||||
def _top_level_blocks(text: str) -> list[str]:
|
||||
"""Top-level markdown blocks, taken from token line maps so fences stay whole."""
|
||||
lines = text.split("\n")
|
||||
tokens = _md.parse(text)
|
||||
starts = sorted({token.map[0] for token in tokens if token.level == 0 and token.map})
|
||||
if not starts:
|
||||
return [text]
|
||||
bounds = [*starts, len(lines)]
|
||||
blocks = ["\n".join(lines[bounds[i] : bounds[i + 1]]).strip() for i in range(len(starts))]
|
||||
return [block for block in blocks if block]
|
||||
|
||||
|
||||
def _overlap_tail(blocks: list[str]) -> list[str]:
|
||||
"""Trailing whole blocks of the emitted chunk, up to OVERLAP tokens."""
|
||||
tail: list[str] = []
|
||||
budget = OVERLAP
|
||||
for block in reversed(blocks):
|
||||
cost = token_estimate(block)
|
||||
if cost > budget:
|
||||
break
|
||||
tail.insert(0, block)
|
||||
budget -= cost
|
||||
return tail
|
||||
|
||||
|
||||
def _embed_text(breadcrumb: str, text: str) -> str:
|
||||
"""The stored chunk text carries its breadcrumb, so the vector sees the context."""
|
||||
return f"{breadcrumb}\n\n{text}"
|
||||
Reference in New Issue
Block a user