stack/src/llm/chunk.py

"""Markdown-aware chunking with deterministic ids.

Ids are ``{item_key}:{content_hash[:12]}:{seq:04d}`` — same input, same
ids — so pgvector upserts are idempotent and ``index_state`` can skip
unchanged items by comparing hashes.

Sizes are in characters (~4 chars/token; 2000 chars ≈ 500 tokens, inside
every candidate embed model's window).
"""

from __future__ import annotations

import hashlib
import re
from dataclasses import dataclass
from typing import Mapping

from pfs.anchors import anchor_metadata


@dataclass(frozen=True)
class Paragraph:
    """One FR paragraph anchor (mirrors ``bib.frlink.Anchor``)."""

    p_id: int
    page: int
    ordinal: int
    text: str


@dataclass(frozen=True)
class Doc:
    key: str
    text: str
    metadata: dict[str, str]
    # Rules: the anchor paragraphs that make up ``text``; when set,
    # chunk_doc packs whole paragraphs and stamps p_id/page per chunk.
    paragraphs: tuple[Paragraph, ...] = ()
    # (section heading, local file path) pairs for PDF page lookup —
    # ``llm.pages.enrich_pdf_pages`` matches chunks' ``section`` to these.
    files: tuple[tuple[str, str], ...] = ()


@dataclass(frozen=True)
class Chunk:
    id: str
    text: str
    metadata: dict[str, str]


def content_hash(text: str) -> str:
    return hashlib.sha256(text.encode()).hexdigest()


_FRONTMATTER = re.compile(r"\A---\n.*?\n---\n", re.DOTALL)
_HEADING = re.compile(r"^#{1,6} ", re.MULTILINE)
# C0 control chars except tab/newline/carriage-return. Extracted PDF text
# carries NUL (0x00) and form-feeds that Postgres text columns reject
# ("cannot contain NUL bytes") and that add no semantic value to embeddings.
_CONTROL = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f]")


_HEADING_LINE = re.compile(r"^#{1,6}\s+(.*)$")


def _sections(text: str) -> list[tuple[str, str]]:
    """Split at markdown headings → (heading text, section incl. heading)."""
    starts = [m.start() for m in _HEADING.finditer(text)]
    if not starts:
        return [("", text)]
    bounds = ([0] if starts[0] != 0 else []) + starts + [len(text)]
    out = []
    for a, b in zip(bounds, bounds[1:]):
        section = text[a:b]
        m = _HEADING_LINE.match(section.split("\n", 1)[0])
        out.append((m.group(1).strip() if m else "", section))
    return out


def _hard_wrap(para: str, target: int, overlap: int) -> list[str]:
    """Slice an oversized paragraph into <= target windows, each window's
    leading `overlap` chars equal to the previous window's trailing
    `overlap` chars."""
    n = len(para)
    pieces: list[str] = []
    start = 0
    while True:
        end = min(start + target, n)
        pieces.append(para[start:end])
        if end == n:
            break
        start = end - overlap
    return pieces


def _pack(section: str, target: int, overlap: int) -> list[str]:
    """Greedily pack paragraphs; hard-wrap oversized ones with overlap."""
    paras = [p for p in re.split(r"\n\n+", section) if p.strip()]
    pieces: list[str] = []
    buf = ""
    for para in paras:
        if buf and len(buf) + len(para) + 2 > target:
            pieces.append(buf)
            buf = ""
        if len(para) > target:
            pieces.extend(_hard_wrap(para, target, overlap))
            continue
        buf = f"{buf}\n\n{para}" if buf else para
    if buf:
        pieces.append(buf)
    return pieces


def chunk_doc(
    doc: Doc,
    *,
    target_chars: int = 2000,
    overlap_chars: int = 200,
    code_index: Mapping[str, tuple[str, ...]] | None = None,
) -> list[Chunk]:
    """Chunk *doc* into <= target_chars windows with overlap between them.

    Rule docs (``doc.paragraphs`` set) pack whole FR paragraphs instead
    and carry ``p_id``/``p_id_last``/``page``/``ordinal`` metadata.
    Every chunk carries ``section`` — the markdown heading it sits under
    (``""`` when none), which is how comment chunks know their attachment.

    *code_index* is ``pfs.anchors.code_family_index(families)`` — pass
    the same one the restamp backfill builds (hand families plus
    ``pfs.families.refresh_from`` against the DuckDB replica) so a chunk
    stamped at index time and the same chunk restamped later carry
    identical ``families`` metadata (Ruling F1). ``None`` (the default)
    falls back to a live scan of the import-time ``pfs.families.FAMILIES``
    registry — only hand families, no CPT-derived ones.

    Raises ``ValueError`` when ``overlap_chars >= target_chars`` (an
    overlap that large or larger would never let the window advance).
    """
    if overlap_chars >= target_chars:
        raise ValueError("overlap_chars must be smaller than target_chars")
    if doc.paragraphs:
        return _chunk_paragraphs(doc, target_chars, overlap_chars, code_index)
    body = _CONTROL.sub("", _FRONTMATTER.sub("", doc.text)).strip()
    if not body:
        return []
    prefix = f"{doc.key}:{content_hash(doc.text)[:12]}"
    pieces = [
        (heading, piece)
        for heading, section in _sections(body)
        for piece in _pack(section, target_chars, overlap_chars)
    ]
    return [
        Chunk(
            id=f"{prefix}:{seq:04d}",
            text=piece,
            metadata={
                **doc.metadata,
                "item_key": doc.key,
                "seq": str(seq),
                "section": heading,
                # The heading is scanned with the piece: a section titled
                # "## G0556 — APCM" is about G0556 all the way down, but
                # only its first chunk repeats the code.
                **anchor_metadata(f"{heading} {piece}", code_index=code_index),
            },
        )
        for seq, (heading, piece) in enumerate(pieces)
    ]


def _chunk_paragraphs(
    doc: Doc,
    target: int,
    overlap: int,
    code_index: Mapping[str, tuple[str, ...]] | None = None,
) -> list[Chunk]:
    """Greedy pack of whole paragraphs; an oversized paragraph is
    hard-wrapped with overlap, every piece keeping its own anchor."""
    packed: list[tuple[str, Paragraph, Paragraph]] = []  # text, first, last
    buf: list[Paragraph] = []

    def flush() -> None:
        if buf:
            packed.append(("\n\n".join(p.text for p in buf), buf[0], buf[-1]))
            buf.clear()

    for para in doc.paragraphs:
        text = _CONTROL.sub("", para.text).strip()
        if not text:
            continue
        para = Paragraph(para.p_id, para.page, para.ordinal, text)
        if len(text) > target:
            flush()
            for piece in _hard_wrap(text, target, overlap):
                packed.append((piece, para, para))
            continue
        if buf and sum(len(p.text) + 2 for p in buf) + len(text) > target:
            flush()
        buf.append(para)
    flush()
    prefix = f"{doc.key}:{content_hash(doc.text)[:12]}"
    return [
        Chunk(
            id=f"{prefix}:{seq:04d}",
            text=text,
            metadata={
                **doc.metadata,
                "item_key": doc.key,
                "seq": str(seq),
                "section": "",
                **anchor_metadata(text, code_index=code_index),
                "p_id": str(first.p_id),
                "p_id_last": str(last.p_id),
                "page": str(first.page),
                "ordinal": str(first.ordinal),
            },
        )
        for seq, (text, first, last) in enumerate(packed)
    ]