""" Document chunking strategies for RAG. Implements semantic chunking with overlap, metadata enrichment, and configurable strategies for different content types. """ import re from typing import List, Dict, Any, Optional from dataclasses import dataclass import logging logger = logging.getLogger(__name__) @dataclass class DocumentChunk: """Represents a single document chunk with metadata.""" content: str metadata: Dict[str, Any] chunk_id: str def to_dict(self) -> Dict[str, Any]: """Convert to dictionary for storage.""" return { "content": self.content, "metadata": self.metadata, "chunk_id": self.chunk_id } class SemanticChunker: """ Smart chunking that preserves semantic meaning. Features: - Splits on natural boundaries (paragraphs, sentences) - Maintains context with overlap - Preserves code blocks intact - Enriches chunks with metadata """ def __init__( self, chunk_size: int = 600, chunk_overlap: int = 100, preserve_code_blocks: bool = True ): self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap self.preserve_code_blocks = preserve_code_blocks def chunk_document( self, text: str, metadata: Optional[Dict[str, Any]] = None ) -> List[DocumentChunk]: """ Split document into semantically meaningful chunks. Args: text: Document text to chunk metadata: Optional metadata to attach to all chunks Returns: List of DocumentChunk objects """ if not text or not text.strip(): logger.warning("Empty text provided for chunking") return [] metadata = metadata or {} # Extract and preserve code blocks code_blocks = [] if self.preserve_code_blocks: text, code_blocks = self._extract_code_blocks(text) # Split into paragraphs first paragraphs = self._split_paragraphs(text) # Create chunks chunks = [] current_chunk = [] current_size = 0 for i, para in enumerate(paragraphs): para_tokens = self._estimate_tokens(para) # If single paragraph exceeds chunk size, split by sentences if para_tokens > self.chunk_size: if current_chunk: chunks.append(self._create_chunk( current_chunk, metadata, len(chunks) )) current_chunk = [] current_size = 0 # Split long paragraph sentence_chunks = self._split_long_paragraph(para, metadata, len(chunks)) chunks.extend(sentence_chunks) continue # Add paragraph to current chunk if current_size + para_tokens <= self.chunk_size: current_chunk.append(para) current_size += para_tokens else: # Save current chunk if current_chunk: chunks.append(self._create_chunk( current_chunk, metadata, len(chunks) )) # Start new chunk with overlap overlap_text = self._get_overlap_text(current_chunk) current_chunk = [overlap_text, para] if overlap_text else [para] current_size = self._estimate_tokens(overlap_text) + para_tokens # Add remaining chunk if current_chunk: chunks.append(self._create_chunk( current_chunk, metadata, len(chunks) )) # Reinsert code blocks if code_blocks: chunks = self._reinsert_code_blocks(chunks, code_blocks) logger.info(f"Created {len(chunks)} chunks from document") return chunks def _extract_code_blocks(self, text: str) -> tuple[str, List[Dict[str, str]]]: """Extract code blocks to preserve them intact.""" code_pattern = r'```[\s\S]*?```|`[^`]+`' code_blocks = [] def replace_code(match): placeholder = f"__CODE_BLOCK_{len(code_blocks)}__" code_blocks.append({ "placeholder": placeholder, "content": match.group(0) }) return placeholder text_without_code = re.sub(code_pattern, replace_code, text) return text_without_code, code_blocks def _reinsert_code_blocks( self, chunks: List[DocumentChunk], code_blocks: List[Dict[str, str]] ) -> List[DocumentChunk]: """Reinsert code blocks into chunks.""" for chunk in chunks: for code_block in code_blocks: chunk.content = chunk.content.replace( code_block["placeholder"], code_block["content"] ) return chunks def _split_paragraphs(self, text: str) -> List[str]: """Split text into paragraphs.""" # Split on double newlines or more paragraphs = re.split(r'\n\s*\n', text) return [p.strip() for p in paragraphs if p.strip()] def _split_long_paragraph( self, paragraph: str, metadata: Dict[str, Any], start_idx: int ) -> List[DocumentChunk]: """Split a long paragraph by sentences.""" sentences = re.split(r'(?<=[.!?])\s+', paragraph) chunks = [] current_chunk = [] current_size = 0 for sentence in sentences: sentence_tokens = self._estimate_tokens(sentence) if current_size + sentence_tokens <= self.chunk_size: current_chunk.append(sentence) current_size += sentence_tokens else: if current_chunk: chunks.append(self._create_chunk( current_chunk, metadata, start_idx + len(chunks) )) current_chunk = [sentence] current_size = sentence_tokens if current_chunk: chunks.append(self._create_chunk( current_chunk, metadata, start_idx + len(chunks) )) return chunks def _create_chunk( self, text_segments: List[str], metadata: Dict[str, Any], chunk_idx: int ) -> DocumentChunk: """Create a DocumentChunk from text segments.""" content = "\n\n".join(text_segments) # Enrich metadata enriched_metadata = { **metadata, "chunk_index": chunk_idx, "chunk_size": len(content), "has_code": "```" in content or "`" in content, } chunk_id = f"{metadata.get('source', 'unknown')}_{chunk_idx}" return DocumentChunk( content=content, metadata=enriched_metadata, chunk_id=chunk_id ) def _get_overlap_text(self, chunks: List[str]) -> str: """Get overlap text from previous chunks.""" if not chunks: return "" combined = " ".join(chunks[-2:]) tokens = self._estimate_tokens(combined) if tokens <= self.chunk_overlap: return combined # Truncate to overlap size words = combined.split() overlap_words = words[-(self.chunk_overlap // 4):] return " ".join(overlap_words) @staticmethod def _estimate_tokens(text: str) -> int: """Rough token estimation (1 token ≈ 4 characters).""" return len(text) // 4 def create_chunker(chunk_size: int = 600, chunk_overlap: int = 100) -> SemanticChunker: """Factory function to create a chunker instance.""" return SemanticChunker( chunk_size=chunk_size, chunk_overlap=chunk_overlap, preserve_code_blocks=True )