rag-from-scratch/chunks/packer.py

from __future__ import annotations
import re
from chunks.models import Block, Chunk
from chunks.config import ChunkerConfig

_SENTENCE_RE = re.compile(r"(?<=[.!?»])\s+(?=[A-ZÀÈÉÌÒÙ\"])")
_CHUNK_COUNTER = 0


def _reset_counter() -> None:
    global _CHUNK_COUNTER
    _CHUNK_COUNTER = 0


def _make_chunk_id() -> str:
    global _CHUNK_COUNTER
    _CHUNK_COUNTER += 1
    return f"chk_{_CHUNK_COUNTER:06d}"


def _header_prefix(header_path: list[dict], depth: int) -> str:
    return " > ".join(h["text"] for h in header_path[:depth])


def _build_chunk(blocks: list[Block], index: int, config: ChunkerConfig,
                 content_type: str = "section_fragment") -> Chunk:
    content = "\n\n".join(b.content for b in blocks)
    header_path = blocks[0].header_path
    prefix = _header_prefix(header_path, config.context_depth)
    embedding = f"{prefix}\n\n{content}" if prefix else content
    total_chars = len(content)
    return Chunk(
        chunk_id=_make_chunk_id(),
        chunk_index=index,
        content_original=content,
        content_for_embedding=embedding,
        content_type=content_type,
        chars=total_chars,
        start_line=blocks[0].start_line,
        end_line=blocks[-1].end_line,
        header_path=header_path,
        block_ids=[b.id for b in blocks],
        flags={
            "has_code":    any(b.kind == "code"  for b in blocks),
            "has_table":   any(b.kind == "table" for b in blocks),
            "has_math":    any(b.kind == "math"  for b in blocks),
            "is_overflow": total_chars > config.max_chars,
            "is_sparse":   total_chars < config.min_chars,
        },
        neighbors={"previous_chunk_id": None, "next_chunk_id": None},
        assets=[],
    )


def _split_paragraph(block: Block, config: ChunkerConfig) -> list[Block]:
    sentences = _SENTENCE_RE.split(block.content)
    sub_blocks: list[Block] = []
    accumulated = ""

    for sent in sentences:
        candidate = (accumulated + " " + sent).strip() if accumulated else sent
        if len(candidate) > config.max_chars and accumulated:
            text = accumulated.strip()
            sub_blocks.append(Block(
                id=f"{block.id}_s{len(sub_blocks) + 1}", kind="paragraph",
                content=text, plain_text=text, atomic=False,
                start_line=block.start_line, end_line=block.end_line,
                header_path=block.header_path, chars=len(text),
            ))
            accumulated = sent
        else:
            accumulated = candidate

    if accumulated:
        text = accumulated.strip()
        sub_blocks.append(Block(
            id=f"{block.id}_s{len(sub_blocks) + 1}", kind="paragraph",
            content=text, plain_text=text, atomic=False,
            start_line=block.start_line, end_line=block.end_line,
            header_path=block.header_path, chars=len(text),
        ))
    return sub_blocks if sub_blocks else [block]


def pack(blocks: list[Block], config: ChunkerConfig, stem: str) -> list[Chunk]:
    _reset_counter()
    chunks: list[Chunk] = []
    chunk_index = 0

    # Espandi paragrafi sovradimensionati
    expanded: list[Block] = []
    for b in blocks:
        if not b.atomic and b.kind == "paragraph" and b.chars > config.max_chars:
            expanded.extend(_split_paragraph(b, config))
        else:
            expanded.append(b)

    # Raggruppa per header_path
    groups: list[list[Block]] = []
    cur_group: list[Block] = []
    cur_key: str | None = None
    for b in expanded:
        key = str(b.header_path)
        if key != cur_key:
            if cur_group:
                groups.append(cur_group)
            cur_group = [b]
            cur_key = key
        else:
            cur_group.append(b)
    if cur_group:
        groups.append(cur_group)

    for group in groups:
        accumulated: list[Block] = []
        accumulated_chars = 0

        def flush() -> None:
            nonlocal accumulated, accumulated_chars, chunk_index
            if not accumulated:
                return
            chunks.append(_build_chunk(accumulated, chunk_index, config))
            chunk_index += 1
            accumulated = []
            accumulated_chars = 0

        for b in group:
            if b.kind == "thematic_break":
                flush()
                continue

            if b.atomic:
                if accumulated_chars + b.chars <= config.max_chars:
                    accumulated.append(b)
                    accumulated_chars += b.chars
                else:
                    flush()
                    ctype = "overflow" if b.chars > config.max_chars else "atomic_block"
                    chunks.append(_build_chunk([b], chunk_index, config, content_type=ctype))
                    chunk_index += 1
            else:
                # Flush preventivo se aggiungere questo blocco supererebbe max_chars
                # oppure supererebbe il target (con abbastanza contenuto già accumulato)
                if accumulated and accumulated_chars >= config.min_chars:
                    if (accumulated_chars + b.chars > config.max_chars
                            or accumulated_chars + b.chars > config.target_chars):
                        flush()
                accumulated.append(b)
                accumulated_chars += b.chars

        # Flush residuo — merge con precedente se troppo piccolo
        if accumulated:
            if (accumulated_chars < config.min_chars and chunks
                    and chunks[-1].header_path == accumulated[0].header_path):
                prev = chunks[-1]
                merged = prev.content_original + "\n\n" + "\n\n".join(b.content for b in accumulated)
                prefix = _header_prefix(prev.header_path, config.context_depth)
                prev.content_original = merged
                prev.content_for_embedding = f"{prefix}\n\n{merged}" if prefix else merged
                prev.chars = len(merged)
                prev.end_line = accumulated[-1].end_line
                prev.block_ids.extend(b.id for b in accumulated)
                prev.flags["is_sparse"] = prev.chars < config.min_chars
            else:
                flush()

    # Popola neighbors
    for idx, chunk in enumerate(chunks):
        chunk.neighbors["previous_chunk_id"] = chunks[idx - 1].chunk_id if idx > 0 else None
        chunk.neighbors["next_chunk_id"] = chunks[idx + 1].chunk_id if idx < len(chunks) - 1 else None

    return chunks