ragbandit-core 0.1.1__py3-none-any.whl

ragbandit/__init__.py

@@ -0,0 +1,26 @@
+"""ragbandit core package.
+
+This package contains sub-modules for document processing,
+RAG pipeline configuration/execution, and evaluation utilities.
+Only lightweight interfaces and shared utilities are defined here;
+heavy logic resides in sub-packages.
+"""
+
+from importlib import metadata as _metadata
+
+__version__: str
+try:
+    __version__ = _metadata.version("ragbandit-core")
+except _metadata.PackageNotFoundError:  # pragma: no cover
+    __version__ = "0.0.0+dev"
+
+# Re-export public interfaces so that users can simply:
+#   from ragbandit import DocumentProcessor, RAGConfig, RAGPipeline, evaluate
+
+# from ragbandit.documents import DocumentPipeline
+
+
+__all__ = [
+    "__version__",
+    # "DocumentPipeline",
+]

ragbandit/config/__init__.py

@@ -0,0 +1,3 @@
+"""
+Configuration module for ragbandit package.
+"""

ragbandit/config/llms.py

@@ -0,0 +1,34 @@
+"""
+LLM configuration settings for ragbandit.
+
+This module defines default settings and constants for LLM interactions.
+"""
+
+# Default model settings
+DEFAULT_MODEL = "mistral-small-latest"
+DEFAULT_TEMPERATURE = 0.0
+
+# Retry settings
+DEFAULT_MAX_RETRIES = 3
+DEFAULT_RETRY_DELAY = 1.0  # seconds
+DEFAULT_BACKOFF_FACTOR = 2.0  # exponential backoff factor
+DEFAULT_TIMEOUT = 30.0  # seconds
+
+# Token limits
+MAX_PROMPT_TOKENS = {
+    "mistral-small-latest": 8000,
+    "mistral-medium-latest": 32000,
+    "mistral-large-latest": 32000,
+    "gpt-3.5-turbo": 4096,
+    "gpt-4": 8192,
+    "gpt-4-turbo": 128000,
+}
+
+# System prompts
+DEFAULT_SYSTEM_PROMPT = """You are a helpful AI assistant."""
+
+# Response formats
+JSON_FORMAT_INSTRUCTION = """
+Your response must be valid JSON that matches the following schema:
+{schema}
+"""

ragbandit/config/pricing.py

@@ -0,0 +1,38 @@
+"""
+Pricing configuration for LLM API calls.
+
+This module contains pricing constants for various
+LLM models and embedding models.
+"""
+
+# Token cost rates per 1M tokens (in USD)
+# Based on Mistral AI pricing as of July 2025
+MODEL_COSTS = {
+    # Format: "model_name": (input_cost_per_1M, output_cost_per_1M)
+    "mistral-small-latest": (2.00, 6.00),
+    "mistral-medium-latest": (6.00, 18.00),
+    "mistral-large-latest": (12.00, 36.00),
+    # Add other models as needed
+}
+
+# Embedding model costs per 1M tokens
+EMBEDDING_COSTS = {
+    # Format: "model_name": cost_per_1M_tokens
+    "mistral-embed": 0.10,
+    "text-embedding-3-small": 0.02,
+    "text-embedding-3-large": 0.13,
+    # Add other embedding models as needed
+}
+
+# OCR model costs per page (in EUR)
+OCR_MODEL_COSTS = {
+    # Format: "model_name": cost_per_page
+    "mistral-ocr-latest": 0.001,  # 1 EUR per 1000 pages
+    # Add other OCR models as needed
+}
+
+# Default OCR model to use if the specified model is not in OCR_MODEL_COSTS
+DEFAULT_OCR_MODEL = "mistral-ocr-latest"
+
+# Default model to use if the specified model is not in MODEL_COSTS
+DEFAULT_MODEL = "mistral-small-latest"

ragbandit/documents/__init__.py

