chunks.py

import re
import uuid

def chunk_text(
    text: str,
    max_chunk_size: int = 500,
    overlap: int = 50
) -> list[str]:
    """
    Splits text into overlapping chunks that end at sentence or paragraph
    boundaries, never in the middle of a word.

    Args:
        text:           The input text to chunk.
        max_chunk_size: Maximum number of characters per chunk (default 500).
        overlap:        Number of characters to overlap between chunks (default 50).

    Returns:
        A list of text chunks.
    """
    if not text or not text.strip():
        return []

    if len(text) <= max_chunk_size:
        return [text.strip()]

    # Regex pattern that matches the end of a sentence or a paragraph.
    # Priority (highest → lowest):
    #   1. Paragraph break (one or more blank lines)
    #   2. Sentence-ending punctuation followed by whitespace or end-of-string
    #   3. Any whitespace (word boundary — last resort)
    boundary_pattern = re.compile(
        r'(\n\s*\n'               # paragraph break
        r'|[.!?…]+["\']?\s+'     # sentence-ending punctuation + space
        r'|\s+)',                 # any whitespace (word boundary)
        re.MULTILINE
    )

    def find_best_boundary(text_slice: str, ideal_end: int) -> int:
        """
        Starting from `ideal_end`, scan backwards for the nearest sentence /
        paragraph boundary. Falls back to the nearest word boundary, and as a
        last resort returns `ideal_end` unchanged (no mid-word split is still
        possible when the window contains no spaces at all).
        """
        search_region = text_slice[:ideal_end]

        # Collect all boundary positions within the slice
        boundaries = []
        for match in boundary_pattern.finditer(search_region):
            end_pos = match.end()
            # Classify the boundary type for prioritisation
            matched = match.group()
            if re.match(r'\n\s*\n', matched):
                priority = 0   # paragraph — best
            elif re.match(r'[.!?…]', matched):
                priority = 1   # sentence — good
            else:
                priority = 2   # word boundary — fallback
            boundaries.append((end_pos, priority))

        if not boundaries:
            return ideal_end  # no spaces found; return as-is

        # Prefer sentence/paragraph boundaries closest to (but not past) ideal_end
        for desired_priority in (0, 1, 2):
            candidates = [
                pos for pos, pri in boundaries if pri == desired_priority
            ]
            if candidates:
                return candidates[-1]   # latest boundary at that priority level

        return ideal_end

    chunks: list[str] = []
    start = 0
    text_len = len(text)

    while start < text_len:
        end = start + max_chunk_size

        if end >= text_len:
            # Last chunk — take everything that remains
            chunk = text[start:].strip()
            if chunk:
                chunks.append(chunk)
            break

        # Find the best boundary to cut at
        best_end = find_best_boundary(text, end)

        # Safety: if best_end didn't move far enough, advance at least one word
        if best_end <= start:
            next_space = text.find(' ', end)
            best_end = next_space + 1 if next_space != -1 else text_len

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

        # Next chunk starts `overlap` characters before the cut point
        next_start = best_end - overlap

        if next_start <= start:
            # Safety: advance to the next word boundary, not just +1
            next_space = text.find(' ', start + 1)
            next_start = next_space + 1 if next_space != -1 else text_len

        start = next_start

    return chunks


def build_chunks(documents):
    chunks = []
    ids = []
    metadata = []

    for doc in documents:
        chunks_ = [c for c in chunk_text(doc["text"], max_chunk_size=500, overlap=50) if len(c) >= 100]
        chunks.extend(chunks_)
        ids_ = [str(uuid.uuid4()) for _ in range(len(chunks_))]
        ids.extend(ids_)
        metadata.extend({"source": doc["source"], "chunk_index": i} for i in range(len(chunks_)))

    print(f"Total chunks produced: {len(chunks)}")
    return chunks, ids, metadata