ragbandit-core 0.1.1__py3-none-any.whl

ragbandit/documents/chunkers/semantic_chunker.py

@@ -0,0 +1,205 @@
+from datetime import datetime, timezone
+from pydantic import BaseModel
+from ragbandit.schema import (
+    ProcessingResult,
+    Chunk,
+    ChunkMetadata,
+    ChunkingResult,
+)
+from ragbandit.utils.token_usage_tracker import TokenUsageTracker
+
+from ragbandit.prompt_tools.semantic_chunker_tools import (
+    find_semantic_break_tool,
+)
+
+from ragbandit.documents.chunkers.base_chunker import BaseChunker
+
+
+class SemanticBreak(BaseModel):
+    semantic_break: str
+
+
+class SemanticChunker(BaseChunker):
+    """
+    A document chunker that uses semantic understanding to split documents
+    into coherent chunks based on content.
+    """
+
+    def __init__(
+        self,
+        min_chunk_size: int = 500,
+        name: str | None = None,
+        api_key: str | None = None
+    ):
+        """
+        Initialize the semantic chunker.
+
+        Args:
+            min_chunk_size: Minimum size for chunks
+                            (smaller chunks will be merged)
+            name: Optional name for the chunker
+            api_key: Mistral API Key
+        """
+        super().__init__(name, api_key)
+        self.min_chunk_size = min_chunk_size
+
+    def semantic_chunk_pages(
+        self, pages: list, usage_tracker: TokenUsageTracker | None = None
+    ) -> list[Chunk]:
+        """
+        Chunk pages semantically using LLM-based semantic breaks.
+
+        Args:
+            pages: List of page dictionaries with markdown content
+            usage_tracker: Optional tracker for token usage
+
+        Returns:
+            A list of chunk dictionaries
+        """
+        i = 0
+        full_text = pages[i].markdown
+        chunks: list[Chunk] = []
+
+        while i < len(pages):
+            # If we have "remainder" from the last iteration,
+            # it might be appended here
+            break_lead = find_semantic_break_tool(
+                api_key=self.api_key,
+                text=full_text,
+                usage_tracker=usage_tracker
+            )
+
+            if break_lead == "NO_BREAK":
+                # This means the LLM found no break;
+                # treat the entire `full_text` as one chunk
+                meta = ChunkMetadata(page_index=i, images=[], extra={})
+                chunks.append(Chunk(text=full_text, metadata=meta))
+                # Move to the next page
+                i += 1
+                if i < len(pages):
+                    full_text = pages[i].markdown
+                else:
+                    break
+            else:
+                # Attempt to find the snippet in the text
+                idx = full_text.find(break_lead)
+
+                # If exact match fails, try progressively shorter versions
+                if idx == -1 and len(break_lead) > 0:
+                    current_break_lead = break_lead
+                    min_length = 10  # Minimum characters to try matching
+
+                    # Try progressively shorter versions
+                    # of the break_lead until we find a match
+                    # or reach the minimum length
+                    while idx == -1 and len(current_break_lead) >= min_length:
+                        # Cut the break_lead in half and try again
+                        current_break_lead = current_break_lead[
+                            : len(current_break_lead) // 2
+                        ]
+                        idx = full_text.find(current_break_lead)
+
+                if idx == -1:
+                    # If we still can't find the snippet after
+                    # trying shorter versions,
+                    # fallback: chunk everything as is
+                    meta = ChunkMetadata(page_index=i, images=[], extra={})
+                    chunks.append(Chunk(text=full_text, metadata=meta))
+                    i += 1
+                    if i < len(pages):
+                        full_text = pages[i].markdown
+                    else:
+                        break
+                else:
+                    # We found a break
+                    chunk_text = full_text[:idx]
+                    remainder = full_text[idx:]
+                    meta = ChunkMetadata(page_index=i, images=[], extra={})
+                    chunks.append(Chunk(text=chunk_text, metadata=meta))
+
+                    # Now we see if remainder is too small
+                    if len(remainder) < 1500:  # ~some threshold
+                        i += 1
+                        if i < len(pages):
+                            # Combine remainder with next page
+                            remainder += "\n" + pages[i].markdown
+                    # remainder becomes the new full_text
+                    full_text = remainder
+
+                    # If we used up the last page, break
+                    if i >= len(pages):
+                        # Possibly chunk the remainder if it's not empty
+                        if len(full_text.strip()) > 0:
+                            meta = ChunkMetadata(
+                                page_index=min(i, len(pages) - 1),
+                                images=[],
+                                extra={},
+                            )
+                            chunks.append(Chunk(text=full_text, metadata=meta))
+                        break
+
+        return chunks
+
+    def chunk(
+        self,
+        proc_result: ProcessingResult,
+        usage_tracker: TokenUsageTracker | None = None,
+    ) -> ChunkingResult:
+        """
+        Chunk the document using semantic chunking.
+
+        Args:
+            proc_result: The ProcessingResult containing
+                      document content to chunk
+            usage_tracker: Tracker for token usage during chunking
+
+        Returns:
+            A list of chunk dictionaries
+        """
+        self.logger.info("Starting semantic chunking")
+
+        # Get the pages from the response
+        pages = proc_result.pages
+
+        # Perform semantic chunking
+        chunks = self.semantic_chunk_pages(pages, usage_tracker)
+
+        # Attach image data to chunks using shared helper
+        chunks = self.attach_images(chunks, proc_result)
+
+        # Merge small chunks if needed
+        chunks = self.process_chunks(chunks)
+
+        return ChunkingResult(
+            processed_at=datetime.now(timezone.utc),
+            chunks=chunks,
+            metrics=usage_tracker.get_summary() if usage_tracker else None,
+        )
+
+    def process_chunks(
+        self, chunks: list[Chunk]
+    ) -> list[Chunk]:
+        """
+        Process chunks after initial chunking - merge small chunks.
+
+        Args:
+            chunks: The initial chunks produced by the chunk method
+
+        Returns:
+            Processed chunks with small chunks merged
+        """
+        # Check if any chunks are too small
+        min_len = min([len(c.text) for c in chunks]) if chunks else 0
+
+        # Merge small chunks if needed
+        if min_len < self.min_chunk_size:
+            self.logger.info(
+                f"Found chunks smaller than {self.min_chunk_size} characters. "
+                "Merging..."
+            )
+            chunks = self.merge_small_chunks(
+                chunks, min_size=self.min_chunk_size
+            )
+            self.logger.info(f"After merging: {len(chunks)} chunks")
+
+        return chunks

