src/chunker.py

# ============================================================
# FILE: src/chunker.py
# ============================================================
# PURPOSE:
# Split documents into chunks for retrieval.
#
# WHY CHUNKING MATTERS:
# RAG systems retrieve chunks, not whole documents.
#
# Good chunks should:
# - preserve meaning
# - include enough context
# - avoid being too long
# - avoid being too short
# - preserve useful headings when possible
#
# Bad chunking is one of the most common reasons RAG systems fail.
# ============================================================

import hashlib
from dataclasses import dataclass
from typing import List

from src.document_loader import Document

@dataclass
class Chunk:
    """
    Represents one retrievable text chunk.

    id:stable unique ID used by the vector database
    text:chunk content
    source:original document path
    chunk_index: chunk position inside the source document
    character_count: useful for debugging
    """

    id: str
    text: str
    source: str
    chunk_index: int
    character_count: int


def create_stable_chunk_id(source: str, chunk_index: int, text: str) -> str:
    """
    Create a stable unique ID for a chunk.

    Why stable IDs matter:
    - easier updates
    - easier deletes
    - easier debugging
    - easier source tracing
    """

    raw_id = f"{source}|{chunk_index}|{text}"
    return hashlib.md5(raw_id.encode("utf-8")).hexdigest()


def split_long_text_by_characters(
    text: str,
    chunk_size: int,
    chunk_overlap: int,
) -> List[str]:
    """
    Split very long text into overlapping character chunks.

    This is used only when a single paragraph is too large.
    """

    if chunk_overlap >= chunk_size:
        raise ValueError("chunk_overlap must be smaller than chunk_size.")

    chunks = []
    start = 0
    text_length = len(text)

    while start < text_length:
        end = start + chunk_size
        chunk = text[start:end].strip()

        if chunk:
            chunks.append(chunk)

        start = end - chunk_overlap

    return chunks


def chunk_text_by_paragraphs(
    text: str,
    chunk_size: int,
    chunk_overlap: int,
) -> List[str]:
    """
    Paragraph-aware chunking.

    This tries to keep paragraphs together instead of blindly cutting text.

    How it works:
    1. Split text by blank lines.
    2. Add paragraphs into the current chunk until chunk_size is reached.
    3. Start a new chunk when the next paragraph would exceed chunk_size.
    4. If a paragraph itself is too large, split it by characters.
    """

    paragraphs = [paragraph.strip() for paragraph in text.split("\n\n") if paragraph.strip()]
    chunks = []
    current_chunk = ""

    for paragraph in paragraphs:
        if len(paragraph) > chunk_size:
            if current_chunk:
                chunks.append(current_chunk.strip())
                current_chunk = ""

            long_chunks = split_long_text_by_characters(
                text=paragraph,
                chunk_size=chunk_size,
                chunk_overlap=chunk_overlap,
            )

            chunks.extend(long_chunks)
            continue

        candidate_chunk = paragraph if not current_chunk else current_chunk + "\n\n" + paragraph

        if len(candidate_chunk) <= chunk_size:
            current_chunk = candidate_chunk
        else:
            if current_chunk:
                chunks.append(current_chunk.strip())

            current_chunk = paragraph

    if current_chunk:
        chunks.append(current_chunk.strip())

    return chunks


def build_chunks_from_documents(
    documents: List[Document],
    chunk_size: int,
    chunk_overlap: int,
) -> List[Chunk]:
    """
    Convert loaded documents into retrievable chunks.
    """

    chunks = []

    for document in documents:
        text_chunks = chunk_text_by_paragraphs(
            text=document.text,
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
        )

        for chunk_index, chunk_text in enumerate(text_chunks):
            chunk_id = create_stable_chunk_id(
                source=document.source,
                chunk_index=chunk_index,
                text=chunk_text,
            )

            chunks.append(
                Chunk(
                    id=chunk_id,
                    text=chunk_text,
                    source=document.source,
                    chunk_index=chunk_index,
                    character_count=len(chunk_text),
                )
            )

    return chunks