03_chunking.py

"""
03_chunking.py
--------------
Stage 3 of the RAG pipeline: CHUNKING.

Splits cleaned documents into smaller overlapping text chunks so that each
chunk is small enough to embed meaningfully and retrieve precisely.

Strategy: fixed-size character chunking with overlap. This is simple,
dependency-free, and works well for a student-scale RAG project.

Run directly to see a demo of this stage:
    python 03_chunking.py
"""

from __future__ import annotations

import importlib.util
import sys
import os
from dataclasses import dataclass, field
from typing import List

_THIS_DIR = os.path.dirname(os.path.abspath(__file__))

# Default chunking parameters. Kept as module-level constants so other
# stages (and streamlit_app.py) can reference or override them easily.
CHUNK_SIZE = 800
CHUNK_OVERLAP = 120


def _import_module_from_file(module_name: str, filename: str):
    path = os.path.join(_THIS_DIR, filename)
    spec = importlib.util.spec_from_file_location(module_name, path)
    module = importlib.util.module_from_spec(spec)
    sys.modules[module_name] = module  # required so dataclasses can resolve cls.__module__
    spec.loader.exec_module(module)
    return module


_preprocessing_module = _import_module_from_file("preprocessing_stage", "02_preprocessing.py")
_documents_module = _import_module_from_file("documents_stage", "01_documents.py")
CleanDocument = _preprocessing_module.CleanDocument


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

    chunk_id: str
    doc_id: str
    source: str
    text: str
    metadata: dict = field(default_factory=dict)


def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> List[str]:
    """
    Split `text` into overlapping chunks of at most `chunk_size` characters.

    Splitting is done on whitespace boundaries where possible so words are
    not cut in half.
    """
    if chunk_size <= 0:
        raise ValueError("chunk_size must be a positive integer")
    if overlap < 0 or overlap >= chunk_size:
        raise ValueError("overlap must be non-negative and smaller than chunk_size")

    text = text.strip()
    if not text:
        return []

    chunks: List[str] = []
    start = 0
    text_length = len(text)

    while start < text_length:
        end = min(start + chunk_size, text_length)

        # Try to break on a whitespace boundary instead of mid-word,
        # as long as we're not at the very end of the text.
        if end < text_length:
            last_space = text.rfind(" ", start, end)
            if last_space > start:
                end = last_space

        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)

        if end >= text_length:
            break

        start = max(end - overlap, start + 1)

    return chunks


def chunk_documents(
    documents: List[CleanDocument],
    chunk_size: int = CHUNK_SIZE,
    overlap: int = CHUNK_OVERLAP,
) -> List[Chunk]:
    """Chunk every document in `documents` and return a flat list of Chunks."""
    all_chunks: List[Chunk] = []

    for doc in documents:
        pieces = chunk_text(doc.text, chunk_size=chunk_size, overlap=overlap)
        for i, piece in enumerate(pieces):
            chunk_id = f"{doc.doc_id}::chunk_{i}"
            all_chunks.append(
                Chunk(
                    chunk_id=chunk_id,
                    doc_id=doc.doc_id,
                    source=doc.source,
                    text=piece,
                    metadata={**doc.metadata, "chunk_index": i},
                )
            )

    return all_chunks


if __name__ == "__main__":
    raw_docs = _documents_module.load_documents()
    cleaned_docs = _preprocessing_module.preprocess_documents(raw_docs)
    chunks = chunk_documents(cleaned_docs)
    print(f"Created {len(chunks)} chunk(s) from {len(cleaned_docs)} document(s):")
    for chunk in chunks[:10]:
        preview = chunk.text[:100].replace("\n", " ")
        print(f"  - {chunk.chunk_id} ({len(chunk.text)} chars): {preview}...")
    if len(chunks) > 10:
        print(f"  ... and {len(chunks) - 10} more chunk(s)")