src/chunking.py

"""
chunking.py — Text splitting with configurable size and overlap.

CHUNKING TRADEOFFS:

Too SMALL (< 100 tokens):
  - Embeddings lose sentence context → poor retrieval
  - More vectors in FAISS → slower at scale
  - Single sentences lack enough context to answer questions

Too LARGE (> 800 tokens):
  - One vector covers too many topics → diluted signal
  - Retrieved chunk floods prompt with irrelevant text
  - Increases token cost per query

OVERLAP purpose:
  - Ensures content at chunk boundaries appears in at least one complete chunk
  - Without overlap, answers straddling two chunks are missed
  - Too much overlap (>20%) creates near-duplicate vectors
"""

import logging
from src.utils import count_tokens_estimate

logger = logging.getLogger("enterprise-rag.chunking")

DEFAULT_CHUNK_SIZE = 512
DEFAULT_CHUNK_OVERLAP = 64
CHARS_PER_TOKEN = 4


def chunk_text(
    text: str,
    chunk_size: int = DEFAULT_CHUNK_SIZE,
    chunk_overlap: int = DEFAULT_CHUNK_OVERLAP,
    min_chunk_chars: int = 80,
) -> list:
    """
    Split text into overlapping chunks.

    Returns list of dicts:
        text        — chunk content
        chunk_index — position in document
        token_count — estimated token count
        char_start  — start offset in original text
        char_end    — end offset in original text
    """
    if not text or not text.strip():
        return []

    char_size = chunk_size * CHARS_PER_TOKEN
    char_overlap = chunk_overlap * CHARS_PER_TOKEN

    chunks = []
    start = 0
    idx = 0
    n = len(text)

    while start < n:
        end = min(start + char_size, n)

        # Try to end at a natural boundary: paragraph > sentence > word
        if end < n:
            para = text.rfind("\n\n", start, end)
            sent = max(
                text.rfind(". ", start, end),
                text.rfind(".\n", start, end),
                text.rfind("! ", start, end),
                text.rfind("? ", start, end),
            )
            if para > start + char_size // 2:
                end = para + 2
            elif sent > start + char_size // 2:
                end = sent + 2

        content = text[start:end].strip()

        if len(content) >= min_chunk_chars:
            chunks.append({
                "text": content,
                "chunk_index": idx,
                "token_count": count_tokens_estimate(content),
                "char_start": start,
                "char_end": end,
            })
            idx += 1

        start = end - char_overlap
        if start >= n:
            break

    logger.info(f"Created {len(chunks)} chunks from {n} chars")
    return chunks


def chunk_statistics(chunks: list) -> dict:
    """Summary stats for the metrics panel."""
    if not chunks:
        return {"count": 0, "avg_tokens": 0, "min_tokens": 0, "max_tokens": 0, "total_tokens": 0}
    counts = [c["token_count"] for c in chunks]
    return {
        "count": len(chunks),
        "avg_tokens": round(sum(counts) / len(counts), 1),
        "min_tokens": min(counts),
        "max_tokens": max(counts),
        "total_tokens": sum(counts),
    }