chroma/chunking.py

"""Split agenda text into overlapping chunks for embedding/summarization.

A small, dependency-free recursive splitter: it prefers to break on paragraph
then line then sentence then word boundaries, packing pieces up to ``chunk_size``
characters with ``overlap`` characters carried from the end of one chunk into the
start of the next (so context isn't lost across a boundary).
"""

from __future__ import annotations

import re

# Boundary separators tried in order, coarsest first.
_SEPARATORS = ["\n\n", "\n", ". ", " "]


def _split_keep(text: str, sep: str) -> list[str]:
    """Split on ``sep`` but keep the separator attached to each preceding piece."""
    if sep == "":
        return list(text)
    parts = text.split(sep)
    out = [p + sep for p in parts[:-1]]
    if parts[-1]:
        out.append(parts[-1])
    return out


def _pieces(text: str) -> list[str]:
    """Break text into the smallest reasonable units we will pack into chunks."""
    units = [text]
    for sep in _SEPARATORS:
        nxt: list[str] = []
        for u in units:
            nxt.extend(_split_keep(u, sep) if len(u) > 1 else [u])
        units = nxt
    return [u for u in units if u]


def chunk_text(
    text: str, *, chunk_size: int = 1500, overlap: int = 200
) -> list[str]:
    """Split ``text`` into overlapping chunks of about ``chunk_size`` characters.

    Parameters
    ----------
    chunk_size:
        Target maximum characters per chunk.
    overlap:
        Characters of trailing context repeated at the start of the next chunk.
    """
    text = (text or "").strip()
    if not text:
        return []
    if overlap >= chunk_size:
        raise ValueError("overlap must be smaller than chunk_size")
    if len(text) <= chunk_size:
        return [text]

    pieces = _pieces(text)
    chunks: list[str] = []
    cur = ""
    for piece in pieces:
        # A single piece longer than chunk_size: hard-split it.
        if len(piece) > chunk_size:
            if cur:
                chunks.append(cur)
                cur = ""
            for i in range(0, len(piece), chunk_size - overlap):
                chunks.append(piece[i : i + chunk_size])
            continue
        if len(cur) + len(piece) <= chunk_size:
            cur += piece
        else:
            chunks.append(cur)
            tail = cur[-overlap:] if overlap else ""
            cur = tail + piece
    if cur.strip():
        chunks.append(cur)

    # Normalize whitespace edges; drop empties.
    return [re.sub(r"[ \t]+\n", "\n", c).strip() for c in chunks if c.strip()]