chunking.py

"""Split cleaned documents into overlapping text chunks."""

from config import CHUNK_SIZE, CHUNK_OVERLAP


def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[str]:
    """Slide a window over text to produce overlapping chunks."""
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunks.append(text[start:end].strip())
        start += chunk_size - overlap
    return [c for c in chunks if len(c) > 30]  # drop tiny trailing chunks


def chunk_documents(documents: list[dict]) -> list[dict]:
    """Chunk every document and tag each chunk with its source."""
    all_chunks = []
    for doc in documents:
        text_chunks = chunk_text(doc["text"])
        for i, chunk in enumerate(text_chunks):
            all_chunks.append({
                "source": doc["source"],
                "chunk_id": f"{doc['source']}__chunk_{i}",
                "text": chunk,
            })
    print(f"[chunking] Produced {len(all_chunks)} chunks.")
    return all_chunks