@@ -0,0 +1,66 @@
+"""
+Document processing module for handling, analyzing, and transforming documents.
+
+This package provides tools for OCR, chunking,
+embedding, and processing documents.
+"""
+
+# Import key components from subdirectories
+from ragbandit.documents.document_pipeline import DocumentPipeline
+
+# Import from chunkers
+from ragbandit.documents.chunkers import (
+    BaseChunker,
+    FixedSizeChunker,
+    SemanticChunker,
+    SemanticBreak
+)
+
+# Import from processors
+from ragbandit.documents.processors import (
+    BaseProcessor,
+    FootnoteProcessor,
+    ReferencesProcessor
+)
+
+# Import from embedders
+from ragbandit.documents.embedders import (
+    BaseEmbedder,
+    MistralEmbedder
+)
+
+# Import from OCR
+from ragbandit.documents.ocr import (
+    BaseOCR,
+    MistralOCRDocument
+)
+
+# Import from utils
+from ragbandit.documents.utils import SecureFileHandler
+
+__all__ = [
+    # Main pipeline
+    "DocumentPipeline",
+
+    # Chunkers
+    "BaseChunker",
+    "FixedSizeChunker",
+    "SemanticChunker",
+    "SemanticBreak",
+
+    # Processors
+    "BaseProcessor",
+    "FootnoteProcessor",
+    "ReferencesProcessor",
+
+    # Embedders
+    "BaseEmbedder",
+    "MistralEmbedder",
+
+    # OCR
+    "BaseOCR",
+    "MistralOCRDocument",
+
+    # Utils
+    "SecureFileHandler"
+]

ragbandit/documents/chunkers/__init__.py

@@ -0,0 +1,18 @@
+"""
+Chunker implementations for document processing.
+
+This module provides various chunking strategies for documents.
+"""
+
+from ragbandit.documents.chunkers.base_chunker import BaseChunker
+from ragbandit.documents.chunkers.fixed_size_chunker import FixedSizeChunker
+from ragbandit.documents.chunkers.semantic_chunker import (
+    SemanticChunker, SemanticBreak
+)
+
+__all__ = [
+    "BaseChunker",
+    "FixedSizeChunker",
+    "SemanticChunker",
+    "SemanticBreak"
+]

ragbandit/documents/chunkers/base_chunker.py

