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