ragbandit/documents/document_pipeline.py

@@ -0,0 +1,350 @@
+"""
+Document processing pipeline that orchestrates multiple document processors.
+
+This module provides the main DocumentPipeline class that manages the execution
+of document processors in sequence, chunking, and embedding.
+"""
+
+import logging
+import traceback
+from datetime import datetime, timezone
+from typing import Callable
+import time
+from dataclasses import dataclass
+from ragbandit.schema import (
+    OCRResult,
+    ProcessingResult,
+    ChunkingResult,
+    EmbeddingResult,
+    DocumentPipelineResult,
+    TimingMetrics,
+    StepReport,
+    StepStatus,
+)
+
+from ragbandit.documents.ocr import BaseOCR
+from ragbandit.documents.processors.base_processor import BaseProcessor
+from ragbandit.documents.chunkers.base_chunker import BaseChunker
+from ragbandit.documents.embedders.base_embedder import BaseEmbedder
+from ragbandit.utils.token_usage_tracker import TokenUsageTracker
+
+from ragbandit.utils.in_memory_log_handler import InMemoryLogHandler
+
+
+class DocumentPipeline:
+    """Pipeline for processing documents through a
+    sequence of document processors, chunkers, and embedders.
+
+    The pipeline manages the execution of document processors in sequence,
+    where each processor receives the output of the previous processor.
+    The pipeline also tracks token usage and costs for each document.
+    """
+
+    @dataclass
+    class _PipelineStep:
+        key: str  # "ocr" | "processing" | "chunking" | "embedding"
+        run: Callable[[], object]
+        on_success: Callable[[object], None]
+
+    def __init__(
+        self,
+        ocr_processor: BaseOCR | None = None,
+        processors: list[BaseProcessor] | None = None,
+        chunker: BaseChunker | None = None,
+        embedder: BaseEmbedder | None = None,
+        logger: logging.Logger | None = None,
+    ):
+        """Initialize a new document processing pipeline.
+
+        All components are optional to allow running
+        individual steps independently.
+        For full pipeline execution via process(),
+        all components must be provided.
+
+        Args:
+            ocr_processor: OCR processor to use (required for run_ocr
+                           and process)
+            processors: List of document processors to execute in
+                        sequence
+            chunker: Chunker to use for document chunking (required
+                     for run_chunker and process)
+            embedder: Embedder to use for chunk embedding (required
+                      for run_embedder and process)
+            logger: Optional logger for pipeline events
+        """
+        self.ocr_processor = ocr_processor
+        self.processors = processors or []
+        self.chunker = chunker
+        self.embedder = embedder
+
+        # Set up logging with more explicit configuration
+        self.logger = logger or logging.getLogger(__name__)
+
+        self._transcript = InMemoryLogHandler(level=logging.DEBUG)
+        root_logger = logging.getLogger()
+        root_logger.addHandler(self._transcript)
+
+        # Ensure we're generating logs
+        self.logger.info("DocumentPipeline initialized")
+
+    def add_processor(self, processor: BaseProcessor) -> None:
+        """Add a processor to the pipeline.
+
+        Args:
+            processor: The document processor to add
+        """
+        self.processors.append(processor)
+        self.logger.info(f"Added processor: {processor}")
+
+    def _fresh_buffer(self):
+        self._transcript.clear()
+        # Ensure the handler is still attached
+        root_logger = logging.getLogger()
+        if self._transcript not in root_logger.handlers:
+            root_logger.addHandler(self._transcript)
+
+    def run_ocr(self, pdf_filepath: str) -> OCRResult:
+        """Perform OCR on a PDF file using the configured OCR processor.
+
+        Args:
+            pdf_filepath: Path to the PDF file to process
+
+        Returns:
+            OCRResult: The OCR result from the processor
+
+        Raises:
+            ValueError: If ocr_processor is not configured
+        """
+        if not self.ocr_processor:
+            raise ValueError("ocr_processor is required for OCR operation")
+        return self.ocr_processor.process(pdf_filepath)
+
+    def run_processors(
+        self,
+        ocr_result: OCRResult,
+    ) -> list[ProcessingResult]:
+        """Process a document through the processors pipeline.
+
+        Args:
+            ocr_result: The initial OCR result to process
+
+        Returns:
+            A list of ProcessingResult with additional metadata
+            from all processors
+        """
+        processing_results: list[ProcessingResult] = []
+
+        # Start the processor chain with the raw OCRResult; each processor
+        # is responsible for converting it to ProcessingResult if needed.
+        prev_result = ocr_result
+
+        # Process the document through each processor in sequence
+        for processor in self.processors:
+            self.logger.info(f"Running processor: {processor}")
+
+            # Give each processor its own usage tracker
+            proc_usage = TokenUsageTracker()
+
+            start_processing = time.perf_counter()
+            proc_result = processor.process(prev_result, proc_usage)
+            end_processing = time.perf_counter()
+
+            # Attach token usage summary to metrics
+            proc_result.metrics = proc_usage.get_summary()
+            proc_duration = end_processing - start_processing
+            proc_result.processing_duration = proc_duration
+
+            processing_results.append(proc_result)
+            prev_result = proc_result
+            self.logger.info(f"{processor} completed successfully")
+
+        return processing_results
+
+    def run_chunker(
+        self,
+        doc: ProcessingResult | OCRResult,
+    ) -> ChunkingResult:
+        """Chunk the document using the configured chunker.
+
+        Args:
+            doc: The ProcessingResult or OCRResult to chunk
+
+        Returns:
+            A ChunkingResult object
+
+        Raises:
+            ValueError: If chunker is not configured
+        """
+        if not self.chunker:
+            raise ValueError("chunker is required for chunking operation")
+        proc_result = (
+            doc
+            if isinstance(doc, ProcessingResult)
+            else BaseProcessor.ensure_processing_result(doc)
+        )
+        usage_tracker = TokenUsageTracker()
+        # Generate chunks via chunker -> returns ChunkingResult
+        chunk_result = self.chunker.chunk(proc_result, usage_tracker)
+        return chunk_result
+
+    def run_embedder(
+        self,
+        chunk_result: ChunkingResult,
+    ) -> EmbeddingResult:
+        """Embed chunks using the configured embedder.
+
+        Args:
+            chunk_result: The ChunkingResult to embed
+
+        Returns:
+            An EmbeddingResult containing embeddings for each chunk
+
+        Raises:
+            ValueError: If embedder is not configured
+        """
+        if not self.embedder:
+            raise ValueError("embedder is required for embedding operation")
+        usage_tracker = TokenUsageTracker()
+        embedding_result = self.embedder.embed_chunks(
+            chunk_result, usage_tracker
+        )
+        return embedding_result
+
+    def _run_step(
+        self,
+        step: _PipelineStep,
+        dpr: DocumentPipelineResult,
+        start_total: float,
+    ) -> tuple[bool, object | None]:
+        key = step.key  # e.g. "ocr"
+        self.logger.info(f"Starting {key} step…")
+        start = time.perf_counter()
+        try:
+            result = step.run()
+            setattr(dpr.step_report, key, StepStatus.success)
+            setattr(dpr.timings, key, time.perf_counter() - start)
+            step.on_success(result)
+            self.logger.info(f"Step {key} completed")
+            return True, result
+        except Exception as exc:
+            tb = traceback.format_exc()
+            self.logger.error(f"Step {key} failed: {exc}\n{tb}")
+            setattr(dpr.step_report, key, StepStatus.failed)
+            setattr(dpr.timings, key, time.perf_counter() - start)
+            dpr.timings.total_duration = time.perf_counter() - start_total
+            return False, None
+
+    def process(
+        self,
+        pdf_filepath: str
+    ) -> DocumentPipelineResult:
+        """Run the configured pipeline steps in order.
+
+        Raises:
+            ValueError: If any required component is not configured
+        """
+        # Validate all components are present for full pipeline execution
+        if not self.ocr_processor:
+            raise ValueError(
+                "ocr_processor is required for full pipeline execution"
+            )
+        if not self.chunker:
+            raise ValueError(
+                "chunker is required for full pipeline execution"
+            )
+        if not self.embedder:
+            raise ValueError(
+                "embedder is required for full pipeline execution"
+            )
+
+        start_total = time.perf_counter()
+        dpr = DocumentPipelineResult(
+            source_file_path=pdf_filepath,
+            processed_at=datetime.now(timezone.utc),
+            pipeline_config={
+                "ocr": str(self.ocr_processor),
+                "processors": [str(p) for p in self.processors],
+                "chunker": str(self.chunker),
+                "embedder": str(self.embedder),
+            },
+            timings=TimingMetrics(),
+            total_metrics=[],
+            step_report=StepReport(),
+        )
+
+        # ---------------- helpers ----------------
+        def _on_success(attr):
+            def handler(res):
+                # 1. Set the result (save the step result to DPR)
+                setattr(dpr, attr, res)
+                # 2. Set the metrics of the result to total metrics
+                # - res may be a single result or a list of results
+                if isinstance(res, list):
+                    dpr.total_metrics.extend(
+                        [r.metrics for r in res if r.metrics]
+                    )
+                else:
+                    if isinstance(res.metrics, list):
+                        dpr.total_metrics.extend(res.metrics or [])
+                    else:
+                        dpr.total_metrics.append(res.metrics)
+            return handler
+
+        # placeholders for passing results between steps
+        ocr_res: OCRResult | None = None
+        proc_results: list[ProcessingResult] | None = None
+        chunk_res: ChunkingResult | None = None
+
+        # ---------------- step table ----------------
+        steps = [
+            self._PipelineStep(
+                "ocr",
+                lambda: self.run_ocr(pdf_filepath),
+                _on_success("ocr_result"),
+            ),
+            self._PipelineStep(
+                "processing",
+                lambda: self.run_processors(ocr_res),
+                _on_success("processing_results"),
+            ),
+            self._PipelineStep(
+                "chunking",
+                lambda: self.run_chunker(
+                    proc_results[-1] if proc_results else ocr_res
+                ),
+                _on_success("chunking_result"),
+            ),
+            self._PipelineStep(
+                "embedding",
+                lambda: self.run_embedder(chunk_res),
+                _on_success("embedding_result"),
+            ),
+        ]
+
+        try:
+            for st in steps:
+                ok, res = self._run_step(st, dpr, start_total)
+                if not ok:
+                    return dpr
+
+                # propagate outputs for later steps
+                if st.key == "ocr":
+                    ocr_res = res  # type: ignore
+                elif st.key == "processing":
+                    proc_results = res  # type: ignore
+                elif st.key == "chunking":
+                    chunk_res = res  # type: ignore
+
+            # aggregate total cost once
+            dpr.total_cost_usd = sum(
+                m.total_cost_usd  # type: ignore[attr-defined]
+                for m in dpr.total_metrics
+                if m and getattr(m, "total_cost_usd", None) is not None
+            )
+
+            dpr.timings.total_duration = time.perf_counter() - start_total
+            self.logger.info("Document processing completed.")
+            return dpr
+        finally:
+            dpr.logs = self._transcript.dump()
+            logging.getLogger().removeHandler(self._transcript)

