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

"""Markdown-aware chunker (PLAN §5 chunking policy).

Pure functions, no I/O — fully unit-testable.

Policy
------
* **Sections** are split on ATX headings of level ≥ 2 (``## ``/``### ``/…).
* A section that fits in ``target_chars`` becomes a single chunk.
* A longer section is sub-split at paragraph boundaries (blank lines outside
  code fences); each chunk after the first starts with the trailing
  ``overlap_chars`` of the previous chunk so context survives the cut.
* Every chunk keeps its nearest preceding heading line (the section anchor),
  so a retrieval hit is always readable in context.
* **Code fences** (``` / ~~~) are atomic: a chunk boundary never falls
  inside one, and lines inside a fence are never mistaken for headings or
  paragraph breaks. One exception: a fence *larger than* :data:`HARD_MAX_CHARS`
  is split by line, because aipi's local embedding model rejects requests
  over ~1024 input tokens and a single 5000-char code block would blow
  past that on its own.
"""
from __future__ import annotations

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()]