361 lines
14 KiB
Python
361 lines
14 KiB
Python
"""Format-aware chunker (PLAN §5 chunking policy).
|
|
|
|
Pure functions, no I/O — fully unit-testable. Stdlib only.
|
|
|
|
:func:`chunk_document` dispatches on the file's lowercased suffix;
|
|
per-format policies:
|
|
|
|
* **md / markdown** — sections are split on ATX headings of level ≥ 2
|
|
(``## ``/``### ``/…); a section that fits in ``target_chars`` becomes a
|
|
single chunk, a longer one is sub-split at paragraph boundaries (blank
|
|
lines outside code fences) with ``overlap_chars`` carry-over, and every
|
|
chunk keeps its nearest preceding heading line (the section anchor). Code
|
|
fences are atomic (a boundary never falls inside one) except a fence
|
|
larger than :data:`HARD_MAX_CHARS`, which is split by line.
|
|
* **yaml / yml** — blocks start at ``---`` document separators and at
|
|
top-level (indent-0) ``key:`` lines; every chunk keeps its key lines as
|
|
anchors, so a hit is always readable in context.
|
|
* **json** — pretty-printed (``json.dumps(obj, indent=2)``) and split on
|
|
top-level keys (one ``{key: value}`` block per key); unparseable input
|
|
falls back to paragraph packing.
|
|
* **py** — split at top-level defs/classes via the stdlib ``ast`` (the
|
|
module preamble — imports, constants — is its own block); an oversized
|
|
definition falls back to line packing.
|
|
* **txt** (and any unknown suffix) — paragraph packing.
|
|
|
|
Every format honors :data:`HARD_MAX_CHARS` (1200 — the aipi ~1024-token
|
|
request cap) and the target/overlap settings; oversized blocks are split
|
|
by line so no chunk can exceed the cap.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import ast
|
|
import json
|
|
import re
|
|
from collections.abc import Sequence
|
|
|
|
#: ATX heading of level ≥ 2 — the section splitter (PLAN §5).
|
|
_HEADING_RE = re.compile(r"^(#{2,6})\s+\S")
|
|
#: First H1 — the document title (used by :func:`extract_title`).
|
|
_H1_RE = re.compile(r"^#\s+\S")
|
|
#: Opening/closing code fence (three or more backticks or tildes).
|
|
_FENCE_RE = re.compile(r"^(`{3,}|~{3,})")
|
|
#: Absolute per-chunk size cap (chars). Measured worst case in the real KB:
|
|
#: punctuation-heavy machine output (``lspci`` dumps in fences) tokenizes at
|
|
#: ~1.4 chars/token, so 1200 chars stays under the endpoint's ~1024-token
|
|
#: per-request input cap even as a solo-chunk request.
|
|
HARD_MAX_CHARS = 1200
|
|
|
|
|
|
def extract_title(content: str, fallback: str = "") -> str:
|
|
"""First markdown H1, else *fallback* (the importer passes the file stem)."""
|
|
for line in content.splitlines():
|
|
m = _H1_RE.match(line)
|
|
if m:
|
|
return line.lstrip("#").strip() or fallback
|
|
return fallback
|
|
|
|
|
|
def _iter_fence_state(lines: Sequence[str]) -> list[bool]:
|
|
"""Per-line flags: ``True`` on a fence delimiter or inside a fence."""
|
|
flags: list[bool] = []
|
|
marker: str | None = None
|
|
for line in lines:
|
|
stripped = line.lstrip()
|
|
if marker is None:
|
|
m = _FENCE_RE.match(stripped)
|
|
if m:
|
|
marker = m.group(1)[:3]
|
|
flags.append(True)
|
|
else:
|
|
flags.append(False)
|
|
else:
|
|
flags.append(True)
|
|
if stripped.startswith(marker):
|
|
marker = None
|
|
return flags
|
|
|
|
|
|
def _section_ranges(lines: Sequence[str], flags: Sequence[bool]) -> list[tuple[int, int]]:
|
|
"""(start, end) line-index ranges of every heading-delimited section."""
|
|
starts = [0]
|
|
for i, line in enumerate(lines):
|
|
if i > 0 and not flags[i] and _HEADING_RE.match(line):
|
|
starts.append(i)
|
|
return list(zip(starts, [*starts[1:], len(lines)], strict=True))
|
|
|
|
|
|
def _paragraph_blocks(lines: Sequence[str], flags: Sequence[bool]) -> list[str]:
|
|
"""Blocks of lines separated by blank lines *outside* fences.
|
|
|
|
A fenced code block (possibly containing blank lines) is one block.
|
|
"""
|
|
blocks: list[str] = []
|
|
cur: list[str] = []
|
|
for line, in_fence in zip(lines, flags, strict=True):
|
|
if in_fence or line.strip():
|
|
cur.append(line)
|
|
elif cur:
|
|
blocks.append("\n".join(cur))
|
|
cur = []
|
|
if cur:
|
|
blocks.append("\n".join(cur))
|
|
return blocks
|
|
|
|
|
|
def _split_oversized(block: str, hard: int = HARD_MAX_CHARS) -> list[str]:
|
|
"""Split one oversized block (e.g. a huge fenced code block) into line
|
|
groups under *hard* chars. A single line longer than *hard* is chopped
|
|
at char boundaries — better than exceeding the endpoint's token cap."""
|
|
parts: list[str] = []
|
|
cur: list[str] = []
|
|
cur_len = 0
|
|
for line in block.splitlines():
|
|
if len(line) > hard:
|
|
if cur:
|
|
parts.append("\n".join(cur))
|
|
cur, cur_len = [], 0
|
|
parts.extend(line[i : i + hard] for i in range(0, len(line), hard))
|
|
continue
|
|
if cur and cur_len + 1 + len(line) > hard:
|
|
parts.append("\n".join(cur))
|
|
cur, cur_len = [], 0
|
|
cur.append(line)
|
|
cur_len += 1 + len(line)
|
|
if cur:
|
|
parts.append("\n".join(cur))
|
|
return parts
|
|
|
|
|
|
def _pack_blocks(blocks: Sequence[str], target: int, overlap: int) -> list[str]:
|
|
"""Greedy paragraph packing; consecutive chunks share ``overlap`` chars."""
|
|
# Oversized blocks are split to the *budget*, so a solo-block chunk
|
|
# (anchor + block) cannot exceed the target — let alone the hard cap.
|
|
split_at = min(HARD_MAX_CHARS, max(1, target))
|
|
expanded: list[str] = []
|
|
for block in blocks:
|
|
if len(block) > split_at:
|
|
expanded.extend(_split_oversized(block, split_at))
|
|
else:
|
|
expanded.append(block)
|
|
chunks: list[str] = []
|
|
cur_parts: list[str] = []
|
|
cur_len = 0
|
|
for block in expanded:
|
|
if cur_len == 0:
|
|
cur_parts.append(block)
|
|
cur_len = len(block)
|
|
elif cur_len + 1 + len(block) <= target:
|
|
cur_parts.append(block)
|
|
cur_len += 1 + len(block)
|
|
else:
|
|
chunks.append("\n".join(cur_parts))
|
|
# The tail is only carried over when tail + block still fits the
|
|
# budget — otherwise the chunk would exceed the size cap.
|
|
tail = (
|
|
chunks[-1][-overlap:]
|
|
if (overlap > 0 and overlap + 1 + len(block) <= target)
|
|
else ""
|
|
)
|
|
cur_parts = [tail, block] if tail else [block]
|
|
cur_len = len(tail) + 1 + len(block) if tail else len(block)
|
|
if cur_parts:
|
|
chunks.append("\n".join(cur_parts))
|
|
return chunks
|
|
|
|
|
|
def _chunk_section(
|
|
lines: Sequence[str], flags: Sequence[bool], target: int, overlap: int
|
|
) -> list[str]:
|
|
"""Chunk one section (its heading line included, when it has one)."""
|
|
first = next((i for i, line in enumerate(lines) if line.strip()), None)
|
|
if first is None:
|
|
return []
|
|
anchor: str | None = None
|
|
if not flags[first] and _HEADING_RE.match(lines[first]):
|
|
anchor = lines[first]
|
|
body: Sequence[str] = lines[first + 1 :]
|
|
body_flags: Sequence[bool] = flags[first + 1 :]
|
|
else:
|
|
body = lines
|
|
body_flags = flags
|
|
|
|
full = (anchor + "\n" + "\n".join(body) if anchor is not None else "\n".join(body)).strip()
|
|
if not full:
|
|
return []
|
|
if len(full) <= target:
|
|
return [full]
|
|
|
|
blocks = _paragraph_blocks(body, body_flags)
|
|
if not blocks:
|
|
return [full]
|
|
budget = target - (len(anchor) + 2 if anchor is not None else 0)
|
|
packed = _pack_blocks(blocks, max(1, budget), overlap)
|
|
if anchor is None:
|
|
return packed
|
|
return [(anchor + "\n\n" + part).strip() for part in packed]
|
|
|
|
|
|
def chunk_markdown(
|
|
content: str, target_chars: int = 2000, overlap_chars: int = 200
|
|
) -> list[str]:
|
|
"""Split markdown into retrieval chunks (see module docstring for policy)."""
|
|
if target_chars <= 0:
|
|
raise ValueError("target_chars must be > 0")
|
|
if overlap_chars < 0:
|
|
raise ValueError("overlap_chars must be >= 0")
|
|
# The endpoint's token cap is absolute — a larger target is unsafe.
|
|
target = min(target_chars, HARD_MAX_CHARS)
|
|
overlap = min(overlap_chars, target - 1)
|
|
|
|
lines = content.splitlines()
|
|
flags = _iter_fence_state(lines)
|
|
chunks: list[str] = []
|
|
for start, end in _section_ranges(lines, flags):
|
|
chunks.extend(_chunk_section(lines[start:end], flags[start:end], target, overlap))
|
|
return [c for c in chunks if c.strip()]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Non-markdown formats (A9, revised 2026-08-21): yaml/yml, json, py, txt
|
|
# ---------------------------------------------------------------------------
|
|
|
|
#: YAML document separator (column 0).
|
|
_YAML_DOC_SEP_RE = re.compile(r"^-{3,}\s*$")
|
|
#: Top-level YAML key (column 0, no leading whitespace) — the block anchor.
|
|
_YAML_KEY_RE = re.compile(r"^[A-Za-z0-9_.\-]+\s*:")
|
|
|
|
|
|
def _yaml_blocks(lines: Sequence[str]) -> list[str]:
|
|
"""Group YAML lines into blocks: ``---`` separators and indent-0
|
|
``key:`` lines each start a new block (the key line stays the anchor)."""
|
|
blocks: list[str] = []
|
|
cur: list[str] = []
|
|
for line in lines:
|
|
if cur and (_YAML_DOC_SEP_RE.match(line) or _YAML_KEY_RE.match(line)):
|
|
blocks.append("\n".join(cur))
|
|
cur = []
|
|
cur.append(line)
|
|
if cur:
|
|
blocks.append("\n".join(cur))
|
|
return [b for b in blocks if b.strip()]
|
|
|
|
|
|
def chunk_yaml(content: str, target_chars: int = 2000, overlap_chars: int = 200) -> list[str]:
|
|
"""Split YAML on document separators + top-level keys (see module docstring)."""
|
|
target, overlap = _normalize_target_overlap(target_chars, overlap_chars)
|
|
return _pack_blocks(_yaml_blocks(content.splitlines()), target, overlap)
|
|
|
|
|
|
def _json_blocks(content: str) -> list[str] | None:
|
|
"""Pretty-printed per-top-level-key blocks, or ``None`` if unparseable."""
|
|
try:
|
|
obj = json.loads(content)
|
|
except (ValueError, TypeError):
|
|
return None
|
|
if isinstance(obj, dict):
|
|
return [json.dumps({k: v}, indent=2) for k, v in obj.items()]
|
|
# Top-level list/scalar: nothing to key on — one pretty-printed block.
|
|
return [json.dumps(obj, indent=2)]
|
|
|
|
|
|
def chunk_json(content: str, target_chars: int = 2000, overlap_chars: int = 200) -> list[str]:
|
|
"""Pretty-print JSON and split on top-level keys (unparseable → paragraphs)."""
|
|
target, overlap = _normalize_target_overlap(target_chars, overlap_chars)
|
|
blocks = _json_blocks(content)
|
|
if blocks is None:
|
|
return chunk_text(content, target, overlap)
|
|
return _pack_blocks(blocks, target, overlap)
|
|
|
|
|
|
def _python_blocks(content: str) -> list[str] | None:
|
|
"""Line blocks: module preamble, then one per top-level def/class.
|
|
|
|
Returns ``None`` when the source does not parse (→ line/paragraph
|
|
packing fallback).
|
|
"""
|
|
try:
|
|
tree = ast.parse(content)
|
|
except (SyntaxError, ValueError):
|
|
return None
|
|
lines = content.splitlines()
|
|
tops = [
|
|
node
|
|
for node in tree.body
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef))
|
|
]
|
|
tops.sort(key=lambda n: n.lineno)
|
|
ranges: list[tuple[int, int]] = []
|
|
for node in tops:
|
|
start = node.lineno - 1
|
|
for dec in node.decorator_list:
|
|
start = min(start, dec.lineno - 1)
|
|
end = node.end_lineno or node.lineno # end_lineno is None on odd parses
|
|
ranges.append((start, end)) # 0-based start, 1-based end
|
|
blocks: list[str] = []
|
|
cursor = 0
|
|
for start, end in ranges:
|
|
if start > cursor:
|
|
blocks.append("\n".join(lines[cursor:start]))
|
|
blocks.append("\n".join(lines[start:end]))
|
|
cursor = end
|
|
if cursor < len(lines):
|
|
blocks.append("\n".join(lines[cursor:]))
|
|
return [b for b in blocks if b.strip()]
|
|
|
|
|
|
def chunk_python(content: str, target_chars: int = 2000, overlap_chars: int = 200) -> list[str]:
|
|
"""Split Python on top-level defs/classes (stdlib ``ast``; see module doc)."""
|
|
target, overlap = _normalize_target_overlap(target_chars, overlap_chars)
|
|
blocks = _python_blocks(content)
|
|
if blocks is None:
|
|
return chunk_text(content, target, overlap)
|
|
return _pack_blocks(blocks, target, overlap)
|
|
|
|
|
|
def chunk_text(content: str, target_chars: int = 2000, overlap_chars: int = 200) -> list[str]:
|
|
"""Plain-text paragraph packing (blank lines separate paragraphs)."""
|
|
target, overlap = _normalize_target_overlap(target_chars, overlap_chars)
|
|
lines = content.splitlines()
|
|
return _pack_blocks(_paragraph_blocks(lines, [False] * len(lines)), target, overlap)
|
|
|
|
|
|
def _normalize_target_overlap(target_chars: int, overlap_chars: int) -> tuple[int, int]:
|
|
"""Validate + clamp the size policy (shared by every format)."""
|
|
if target_chars <= 0:
|
|
raise ValueError("target_chars must be > 0")
|
|
if overlap_chars < 0:
|
|
raise ValueError("overlap_chars must be >= 0")
|
|
# The endpoint's token cap is absolute — a larger target is unsafe.
|
|
target = min(target_chars, HARD_MAX_CHARS)
|
|
return target, min(overlap_chars, target - 1)
|
|
|
|
|
|
#: suffix → chunker (A9, revised: md, markdown, txt, yaml, yml, json, py).
|
|
_FORMAT_CHUNKERS = {
|
|
".md": chunk_markdown,
|
|
".markdown": chunk_markdown,
|
|
".txt": chunk_text,
|
|
".yaml": chunk_yaml,
|
|
".yml": chunk_yaml,
|
|
".json": chunk_json,
|
|
".py": chunk_python,
|
|
}
|
|
|
|
|
|
def chunk_document(
|
|
content: str,
|
|
path: str,
|
|
target_chars: int = 2000,
|
|
overlap_chars: int = 200,
|
|
) -> list[str]:
|
|
"""Chunk *content* according to *path*'s lowercased suffix.
|
|
|
|
Unknown suffixes fall back to plain-text paragraph packing (the
|
|
importer only passes A9-format files, so this is belt-and-braces).
|
|
"""
|
|
name = path.rsplit("/", 1)[-1]
|
|
suffix = "." + name.rsplit(".", 1)[-1].lower() if "." in name else ""
|
|
chunker = _FORMAT_CHUNKERS.get(suffix, chunk_text)
|
|
return chunker(content, target_chars, overlap_chars)
|