ragbandit/documents/embedders/__init__.py

@@ -0,0 +1,14 @@
+"""
+Document embedders for generating vector representations of document chunks.
+
+This module provides embedders that convert document chunks
+into vector embeddings for semantic search and similarity comparison.
+"""
+
+from ragbandit.documents.embedders.base_embedder import BaseEmbedder
+from ragbandit.documents.embedders.mistral_embedder import MistralEmbedder
+
+__all__ = [
+    "BaseEmbedder",
+    "MistralEmbedder"
+]

ragbandit/documents/embedders/base_embedder.py

@@ -0,0 +1,82 @@
+from abc import ABC, abstractmethod
+
+# Third-party
+import numpy as np
+
+# Project
+from ragbandit.schema import (
+    ChunkingResult,
+    EmbeddingResult,
+)
+from ragbandit.utils.token_usage_tracker import TokenUsageTracker
+
+
+class BaseEmbedder(ABC):
+    """
+    Abstract base class for document embedders.
+
+    This class defines the interface for embedding document chunks.
+    Concrete implementations should handle the specifics of
+    generating embeddings using different models or providers.
+    """
+
+    def __init__(self, name: str = None):
+        """
+        Initialize the document embedder.
+
+        Args:
+            name: Optional name for the embedder
+        """
+        self.name = name or self.__class__.__name__
+
+        # Set up logging
+        import logging
+        self.logger = logging.getLogger(f"{self.__class__.__name__}")
+
+    @abstractmethod
+    def embed_chunks(
+        self,
+        chunk_result: ChunkingResult,
+        usage_tracker: TokenUsageTracker | None = None,
+    ) -> EmbeddingResult:
+        """
+        Generate embeddings for a ChunkingResult.
+
+        Args:
+            chunk_result: The ChunkingResult whose chunks will be embedded
+            usage_tracker: Optional tracker for token usage
+
+        Returns:
+            An EmbeddingResult containing embedded chunks
+        """
+        raise NotImplementedError
+
+    def cosine_similarity(self, a: np.ndarray, b: np.ndarray) -> float:
+        """
+        Calculate cosine similarity between two embedding vectors.
+
+        Args:
+            a: First embedding vector
+            b: Second embedding vector
+
+        Returns:
+            Cosine similarity (1 = identical, 0 = orthogonal, -1 = opposite)
+        """
+        return (a @ b) / (np.linalg.norm(a) * np.linalg.norm(b))
+
+    def cosine_distance(self, a: np.ndarray, b: np.ndarray) -> float:
+        """
+        Calculate cosine distance between two embedding vectors.
+
+        Args:
+            a: First embedding vector
+            b: Second embedding vector
+
+        Returns:
+            Cosine distance (0 = identical, 2 = opposite)
+        """
+        return 1 - self.cosine_similarity(a, b)
+
+    def __repr__(self) -> str:
+        """Return string representation of the embedder."""
+        return f"{self.__class__.__name__}"