Files
Project-Thoth/processors/rag_ingestion/chunking.py
T

96 lines
3.2 KiB
Python

"""Deterministic, fixed-width character chunking for ingested documents."""
import hashlib
from dataclasses import dataclass
from .models import Chunk, Document
# A rough four-characters-per-token approximation of the work order's suggested
# 800-token chunks with 120-token overlap; these values remain character counts.
DEFAULT_CHUNK_SIZE = 3_200
DEFAULT_CHUNK_OVERLAP = 480
SIZE_UNIT = "characters"
@dataclass(frozen=True, slots=True)
class ChunkConfiguration:
"""Visible settings that determine mechanical chunk boundaries."""
chunk_size: int = DEFAULT_CHUNK_SIZE
chunk_overlap: int = DEFAULT_CHUNK_OVERLAP
def __post_init__(self) -> None:
if self.chunk_size <= 0:
raise ValueError("chunk size must be greater than zero")
if self.chunk_overlap < 0:
raise ValueError("chunk overlap must not be negative")
if self.chunk_overlap >= self.chunk_size:
raise ValueError("chunk overlap must be smaller than chunk size")
def create_chunk_id(
document_id: str,
chunk_number: int,
configuration: ChunkConfiguration,
) -> str:
"""Hash source identity, ordinal position, and boundary configuration."""
identity = (
f"{document_id}:characters:"
f"{configuration.chunk_size}:{configuration.chunk_overlap}:{chunk_number}"
)
return hashlib.sha256(identity.encode("utf-8")).hexdigest()
def chunk_document(
document: Document,
configuration: ChunkConfiguration | None = None,
) -> list[Chunk]:
"""Divide ``document.raw_text`` into ordered character ranges.
The next range starts ``chunk_size - chunk_overlap`` characters after the
current range. Empty input intentionally produces no retrieval units.
"""
settings = configuration or ChunkConfiguration()
source_text = document.raw_text
if not source_text:
return []
step_size = settings.chunk_size - settings.chunk_overlap
boundaries: list[tuple[int, int]] = []
start_offset = 0
while start_offset < len(source_text):
end_offset = min(start_offset + settings.chunk_size, len(source_text))
boundaries.append((start_offset, end_offset))
if end_offset == len(source_text):
break
start_offset += step_size
chunks: list[Chunk] = []
for chunk_number, (start_offset, end_offset) in enumerate(boundaries):
previous_end = boundaries[chunk_number - 1][1] if chunk_number > 0 else 0
next_start = (
boundaries[chunk_number + 1][0]
if chunk_number + 1 < len(boundaries)
else end_offset
)
text = source_text[start_offset:end_offset]
chunks.append(
Chunk(
chunk_id=create_chunk_id(document.document_id, chunk_number, settings),
document_id=document.document_id,
chunk_number=chunk_number,
text=text,
start_offset=start_offset,
end_offset=end_offset,
size=len(text),
size_unit=SIZE_UNIT,
overlap_with_previous=max(0, previous_end - start_offset),
overlap_with_next=max(0, end_offset - next_start),
)
)
return chunks