ruby2210/rag-chatbot

1"""2Content chunking utilities for the RAG Chatbot application.3Handles text segmentation for embedding and retrieval.4"""5import re6from typing import List, Tuple7from ..utils.logging import get_logger8 9 10logger = get_logger(__name__)11 12 13class Chunker:14    """15    Utility class for chunking text content into manageable pieces for embedding.16    Implements various chunking strategies to optimize retrieval performance.17    """18    def __init__(self, default_chunk_size: int = 1000, default_overlap: int = 100):19        self.default_chunk_size = default_chunk_size20        self.default_overlap = default_overlap21 22    def chunk_by_size(self, text: str, chunk_size: int = None, overlap: int = None) -> List[str]:23        """24        Split text into chunks of approximately the specified size with overlap.25        """26        if chunk_size is None:27            chunk_size = self.default_chunk_size28        if overlap is None:29            overlap = self.default_overlap30 31        if len(text) <= chunk_size:32            return [text]33 34        chunks = []35        start = 036 37        while start < len(text):38            end = start + chunk_size39 40            # If this is not the last chunk, try to break at sentence or paragraph boundary41            if end < len(text):42                # Look for good breaking points near the end43                chunk_segment = text[start:end]44 45                # Prefer to break at paragraph boundaries46                last_paragraph_end = chunk_segment.rfind('\n\n')47                if last_paragraph_end == -1 or last_paragraph_end < chunk_size // 2:48                    # If no good paragraph break, look for sentence endings49                    last_sentence_end = max(50                        chunk_segment.rfind('. '),51                        chunk_segment.rfind('?'),52                        chunk_segment.rfind('!'),53                        chunk_segment.rfind('\n')54                    )55 56                    if last_sentence_end == -1 or last_sentence_end < chunk_size // 2:57                        # If no good sentence break, look for word boundaries58                        last_space = chunk_segment.rfind(' ')59                        if last_space > chunk_size // 2:60                            end = start + last_space61                        else:62                            end = start + chunk_size  # Use full chunk size if no good break found63                    else:64                        end = start + last_sentence_end + 165                else:66                    end = start + last_paragraph_end + 267 68            chunk_text = text[start:end].strip()69            if chunk_text:  # Only add non-empty chunks70                chunks.append(chunk_text)71 72            # Move start forward, with overlap73            start = end - overlap if end < len(text) else end74 75        return [chunk for chunk in chunks if chunk.strip()]  # Remove any empty chunks76 77    def chunk_by_headings(self, text: str) -> List[Tuple[str, str]]:78        """79        Split text based on headings (useful for structured documents like markdown).80        Returns a list of (heading, content) tuples.81        """82        # Split text by common heading patterns83        heading_pattern = r'\n(#{1,6}\s+.*?)(?=\n#|\n$)'84        parts = re.split(heading_pattern, '\n' + text, flags=re.MULTILINE)85 86        # Remove the first empty element if it exists87        if parts and parts[0].startswith('\n'):88            parts[0] = parts[0][1:]  # Remove leading newline89 90        chunks_with_headings = []91        current_heading = "Introduction"  # Default heading92 93        for i, part in enumerate(parts):94            if i % 2 == 0:  # This is content95                if part.strip():96                    chunks_with_headings.append((current_heading, part.strip()))97            else:  # This is a heading98                current_heading = part.strip('# ').strip()99 100        return chunks_with_headings101 102    def chunk_by_paragraphs(self, text: str, max_chunk_size: int = None) -> List[str]:103        """104        Split text by paragraphs, combining paragraphs as needed to approach the target size.105        """106        if max_chunk_size is None:107            max_chunk_size = self.default_chunk_size108 109        # Split by paragraph boundaries110        paragraphs = [p.strip() for p in text.split('\n\n') if p.strip()]111 112        if not paragraphs:113            # If no paragraph breaks, fall back to size-based chunking114            return self.chunk_by_size(text, max_chunk_size)115 116        chunks = []117        current_chunk = ""118 119        for paragraph in paragraphs:120            # If adding this paragraph would exceed the size limit121            if len(current_chunk) + len(paragraph) > max_chunk_size and current_chunk:122                # Save the current chunk and start a new one123                chunks.append(current_chunk.strip())124                current_chunk = paragraph125            else:126                # Add the paragraph to the current chunk127                if current_chunk:128                    current_chunk += "\n\n" + paragraph129                else:130                    current_chunk = paragraph131 132        # Add the final chunk if it exists133        if current_chunk:134            chunks.append(current_chunk.strip())135 136        return chunks137 138    def chunk_text(self, text: str, strategy: str = "size", **kwargs) -> List[str]:139        """140        General method to chunk text using the specified strategy.141        Supported strategies: "size", "paragraphs"142        """143        if strategy == "size":144            return self.chunk_by_size(text, kwargs.get('chunk_size'), kwargs.get('overlap'))145        elif strategy == "paragraphs":146            return self.chunk_by_paragraphs(text, kwargs.get('max_chunk_size'))147        else:148            logger.warning(f"Unknown chunking strategy: {strategy}, defaulting to 'size'")149            return self.chunk_by_size(text)150 151 152# Create a singleton instance153chunker = Chunker()