129 lines
4.4 KiB
Python
129 lines
4.4 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
|
|
from app.rag.extractors import ExtractedBlock
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ChunkMetadata:
|
|
content: str
|
|
page_start: int | None = None
|
|
page_end: int | None = None
|
|
section_title: str = ""
|
|
section_path: str = ""
|
|
paragraph_index: int | None = None
|
|
source_ref: str = ""
|
|
|
|
|
|
def chunk_text(text: str, *, chunk_size: int = 800, overlap: int = 120) -> list[str]:
|
|
cleaned = (text or "").strip()
|
|
if not cleaned:
|
|
return []
|
|
if len(cleaned) <= chunk_size:
|
|
return [cleaned]
|
|
chunks: list[str] = []
|
|
start = 0
|
|
while start < len(cleaned):
|
|
end = min(len(cleaned), start + chunk_size)
|
|
piece = cleaned[start:end].strip()
|
|
if piece:
|
|
chunks.append(piece)
|
|
if end >= len(cleaned):
|
|
break
|
|
start = max(0, end - overlap)
|
|
return chunks
|
|
|
|
|
|
def chunk_blocks(
|
|
blocks: list[ExtractedBlock],
|
|
*,
|
|
chunk_size: int = 1400,
|
|
overlap_blocks: int = 1,
|
|
) -> list[ChunkMetadata]:
|
|
normalized = [block for block in blocks if block.text.strip()]
|
|
if not normalized:
|
|
return []
|
|
|
|
chunks: list[ChunkMetadata] = []
|
|
current: list[ExtractedBlock] = []
|
|
current_len = 0
|
|
|
|
for block in normalized:
|
|
block_len = len(block.text)
|
|
if current and current_len + block_len + 2 > chunk_size:
|
|
chunks.append(_build_chunk(current))
|
|
current = current[-overlap_blocks:] if overlap_blocks > 0 else []
|
|
current_len = sum(len(item.text) + 2 for item in current)
|
|
|
|
if block_len > chunk_size:
|
|
if current:
|
|
chunks.append(_build_chunk(current))
|
|
current = []
|
|
current_len = 0
|
|
chunks.extend(_split_large_block(block, chunk_size=chunk_size))
|
|
continue
|
|
|
|
current.append(block)
|
|
current_len += block_len + 2
|
|
|
|
if current:
|
|
chunks.append(_build_chunk(current))
|
|
|
|
return chunks
|
|
|
|
|
|
def _build_chunk(blocks: list[ExtractedBlock]) -> ChunkMetadata:
|
|
page_values = [block.page_start for block in blocks if block.page_start is not None]
|
|
page_end_values = [block.page_end or block.page_start for block in blocks if block.page_start is not None]
|
|
section_path = next((block.section_path for block in reversed(blocks) if block.section_path), "")
|
|
section_title = next((block.section_title for block in reversed(blocks) if block.section_title), "")
|
|
paragraph_values = [block.paragraph_index for block in blocks if block.paragraph_index is not None]
|
|
page_start = min(page_values) if page_values else None
|
|
page_end = max(page_end_values) if page_end_values else None
|
|
source_ref = _format_source_ref(page_start, page_end, section_path or section_title)
|
|
|
|
return ChunkMetadata(
|
|
content="\n\n".join(block.text.strip() for block in blocks if block.text.strip()),
|
|
page_start=page_start,
|
|
page_end=page_end,
|
|
section_title=section_title,
|
|
section_path=section_path,
|
|
paragraph_index=min(paragraph_values) if paragraph_values else None,
|
|
source_ref=source_ref,
|
|
)
|
|
|
|
|
|
def _split_large_block(block: ExtractedBlock, *, chunk_size: int) -> list[ChunkMetadata]:
|
|
pieces = chunk_text(block.text, chunk_size=chunk_size, overlap=120)
|
|
out: list[ChunkMetadata] = []
|
|
for piece in pieces:
|
|
out.append(
|
|
ChunkMetadata(
|
|
content=piece,
|
|
page_start=block.page_start,
|
|
page_end=block.page_end,
|
|
section_title=block.section_title,
|
|
section_path=block.section_path,
|
|
paragraph_index=block.paragraph_index,
|
|
source_ref=_format_source_ref(
|
|
block.page_start,
|
|
block.page_end,
|
|
block.section_path or block.section_title,
|
|
),
|
|
)
|
|
)
|
|
return out
|
|
|
|
|
|
def _format_source_ref(page_start: int | None, page_end: int | None, section_path: str) -> str:
|
|
refs: list[str] = []
|
|
if page_start is not None:
|
|
if page_end and page_end != page_start:
|
|
refs.append(f"стр. {page_start}-{page_end}")
|
|
else:
|
|
refs.append(f"стр. {page_start}")
|
|
if section_path:
|
|
refs.append(section_path)
|
|
return ", ".join(refs)
|