261 lines
8.9 KiB
Python
261 lines
8.9 KiB
Python
#!/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}"
|