Source code for ffai.rag.splitters.base

"""Define the abstract ChunkerBase interface and the TextChunk and HierarchicalTextChunk data classes."""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Any


[docs] @dataclass class TextChunk: """A single chunk of text produced by a chunker. Attributes: content: The text content of the chunk. chunk_index: Zero-based index of this chunk in the sequence. start_char: Character offset where this chunk begins in the source. end_char: Character offset where this chunk ends in the source. metadata: Optional metadata attached to this chunk. """ content: str chunk_index: int start_char: int end_char: int metadata: dict[str, Any] | None = None
[docs] @dataclass class HierarchicalTextChunk(TextChunk): """A text chunk that participates in a parent-child hierarchy. Used by hierarchical chunking strategies that maintain relationships between larger parent chunks and smaller child chunks. Attributes: id: Unique identifier for this chunk. parent_id: Identifier of the parent chunk, or None if this is a root. child_ids: Identifiers of child chunks, or an empty list. hierarchy_level: Depth in the hierarchy (0 = root). """ id: str = "" parent_id: str | None = None child_ids: list[str] | None = None hierarchy_level: int = 0 def __post_init__(self) -> None: if self.child_ids is None: self.child_ids = [] if self.metadata is None: self.metadata = {}
[docs] class ChunkerBase(ABC): """Abstract base class for text chunking strategies. Args: chunk_size: Maximum size of each chunk (interpretation varies by strategy). chunk_overlap: Overlap between consecutive chunks. metadata: Default metadata to attach to all chunks. """ def __init__( self, chunk_size: int = 1000, chunk_overlap: int = 200, metadata: dict[str, Any] | None = None, ) -> None: self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap self.default_metadata = metadata or {}
[docs] @abstractmethod def chunk( self, text: str, metadata: dict[str, Any] | None = None, ) -> list[TextChunk]: """Split text into chunks. Args: text: The text to split. metadata: Optional metadata to attach to each chunk (merged with default). Returns: List of TextChunk objects. """ pass
def _merge_metadata(self, metadata: dict[str, Any] | None) -> dict[str, Any]: """Merge provided metadata with default metadata.""" merged = self.default_metadata.copy() if metadata: merged.update(metadata) return merged def _validate_params(self) -> None: """Validate chunker parameters.""" if self.chunk_size <= 0: raise ValueError("chunk_size must be positive") if self.chunk_overlap < 0: raise ValueError("chunk_overlap cannot be negative") if self.chunk_overlap >= self.chunk_size: raise ValueError("chunk_overlap must be less than chunk_size") @property def name(self) -> str: """Return the chunker strategy name.""" return self.__class__.__name__.replace("Chunker", "").lower()