@@ -0,0 +1,201 @@
+# ----------------------------------------------------------------------
+# Standard library
+import logging
+import re
+from abc import ABC, abstractmethod
+
+# Project
+from ragbandit.schema import (
+    ProcessingResult,
+    Chunk,
+    ChunkingResult,
+    Image,
+)
+from ragbandit.utils.token_usage_tracker import TokenUsageTracker
+
+
+class BaseChunker(ABC):
+    """
+    Base class for document chunking strategies.
+    Subclasses should implement the `chunk()` method to
+    provide specific chunking logic.
+    """
+
+    def __init__(self, name: str | None = None, api_key: str | None = None):
+        """
+        Initialize the chunker.
+
+        Args:
+            name: Optional name for the chunker
+            api_key: API key for LLM services
+        """
+        # Hierarchical names make it easy to filter later:
+        #   chunker.semantic, chunker.fixed_size, etc.
+        base = "chunker"
+        self.logger = logging.getLogger(
+            f"{base}.{name or self.__class__.__name__}"
+        )
+        self.api_key = api_key
+
+    @abstractmethod
+    def chunk(
+        self,
+        document: ProcessingResult,
+        usage_tracker: TokenUsageTracker | None = None,
+    ) -> ChunkingResult:
+        """
+        Chunk the document content from a ProcessingResult.
+
+        Args:
+            document: The ProcessingResult containing
+                      document content to chunk
+            usage_tracker: Optional tracker for token usage during chunking
+
+        Returns:
+            A `ChunkingResult` containing a list of `Chunk` objects and
+            optional metrics.
+        """
+        raise NotImplementedError
+
+    def merge_small_chunks(
+        self, chunks: list[Chunk], min_size: int
+    ) -> list[Chunk]:
+        """
+        Merge small chunks with adjacent chunks to ensure minimum chunk size.
+
+        Args:
+            chunks: The chunks to process
+            min_size: Minimum size for chunks (smaller chunks will be merged)
+
+        Returns:
+            Processed chunks with small chunks merged
+        """
+        if not chunks:
+            return []
+
+        merged = []
+        i = 0
+        n = len(chunks)
+
+        while i < n:
+            current_chunk = chunks[i]
+            current_text = current_chunk.text
+
+            # Check if this chunk is "small"
+            if len(current_text) < min_size:
+                # 1) Try to merge with the NEXT chunk if same page_index
+                next_chunk_exists = (i + 1) < n
+                if next_chunk_exists:
+                    next_chunk_same_page = (
+                        chunks[i + 1].metadata.page_index
+                        == current_chunk.metadata.page_index
+                    )
+                else:
+                    next_chunk_same_page = False
+
+                if i < n - 1 and next_chunk_same_page:
+                    # Merge current with the next chunk
+                    current_chunk.text += (" " + chunks[i + 1].text)
+
+                    # Merge images if they exist
+                    if (
+                        current_chunk.metadata.images
+                        and chunks[i + 1].metadata.images
+                    ):
+                        current_chunk.metadata.images.extend(
+                            chunks[i + 1].metadata.images
+                        )
+
+                    # We've used chunk i+1, so skip it
+                    i += 2
+
+                    # Now this newly merged chunk is complete; add to 'merged'
+                    merged.append(current_chunk)
+                else:
+                    # 2) Otherwise, try to merge with
+                    # PREVIOUS chunk in 'merged'
+                    if merged:
+                        # Merge current chunk into the last chunk in 'merged'
+                        merged[-1].text += (" " + current_chunk.text)
+
+                        # Merge images if they exist
+                        if (
+                            merged[-1].metadata.images
+                            and current_chunk.metadata.images
+                        ):
+                            merged[-1].metadata.images.extend(
+                                current_chunk.metadata.images
+                            )
+                    else:
+                        # If there's no previous chunk in 'merged', just add it
+                        merged.append(current_chunk)
+
+                    i += 1
+            else:
+                # If it's not "small," just add it as-is
+                merged.append(current_chunk)
+                i += 1
+
+        return merged
+
+    def process_chunks(
+        self, chunks: list[Chunk]
+    ) -> list[Chunk]:
+        """
+        Optional post-processing of chunks after initial chunking.
+        This can be overridden by subclasses to
+        implement additional processing.
+
+        Args:
+            chunks: The initial chunks produced by the chunk method
+
+        Returns:
+            Processed chunks
+        """
+        return chunks
+
+    # ------------------------------------------------------------------
+    # Shared helpers
+    def attach_images(
+        self,
+        chunks: list[Chunk],
+        proc_result: ProcessingResult,
+    ) -> list[Chunk]:
+        """Populate each Chunk's metadata.images with inlined image data.
+
+        Looks for `![img-XX.jpeg](img-XX.jpeg)` markers inside the chunk text
+        and copies the matching `image_base64` from the corresponding page's
+        images collection.
+        """
+
+        img_pattern = re.compile(r"!\[img-\d+\.jpeg\]\(img-\d+\.jpeg\)")
+
+        for chunk in chunks:
+            images_in_chunk = img_pattern.findall(chunk.text)
+            if not images_in_chunk:
+                # No image markers, ensure empty list and continue
+                chunk.metadata.images = []
+                continue
+
+            page_idx = chunk.metadata.page_index
+            rel_images = proc_result.pages[page_idx].images or []
+            chunk.metadata.images = []
+
+            for img_tag in images_in_chunk:
+                img_id = img_tag.split("[")[1].split("]")[0]
+                for ocr_img in rel_images:
+                    if ocr_img.id == img_id:
+                        chunk.metadata.images.append(
+                            Image(id=img_id, image_base64=ocr_img.image_base64)
+                        )
+                        break
+
+        return chunks
+
+    def __str__(self) -> str:
+        """Return a string representation of the chunker."""
+        return self.__class__.__name__
+
+    def __repr__(self) -> str:
+        """Return a string representation of the chunker."""
+        return f"{self.__class__.__name__}()"

ragbandit/documents/chunkers/fixed_size_chunker.py

