From d2b8baa8446641ea13b3cf6d92695e1089263f88 Mon Sep 17 00:00:00 2001 From: Davide Grilli Date: Mon, 8 Jun 2026 16:29:24 +0200 Subject: [PATCH] =?UTF-8?q?feat(chunks):=20aggiunge=20packer.py=20?= =?UTF-8?q?=E2=80=94=20Block[]=20=E2=86=92=20Chunk[]=20con=20packing=20min?= =?UTF-8?q?/target/max?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Sonnet 4.6 --- chunks/packer.py | 172 ++++++++++++++++++++++++++++++++++++ tests/chunks/test_packer.py | 118 +++++++++++++++++++++++++ 2 files changed, 290 insertions(+) create mode 100644 chunks/packer.py create mode 100644 tests/chunks/test_packer.py diff --git a/chunks/packer.py b/chunks/packer.py new file mode 100644 index 0000000..0537657 --- /dev/null +++ b/chunks/packer.py @@ -0,0 +1,172 @@ +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 diff --git a/tests/chunks/test_packer.py b/tests/chunks/test_packer.py new file mode 100644 index 0000000..7dfd2e3 --- /dev/null +++ b/tests/chunks/test_packer.py @@ -0,0 +1,118 @@ +import pytest +from chunks.models import Block +from chunks.config import ChunkerConfig +from chunks.packer import pack + + +def make_block(idx: int, content: str, kind: str = "paragraph", atomic: bool = False, + header_path=None) -> Block: + hp = header_path or [{"level": 1, "text": "Titolo"}] + return Block( + id=f"blk_{idx:04d}", kind=kind, content=content, plain_text=content, + atomic=atomic, start_line=idx * 2, end_line=idx * 2 + 1, + header_path=hp, chars=len(content), + ) + + +@pytest.fixture +def cfg(): + return ChunkerConfig(min_chars=10, target_chars=50, max_chars=100) + + +def test_single_block_becomes_single_chunk(cfg): + blocks = [make_block(1, "Testo breve.")] + chunks = pack(blocks, cfg, "test") + assert len(chunks) == 1 + assert chunks[0].content_original == "Testo breve." + + +def test_chunk_id_format(cfg): + blocks = [make_block(1, "Testo.")] + chunks = pack(blocks, cfg, "test") + assert chunks[0].chunk_id.startswith("chk_") + + +def test_neighbors_populated(cfg): + blocks = [make_block(i, "x" * 60, header_path=[{"level": 1, "text": "T"}]) for i in range(1, 4)] + chunks = pack(blocks, cfg, "test") + assert len(chunks) >= 2 + assert chunks[0].neighbors["next_chunk_id"] == chunks[1].chunk_id + assert chunks[1].neighbors["previous_chunk_id"] == chunks[0].chunk_id + assert chunks[-1].neighbors["next_chunk_id"] is None + + +def test_blocks_merged_until_target(cfg): + blocks = [make_block(i, "x" * 20) for i in range(1, 4)] + chunks = pack(blocks, cfg, "test") + assert len(chunks) >= 2 + + +def test_oversized_paragraph_split(cfg): + long_content = ("Questa è la prima frase completa. " * 5).strip() + blocks = [make_block(1, long_content)] + chunks = pack(blocks, cfg, "test") + for c in chunks: + assert c.chars <= cfg.max_chars or c.flags["is_overflow"] + + +def test_atomic_block_not_split(cfg): + atomic_content = "```python\n" + "codice\n" * 20 + "```" + blocks = [make_block(1, atomic_content, kind="code", atomic=True)] + chunks = pack(blocks, cfg, "test") + assert len(chunks) == 1 + assert chunks[0].flags["is_overflow"] is True + assert chunks[0].content_type == "overflow" + + +def test_heading_path_break_creates_new_chunk(cfg): + hp1 = [{"level": 1, "text": "Sezione 1"}] + hp2 = [{"level": 1, "text": "Sezione 2"}] + blocks = [make_block(1, "x" * 30, header_path=hp1), make_block(2, "x" * 30, header_path=hp2)] + chunks = pack(blocks, cfg, "test") + assert len(chunks) == 2 + assert chunks[0].header_path == hp1 + assert chunks[1].header_path == hp2 + + +def test_content_for_embedding_has_header_prefix(cfg): + blocks = [make_block(1, "Contenuto.", header_path=[{"level": 1, "text": "Guida"}, + {"level": 2, "text": "Intro"}])] + chunks = pack(blocks, cfg, "test") + assert "Guida" in chunks[0].content_for_embedding + assert "Intro" in chunks[0].content_for_embedding + assert "Contenuto." in chunks[0].content_for_embedding + + +def test_thematic_break_flushes_chunk(cfg): + blocks = [ + make_block(1, "x" * 30), + make_block(2, "---", kind="thematic_break"), + make_block(3, "x" * 30), + ] + chunks = pack(blocks, cfg, "test") + assert len(chunks) == 2 + + +def test_flags_has_code(cfg): + blocks = [make_block(1, "```\ncode\n```", kind="code", atomic=True)] + chunks = pack(blocks, cfg, "test") + assert chunks[0].flags["has_code"] is True + + +def test_flags_has_table(cfg): + blocks = [make_block(1, "| A | B |\n|---|---|\n| 1 | 2 |", kind="table", atomic=True)] + chunks = pack(blocks, cfg, "test") + assert chunks[0].flags["has_table"] is True + + +def test_chunk_index_sequential(cfg): + blocks = [make_block(i, "x" * 60) for i in range(1, 5)] + chunks = pack(blocks, cfg, "test") + for i, c in enumerate(chunks): + assert c.chunk_index == i + + +def test_assets_empty(cfg): + blocks = [make_block(1, "Testo.")] + chunks = pack(blocks, cfg, "test") + assert chunks[0].assets == []