RAG Chunk Comparator — Python source

"""RAG Chunk Comparator — chunk one document three ways (fixed-size,
sentence-aware, markdown-heading-aware) and compare retrieval stats.

Language: Python (3.9+, standard library only)
Source:   CosmoDev polyglot showcase port of the RAG Chunk Comparator
          tool (slug: rag-chunk-comparator).
Port of src/lib/ragChunkComparator.ts (the canonical TypeScript
          implementation).
Tool page: https://dev.cosmolabs.org/tools/rag-chunk-comparator
License:  display source — part of CosmoDev's polyglot tool pages.

Token sizes inline the tokenEstimator prose heuristic (~4 chars per
token, per non-empty line, minimum one token per line) so this module is
self-contained. The stats that matter for retrieval: chunk count, size
spread, and how often boundaries land on sentence ends — mid-sentence
cuts are the classic recall killer.

Field names use snake_case where the TypeScript uses camelCase
(sizeTokens -> size_tokens, sentenceBoundaryShare ->
sentence_boundary_share); the RangeError becomes ValueError.
"""

from __future__ import annotations

import math
import re
from dataclasses import dataclass
from typing import List, Optional

#: A chunking strategy identifier (one of "fixed", "sentence", "markdown").
ChunkStrategy = str


@dataclass
class ChunkOptions:
    """Target chunk size in tokens; overlap in tokens (fixed strategy only)."""

    size_tokens: int
    overlap_tokens: int = 0


@dataclass
class Chunk:
    """One chunk: text, prose-heuristic token count, and optional heading."""

    index: int
    text: str
    tokens: int
    #: Nearest markdown heading for markdown chunks (None otherwise).
    heading: Optional[str] = None


@dataclass
class StrategyStats:
    count: int
    min_tokens: int
    max_tokens: int
    avg_tokens: int
    #: Share of chunk boundaries that fall on a sentence end (0-1).
    sentence_boundary_share: float


@dataclass
class StrategyResult:
    strategy: ChunkStrategy
    chunks: List[Chunk]
    stats: StrategyStats


def _round_half_up(x: float) -> int:
    """JavaScript's Math.round: halves go up, never to even."""
    return int(math.floor(x + 0.5))


def _tok(s: str) -> int:
    """Prose token estimate: chars/4 per non-empty line, min 1 per line."""
    total = 0
    for line in re.split(r"\r?\n", s):
        if line.strip() == "":
            continue
        total += max(1, _round_half_up(len(line) / 4))
    return total


_SENTENCE_SPLIT = re.compile(r"(?<=[.!?]) +")
_ENDS_SENTENCE = re.compile(r"[.!?][\"')\]]?$")
_HEADING = re.compile(r"^(#{1,6})\s+(.*)$")


def split_sentences(text: str) -> List[str]:
    """Split on sentence enders followed by whitespace or end of text."""
    normalized = re.sub(r"\s+", " ", text).strip()
    return [part for part in _SENTENCE_SPLIT.split(normalized) if part]


def _ends_sentence(s: str) -> bool:
    return _ENDS_SENTENCE.search(s.strip()) is not None


def chunk_fixed(text: str, opts: ChunkOptions) -> List[Chunk]:
    """Greedy character accumulation to a token target (overlapping allowed)."""
    size_tokens = opts.size_tokens
    overlap_tokens = opts.overlap_tokens
    if size_tokens <= 0:
        raise ValueError("sizeTokens must be > 0")
    if overlap_tokens < 0 or overlap_tokens >= size_tokens:
        raise ValueError("overlapTokens must be in [0, sizeTokens)")
    clean = text.strip()
    if not clean:
        return []
    # ~4 chars per prose token: step by tokens, verify with the estimator.
    char_step = max(1, _round_half_up(size_tokens * 4))
    overlap_chars = _round_half_up(overlap_tokens * 4)
    chunks: List[Chunk] = []
    start = 0
    while start < len(clean):
        end = min(start + char_step, len(clean))
        # Prefer cutting at whitespace near the target — but never trim the
        # document's final piece back to a word when it already fits.
        if end < len(clean):
            cut = clean.rfind(" ", 0, end + 1)
            if cut > start:
                end = cut
        piece = clean[start:end].strip()
        if piece:
            chunks.append(Chunk(index=len(chunks), text=piece, tokens=_tok(piece)))
        if end >= len(clean):
            break
        start = max(end - overlap_chars, start + 1)
    return chunks