@@ -0,0 +1,174 @@
+from datetime import datetime, timezone
+
+from ragbandit.schema import (
+    ProcessingResult,
+    Chunk,
+    ChunkMetadata,
+    ChunkingResult,
+)
+from ragbandit.utils.token_usage_tracker import TokenUsageTracker
+from ragbandit.documents.chunkers.base_chunker import BaseChunker
+
+
+class FixedSizeChunker(BaseChunker):
+    """
+    A document chunker that splits documents into fixed-size chunks
+    with optional overlap between chunks.
+    """
+
+    def __init__(
+        self,
+        chunk_size: int = 1000,
+        overlap: int = 200,
+        name: str | None = None,
+    ):
+        """
+        Initialize the fixed size chunker.
+
+        Args:
+            chunk_size: Target size for each chunk in characters
+            overlap: Number of characters to overlap between chunks
+            name: Optional name for the chunker
+        """
+        super().__init__(name)
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+
+    def chunk(
+        self,
+        proc_result: ProcessingResult,
+        usage_tracker: TokenUsageTracker | None = None,
+    ) -> ChunkingResult:
+        """
+        Chunk the document into fixed-size chunks.
+
+        Args:
+            proc_result: The ProcessingResult containing
+                      document content to chunk
+            usage_tracker: Optional tracker for token usage
+                           (not used in this chunker)
+
+        Returns:
+            A ChunkingResult containing Chunk objects
+        """
+        # 1. Generate raw chunks for each page
+        chunks = self._fixed_size_chunk_pages(proc_result)
+
+        # 2. Attach any inline images using BaseChunker helper
+        chunks = self.attach_images(chunks, proc_result)
+
+        # 3. Merge small chunks if needed
+        chunks = self.process_chunks(chunks)
+
+        # 4. Wrap in ChunkingResult
+        return ChunkingResult(
+            processed_at=datetime.now(timezone.utc),
+            chunks=chunks,
+            metrics=None,
+        )
+
+    # ------------------------------------------------------------------
+    # Internal helpers
+    def _fixed_size_chunk_pages(
+        self, proc_result: ProcessingResult
+    ) -> list[Chunk]:
+        """Split each page into fixed-size chunks with optional overlap."""
+
+        self.logger.info(
+            f"Starting fixed-size chunking with size={self.chunk_size}, "
+            f"overlap={self.overlap}"
+        )
+
+        chunks: list[Chunk] = []
+
+        # Process each page
+        for page_index, page in enumerate(proc_result.pages):
+            page_text = page.markdown
+
+            # Skip empty pages
+            if not page_text.strip():
+                continue
+
+            # Create chunks from this page
+            start = 0
+            while start < len(page_text):
+                # Determine end position for this chunk
+                end = min(start + self.chunk_size, len(page_text))
+
+                # If we're not at the end of the text,
+                # try to find a good break point
+                if end < len(page_text):
+                    # Look for a period, question mark, or exclamation mark
+                    # followed by whitespace
+                    # within the last 100 characters of the chunk
+                    search_start = max(end - 100, start)
+                    for i in range(end, search_start, -1):
+                        # Check if we're at a valid position to examine
+                        if i <= 0 or i >= len(page_text):
+                            continue
+
+                        # Check if the previous character is punctuation
+                        # and the current character is whitespace
+                        if (
+                            page_text[i - 1] in [".", "!", "?"]
+                            and page_text[i].isspace()
+                        ):
+                            end = i
+                            break
+
+                # Create the chunk
+                chunk_text = page_text[start:end]
+                meta = ChunkMetadata(
+                    page_index=page_index, images=[], extra={}
+                    )
+                chunks.append(Chunk(text=chunk_text, metadata=meta))
+
+                # Check if we've reached the end of the page text
+                if end >= len(page_text):
+                    # We've processed the entire page, exit the loop
+                    break
+
+                # Move to next chunk start position, accounting for overlap
+                start = end - self.overlap
+
+                # Make sure we're making progress
+                if start <= 0 or start >= len(page_text):
+                    break
+
+        self.logger.info(
+            f"Fixed-size chunking complete. Created {len(chunks)} chunks."
+        )
+
+        return chunks
+
+    def process_chunks(
+        self, chunks: list[Chunk]
+    ) -> list[Chunk]:
+        """
+        Process chunks after initial chunking - merge small chunks if needed.
+
+        Args:
+            chunks: The initial chunks produced by the chunk method
+
+        Returns:
+            Processed chunks with small chunks merged if needed
+        """
+        if not chunks:
+            return chunks
+
+        # Calculate minimum chunk size as a fraction of the target chunk size
+        min_chunk_size = self.chunk_size // 2
+
+        # Check if any chunks are too small
+        min_len = min(len(c.text) for c in chunks)
+
+        # Merge small chunks if needed
+        if min_len < min_chunk_size:
+            self.logger.info(
+                f"Found chunks smaller than {min_chunk_size} characters. "
+                "Merging..."
+            )
+            chunks = self.merge_small_chunks(chunks, min_size=min_chunk_size)
+            self.logger.info(f"After merging: {len(chunks)} chunks")
+
+        return chunks