brain-of-reese/app/rag/chunker.py

"""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.
* **container / network / volume / image / pod / kube / swap / os /
  endpoint / j2** (A9 revised 2026-08-27) — quadlet unit files (TOML) and
  Jinja templates; plain-text paragraph packing (``chunk_text``) — no
  format-specific splitter (owner decision).

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 2026-08-21 and 2026-08-27: md, markdown,
#: txt, yaml, yml, json, py + the quadlet family and j2 — plain text).
_FORMAT_CHUNKERS = {
    ".md": chunk_markdown,
    ".markdown": chunk_markdown,
    ".txt": chunk_text,
    ".yaml": chunk_yaml,
    ".yml": chunk_yaml,
    ".json": chunk_json,
    ".py": chunk_python,
    # A9 revised 2026-08-27 (owner permission): the full Podman quadlet
    # family + Jinja templates — plain-text paragraph packing, no
    # TOML/Jinja-aware splitter (owner decision).
    ".container": chunk_text,
    ".network": chunk_text,
    ".volume": chunk_text,
    ".image": chunk_text,
    ".pod": chunk_text,
    ".kube": chunk_text,
    ".swap": chunk_text,
    ".os": chunk_text,
    ".endpoint": chunk_text,
    ".j2": chunk_text,
}


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)