rag/chuncking.py

# chunking.py
import re
from typing import List, Tuple


def approx_token_count(text: str) -> int:
    """
    Rough token estimate for chunk sizing (heuristic).
    Keep this simple but consistent: ~4 chars per token.
    """
    return max(1, len(text) // 4)


def chunk_pages(
    pages: List[str],
    target_tokens: int = 520,
    overlap_tokens: int = 80,
) -> List[Tuple[str, int, int]]:
    """
    Split pages (list[str]) into overlapping chunks.
    Returns list of tuples: (chunk_text, page_start, page_end)

    Overlap is implemented at paragraph level (keeps page ranges correct).
    """
    chunks: List[Tuple[str, int, int]] = []

    buffer: List[str] = []
    buffer_pages: List[int] = []
    buffer_tokens = 0

    def make_chunk():
        nonlocal buffer, buffer_pages, buffer_tokens
        if not buffer:
            return
        chunk_text = "\n\n".join(buffer).strip()
        page_start = min(buffer_pages)
        page_end = max(buffer_pages)
        chunks.append((chunk_text, page_start, page_end))

    for page_idx, page in enumerate(pages, start=1):
        paragraphs = [p.strip() for p in re.split(r"\n\s*\n", page) if p.strip()]

        for para in paragraphs:
            t = approx_token_count(para)

            # If single paragraph exceeds target, create it as its own chunk
            if buffer_tokens == 0 and t > target_tokens:
                # make chunk with this large paragraph alone
                buffer = [para]
                buffer_pages = [page_idx]
                buffer_tokens = t
                make_chunk()
                buffer, buffer_pages, buffer_tokens = [], [], 0
                continue

            # If adding this paragraph would exceed target, flush current chunk
            if buffer_tokens + t > target_tokens:
                make_chunk()

                # prepare overlap: keep tail paragraphs whose tokens sum >= overlap_tokens
                tail_buffer: List[str] = []
                tail_pages: List[int] = []
                tail_tokens = 0
                # iterate buffer in reverse to pick tail paragraphs
                for p, p_pg in zip(reversed(buffer), reversed(buffer_pages)):
                    pt = approx_token_count(p)
                    tail_buffer.insert(0, p)
                    tail_pages.insert(0, p_pg)
                    tail_tokens += pt
                    if tail_tokens >= overlap_tokens:
                        break

                buffer = tail_buffer
                buffer_pages = tail_pages
                buffer_tokens = tail_tokens

            # append current paragraph
            buffer.append(para)
            buffer_pages.append(page_idx)
            buffer_tokens += t

    # final flush
    if buffer:
        make_chunk()

    return chunks