64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
"""Delimiter-based splitting for non-tabular documents."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from rag_cut.models import Block, BlockType, SplitConfig
|
|
from rag_cut.splitters.default_splitter import _split_by_size
|
|
|
|
ATOMIC_TYPES = {BlockType.TABLE, BlockType.IMAGE}
|
|
|
|
|
|
def partition_blocks_by_delimiter(blocks: list[Block], delimiter: str) -> list[list[Block]]:
|
|
"""Split block stream on delimiter; delimiter text is discarded from chunks."""
|
|
if not delimiter:
|
|
return [blocks] if blocks else []
|
|
|
|
groups: list[list[Block]] = []
|
|
current: list[Block] = []
|
|
|
|
def flush() -> None:
|
|
nonlocal current
|
|
if current:
|
|
groups.append(current)
|
|
current = []
|
|
|
|
for block in blocks:
|
|
if block.type in ATOMIC_TYPES:
|
|
current.append(block)
|
|
continue
|
|
|
|
text = block.text or block.markdown
|
|
if delimiter not in text:
|
|
current.append(block)
|
|
continue
|
|
|
|
parts = text.split(delimiter)
|
|
for i, part in enumerate(parts):
|
|
part = part.strip()
|
|
if part:
|
|
piece = Block(type=block.type, text=part, level=block.level, meta=dict(block.meta))
|
|
current.append(piece)
|
|
if i < len(parts) - 1:
|
|
flush()
|
|
|
|
flush()
|
|
return groups
|
|
|
|
|
|
def split_by_delimiter(blocks: list[Block], config: SplitConfig) -> list[list[Block]]:
|
|
if not config.delimiter:
|
|
raise ValueError("delimiter is required for delimiter split mode")
|
|
|
|
groups = partition_blocks_by_delimiter(blocks, config.delimiter)
|
|
if not groups:
|
|
return _split_by_size(blocks, config)
|
|
|
|
sized: list[list[Block]] = []
|
|
for group in groups:
|
|
rendered_len = sum(len(b.render()) + 2 for b in group)
|
|
if rendered_len <= config.max_chunk_size:
|
|
sized.append(group)
|
|
else:
|
|
sized.extend(_split_by_size(group, config))
|
|
return sized
|