173 lines
6.3 KiB
Python
173 lines
6.3 KiB
Python
|
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
|