app/services/chunker.py

"""Author RAG Chatbot SaaS — Semantic Text Chunker.

Splits document text into overlapping chunks for vector embedding.
Uses sentence-boundary-aware splitting for clean, meaningful chunks.
Config: CHUNK_SIZE=512 tokens, CHUNK_OVERLAP=64 tokens.
"""

from dataclasses import dataclass

import structlog

from app.config import get_settings
from app.utils.token_counter import count_tokens

logger = structlog.get_logger(__name__)
cfg = get_settings()


@dataclass
class TextChunk:
    """A single chunk of text ready for embedding."""

    text: str          # Chunk content
    chunk_index: int   # Position in document (0-indexed)
    char_start: int    # Character start position in original text
    char_end: int      # Character end position in original text
    token_count: int   # Token count for this chunk


def chunk_document(
    text: str,
    chunk_size: int | None = None,
    overlap: int | None = None,
) -> list[TextChunk]:
    """Split document text into overlapping semantic chunks.

    Splits at sentence boundaries when possible to preserve meaning.
    Falls back to character-based splitting if needed.

    Args:
        text: Full document text.
        chunk_size: Max tokens per chunk (defaults to config CHUNK_SIZE).
        overlap: Overlap tokens between consecutive chunks (defaults to config CHUNK_OVERLAP).

    Returns:
        List of TextChunk objects ready for embedding.
    """
    chunk_size = chunk_size or cfg.CHUNK_SIZE
    overlap = overlap or cfg.CHUNK_OVERLAP

    if not text.strip():
        logger.warning("Empty text passed to chunker")
        return []

    sentences = _split_into_sentences(text)
    chunks = _build_chunks(sentences, chunk_size, overlap, text)

    logger.info("Document chunked", total_chunks=len(chunks), chunk_size=chunk_size, overlap=overlap)
    return chunks


def _split_into_sentences(text: str) -> list[str]:
    """Split text into sentences using punctuation-based heuristics.

    Args:
        text: Input text.

    Returns:
        List of sentence strings.
    """
    import re
    # Split on period/exclamation/question mark followed by space+capital or newline
    sentences = re.split(r"(?<=[.!?])\s+(?=[A-Z\"\'])|(?<=\n)\n", text)
    return [s.strip() for s in sentences if s.strip()]


def _build_chunks(
    sentences: list[str],
    chunk_size: int,
    overlap: int,
    original_text: str,
) -> list[TextChunk]:
    """Aggregate sentences into token-bounded chunks with overlap.

    Args:
        sentences: List of sentences from the document.
        chunk_size: Max tokens per chunk.
        overlap: Target overlap tokens between chunks.
        original_text: Original full text (for char offset calculation).

    Returns:
        List of TextChunk objects.
    """
    chunks: list[TextChunk] = []
    current_sentences: list[str] = []
    current_tokens = 0
    overlap_buffer: list[str] = []
    char_cursor = 0

    for sentence in sentences:
        sentence_tokens = count_tokens(sentence)

        # If adding this sentence exceeds chunk_size, finalize current chunk
        if current_tokens + sentence_tokens > chunk_size and current_sentences:
            chunk_text = " ".join(current_sentences)
            char_start = original_text.find(current_sentences[0], char_cursor)
            char_end = char_start + len(chunk_text)

            chunks.append(TextChunk(
                text=chunk_text,
                chunk_index=len(chunks),
                char_start=max(char_start, 0),
                char_end=char_end,
                token_count=current_tokens,
            ))

            # Build overlap buffer from end of current chunk
            overlap_buffer = _build_overlap_buffer(current_sentences, overlap)
            overlap_tokens = sum(count_tokens(s) for s in overlap_buffer)
            current_sentences = overlap_buffer.copy()
            current_tokens = overlap_tokens
            char_cursor = char_start

        current_sentences.append(sentence)
        current_tokens += sentence_tokens

    # Finalize last chunk
    if current_sentences:
        chunk_text = " ".join(current_sentences)
        char_start = original_text.find(current_sentences[0], char_cursor)
        chunks.append(TextChunk(
            text=chunk_text,
            chunk_index=len(chunks),
            char_start=max(char_start, 0),
            char_end=max(char_start, 0) + len(chunk_text),
            token_count=current_tokens,
        ))

    return chunks


def _build_overlap_buffer(sentences: list[str], overlap_tokens: int) -> list[str]:
    """Select trailing sentences that fit within the overlap token budget.

    Args:
        sentences: Current chunk's sentences.
        overlap_tokens: Target overlap size in tokens.

    Returns:
        List of sentences to carry into the next chunk.
    """
    buffer: list[str] = []
    token_count = 0
    for sentence in reversed(sentences):
        sentence_tokens = count_tokens(sentence)
        if token_count + sentence_tokens > overlap_tokens:
            break
        buffer.insert(0, sentence)
        token_count += sentence_tokens
    return buffer