def chunk_by_sentences(text: str, opts: ChunkOptions) -> List[Chunk]:
    """Group whole sentences up to the token target; boundaries never split a sentence."""
    size_tokens = opts.size_tokens
    if size_tokens <= 0:
        raise ValueError("sizeTokens must be > 0")
    sentences = split_sentences(text)
    if not sentences:
        return []
    chunks: List[Chunk] = []
    current: List[str] = []
    current_tokens = 0

    def flush() -> None:
        if not current:
            return
        piece = " ".join(current)
        chunks.append(Chunk(index=len(chunks), text=piece, tokens=_tok(piece)))
        current.clear()

    for sentence in sentences:
        t = _tok(sentence)
        if current_tokens > 0 and current_tokens + t > size_tokens:
            flush()
            current_tokens = 0
        current.append(sentence)
        current_tokens += t
        # A single sentence larger than the target becomes its own chunk.
    flush()
    return chunks


def chunk_markdown(text: str, opts: ChunkOptions) -> List[Chunk]:
    """Split on markdown headings; oversized sections fall back to sentence grouping."""
    size_tokens = opts.size_tokens
    if size_tokens <= 0:
        raise ValueError("sizeTokens must be > 0")
    sections: List[tuple] = []
    current: tuple = (None, [])  # (heading, body lines)
    for line in text.split("\n"):
        m = _HEADING.match(line)
        if m:
            if current[1]:
                sections.append(current)
            current = (m.group(2).strip(), [])
        else:
            current[1].append(line)
    if current[1]:
        sections.append(current)

    chunks: List[Chunk] = []
    for heading, body_lines in sections:
        body = "\n".join(body_lines).strip()
        if not body:
            continue
        whole = f"# {heading}\n{body}" if heading else body
        if _tok(whole) <= size_tokens:
            chunks.append(
                Chunk(index=len(chunks), text=whole, tokens=_tok(whole), heading=heading)
            )
            continue
        # Oversized section: sentence-group the body, stamp every chunk with
        # the heading.
        for c in chunk_by_sentences(body, opts):
            chunks.append(
                Chunk(index=len(chunks), text=c.text, tokens=c.tokens, heading=heading)
            )
    return chunks


def _stats_for(strategy: ChunkStrategy, chunks: List[Chunk]) -> StrategyResult:
    sizes = [c.tokens for c in chunks]
    count = len(chunks)
    min_tokens = min(sizes) if count else 0
    max_tokens = max(sizes) if count else 0
    avg_tokens = _round_half_up(sum(sizes) / count) if count else 0
    boundaries = [_ends_sentence(c.text) for c in chunks[:-1]]
    if boundaries:
        sentence_boundary_share = sum(1 for b in boundaries if b) / len(boundaries)
    else:
        sentence_boundary_share = 1.0  # a single chunk has no internal boundaries to botch
    return StrategyResult(
        strategy=strategy,
        chunks=chunks,
        stats=StrategyStats(
            count=count,
            min_tokens=min_tokens,
            max_tokens=max_tokens,
            avg_tokens=avg_tokens,
            sentence_boundary_share=sentence_boundary_share,
        ),
    )


def compare_strategies(text: str, opts: ChunkOptions) -> dict:
    """Run all three strategies over one document and report comparable stats."""
    return {
        "fixed": _stats_for("fixed", chunk_fixed(text, opts)),
        "sentence": _stats_for("sentence", chunk_by_sentences(text, opts)),
        "markdown": _stats_for("markdown", chunk_markdown(text, opts)),
    }