ragbandit-core 0.1.1__py3-none-any.whl

ragbandit/documents/embedders/mistral_embedder.py

@@ -0,0 +1,129 @@
+from ragbandit.utils import mistral_client_manager
+from ragbandit.utils.token_usage_tracker import TokenUsageTracker
+from ragbandit.documents.embedders.base_embedder import BaseEmbedder
+from ragbandit.schema import (
+    ChunkingResult,
+    EmbeddingResult,
+    ChunkWithEmbedding,
+)
+from datetime import datetime, timezone
+
+
+class MistralEmbedder(BaseEmbedder):
+    """Document embedder that uses Mistral AI's embedding models."""
+
+    def __init__(
+        self,
+        api_key: str,
+        model: str = "mistral-embed",
+        name: str = None,
+    ):
+        """
+        Initialize the Mistral embedder.
+
+        Args:
+            api_key: Mistral API key
+            model: Embedding model to use
+            name: Optional name for the embedder
+        """
+        super().__init__(name)
+        self.model = model
+
+        # Initialize the Mistral client
+        self.client = mistral_client_manager.get_client(api_key)
+
+        self.logger.info(
+            f"Initialized MistralEmbedder with model {self.model}"
+        )
+
+    # ------------------------------------------------------------------
+    # Public API
+    def embed_chunks(
+        self,
+        chunk_result: ChunkingResult,
+        usage_tracker: TokenUsageTracker | None = None,
+    ) -> EmbeddingResult:
+        """Orchestrate the Mistral embedding flow."""
+
+        chunks = chunk_result.chunks
+        if not chunks:
+            self.logger.warning("No chunks to embed")
+            return self._empty_embedding_result()
+
+        try:
+            texts = self._extract_texts(chunks)
+            response = self._call_mistral_embeddings(texts)
+            return self._build_embedding_result(
+                chunks, response, usage_tracker
+            )
+        except Exception as e:
+            self.logger.error(f"Error generating embeddings: {e}")
+            return self._empty_embedding_result(chunks)
+
+    # ------------------------------------------------------------------
+    # Helpers
+    def _extract_texts(
+        self,
+        chunks: list[ChunkWithEmbedding | object],
+    ) -> list[str]:
+        """Extract raw text from chunks."""
+        return [c.text for c in chunks]
+
+    def _call_mistral_embeddings(self, texts: list[str]):
+        """Call the Mistral embeddings API and return the raw response."""
+        self.logger.info("Requesting embeddings from Mistral API")
+        return self.client.embeddings.create(model=self.model, inputs=texts)
+
+    def _build_embedding_result(
+        self,
+        chunks,
+        response,
+        usage_tracker: TokenUsageTracker | None = None,
+    ) -> EmbeddingResult:
+        """Convert API response into EmbeddingResult."""
+
+        embedded_chunks: list[ChunkWithEmbedding] = []
+        for i, data in enumerate(response.data):
+            embedded_chunks.append(
+                ChunkWithEmbedding(
+                    text=chunks[i].text,
+                    metadata=chunks[i].metadata,
+                    embedding=list(data.embedding),
+                    embedding_model=self.model,
+                )
+            )
+
+        if usage_tracker:
+            usage_tracker.add_embedding_tokens(
+                response.usage.prompt_tokens, self.model
+            )
+
+        return EmbeddingResult(
+            processed_at=datetime.now(timezone.utc),
+            chunks_with_embeddings=embedded_chunks,
+            model_name=self.model,
+            metrics=usage_tracker.get_summary() if usage_tracker else None,
+        )
+
+    def _empty_embedding_result(
+        self, chunks: list | None = None
+    ) -> EmbeddingResult:
+        """Return an EmbeddingResult with no embeddings (error fallback)."""
+
+        empty_embeds: list[ChunkWithEmbedding] = []
+        for c in (chunks or []):
+            empty_embeds.append(
+                ChunkWithEmbedding(
+                    text=c.text if hasattr(c, "text") else "",
+                    metadata=c.metadata if hasattr(c, "metadata") else None,
+                    embedding=[],
+                    embedding_model=self.model,
+                )
+            )
+
+        return EmbeddingResult(
+            processed_at=None,
+            chunks_with_embeddings=empty_embeds,
+            model_name=self.model,
+            metrics=None,
+        )

ragbandit/documents/ocr/__init__.py

@@ -0,0 +1,13 @@
+"""
+OCR (Optical Character Recognition) implementations for document processing.
+
+This module provides OCR processors that convert document images to text.
+"""
+
+from ragbandit.documents.ocr.base_ocr import BaseOCR
+from ragbandit.documents.ocr.mistral_ocr import MistralOCRDocument
+
+__all__ = [
+    "BaseOCR",
+    "MistralOCRDocument"
+]

ragbandit/documents/ocr/base_ocr.py

@@ -0,0 +1,136 @@
+import logging
+import os
+from abc import ABC, abstractmethod
+from pathlib import Path
+from io import BytesIO, BufferedReader
+from ragbandit.documents.utils.secure_file_handler import SecureFileHandler
+from ragbandit.schema import OCRResult
+
+
+class BaseOCR(ABC):
+    """Base class for OCR document processing.
+
+    This class provides the interface for OCR processing and a default
+    implementation using Mistral's OCR API.
+    """
+
+    def __init__(self, logger: logging.Logger = None, **kwargs):
+        """Initialize the OCR processor.
+
+        Args:
+            logger: Optional logger for OCR events
+            **kwargs: Additional keyword arguments (e.g., encryption_key)
+        """
+        self.logger = logger or logging.getLogger(__name__)
+        self.kwargs = kwargs
+
+    def validate_pdf(self, pdf_filepath: str) -> str:
+        """Validate that a PDF file exists.
+
+        Args:
+            pdf_filepath: Path to the PDF file to validate
+
+        Returns:
+            str: The basename of the file
+
+        Raises:
+            ValueError: If the file does not exist
+        """
+        file_name = os.path.basename(pdf_filepath)
+        pdf_file_exists = os.path.isfile(pdf_filepath)
+
+        if not pdf_file_exists:
+            self.logger.error(f"PDF file {pdf_filepath} not found")
+            raise ValueError(f"PDF file {pdf_filepath} not found")
+
+        return file_name
+
+    def read_encrypted_file(self, pdf_filepath: str) -> BufferedReader:
+        """Read an encrypted PDF file and return a buffered reader.
+
+        Args:
+            pdf_filepath: Path to the encrypted PDF file
+
+        Returns:
+            BufferedReader: A buffered reader for the decrypted file content
+
+        Raises:
+            ValueError: If encryption_key is not provided in kwargs
+        """
+        self.logger.info("Decrypting for OCR...")
+
+        encryption_key = self.kwargs.get("encryption_key")
+        if not encryption_key:
+            raise ValueError(
+                "encryption_key must be provided in kwargs "
+                "for encrypted file operations. "
+                "Pass encryption_key when initializing the OCR processor."
+            )
+
+        secure_handler = SecureFileHandler(encryption_key=encryption_key)
+        decrypted = secure_handler.read_encrypted_file(Path(pdf_filepath))
+        raw = BytesIO(decrypted)
+        raw.seek(0)
+        return BufferedReader(raw)
+
+    def read_unencrypted_file(self, pdf_filepath: str) -> BufferedReader:
+        """Read an unencrypted PDF file and return a buffered reader.
+
+        Args:
+            pdf_filepath: Path to the unencrypted PDF file
+
+        Returns:
+            BufferedReader: A buffered reader for the file content
+        """
+        self.logger.info("Reading file for OCR...")
+        with open(pdf_filepath, 'rb') as f:
+            content = f.read()
+
+        raw = BytesIO(content)
+        raw.seek(0)
+        return BufferedReader(raw)
+
+    def validate_and_prepare_file(
+        self, pdf_filepath: str, encrypted: bool = True
+    ) -> tuple[str, BufferedReader]:
+        """Validate and prepare a PDF file for OCR processing.
+
+        Args:
+            pdf_filepath: Path to the PDF file to process
+            encrypted: Whether the file is encrypted (default: True)
+
+        Returns:
+            tuple: (file_name, file_reader)
+
+        Raises:
+            ValueError: If the file does not exist
+        """
+        file_name = self.validate_pdf(pdf_filepath)
+
+        if encrypted:
+            reader = self.read_encrypted_file(pdf_filepath)
+        else:
+            reader = self.read_unencrypted_file(pdf_filepath)
+
+        return file_name, reader
+
+    @abstractmethod
+    def process(self, pdf_filepath: str) -> OCRResult:
+        """Process a PDF file through OCR.
+
+        Args:
+            pdf_filepath: Path to the PDF file to process
+
+        Returns:
+            OCRResult: The OCR result from the processor
+        """
+        raise NotImplementedError("Subclasses must implement process method")
+
+    # ----------------------------------------------------------------------
+    def __str__(self) -> str:
+        """Return a string representation of the OCR processor."""
+        return self.__class__.__name__
+
+    def __repr__(self) -> str:
+        """Return a string representation of the OCR processor."""
+        return f"{self.__class__.__name__}()"

ragbandit/documents/ocr/mistral_ocr.py

@@ -0,0 +1,147 @@
+from ragbandit.documents.ocr.base_ocr import BaseOCR
+from ragbandit.schema import (
+    OCRResult, OCRPage, PageDimensions,
+    Image, OCRUsageInfo, PagesProcessedMetrics
+)
+from ragbandit.config.pricing import OCR_MODEL_COSTS
+from mistralai.models.ocrresponse import OCRResponse
+from io import BufferedReader
+from datetime import datetime
+
+import logging
+from ragbandit.utils import mistral_client_manager
+
+
+class MistralOCRDocument(BaseOCR):
+    """OCR document processor using Mistral's API."""
+
+    # Explicit model identifier used for all Mistral OCR requests
+    MODEL_NAME = "mistral-ocr-latest"
+
+    def __init__(self, api_key: str, logger: logging.Logger = None, **kwargs):
+        """
+        Initialize the Mistral OCR processor.
+
+        Args:
+            api_key: Mistral API key
+            logger: Optional logger for OCR events
+            **kwargs: Additional keyword arguments
+                - encryption_key: Optional key for encrypted file operations
+        """
+        # Pass all kwargs to the base class
+        super().__init__(logger=logger, **kwargs)
+        self.client = mistral_client_manager.get_client(api_key)
+
+    # ----------------- Helper methods ----------------- #
+
+    def _upload_file(self, file_name: str, content: BufferedReader):
+        """Upload PDF to Mistral Cloud and return the uploaded file object."""
+        return self.client.files.upload(
+            file={"file_name": file_name, "content": content},
+            purpose="ocr",
+        )
+
+    def _get_signed_url(self, file_id: str) -> str:
+        """Retrieve a signed URL for the previously uploaded file."""
+        return self.client.files.get_signed_url(file_id=file_id).url
+
+    def _run_ocr(self, document_url: str) -> OCRResponse:
+        """Run Mistral OCR on the provided document URL."""
+        return self.client.ocr.process(
+            model=self.MODEL_NAME,
+            document={"type": "document_url", "document_url": document_url},
+            include_image_base64=True,
+        )
+
+    def _delete_file_with_retries(
+        self, file_id: str, max_tries: int = 10
+    ) -> None:
+        """Attempt to delete a file from Mistral Cloud with retries."""
+        while max_tries > 0:
+            resp = self.client.files.delete(file_id=file_id)
+            if resp.deleted:
+                self.logger.info("File deletion successful!")
+                return
+            max_tries -= 1
+        self.logger.error(f"Deleting unsuccessful. ID: {file_id}")
+
+    def _convert_pages(self, ocr_response: OCRResponse) -> list[OCRPage]:
+        """Convert OCRResponse pages to internal OCRPage schema."""
+        pages = []
+        for i, page in enumerate(ocr_response.pages):
+            images = [
+                Image.model_validate(img, from_attributes=True)
+                for img in (page.images or [])
+            ]
+            ocr_page = OCRPage(
+                index=i,
+                markdown=page.markdown,
+                images=images,
+                dimensions=PageDimensions.model_validate(
+                    page.dimensions, from_attributes=True
+                ),
+            )
+            pages.append(ocr_page)
+        return pages
+
+    def _build_usage_info(self, ocr_response: OCRResponse) -> OCRUsageInfo:
+        """Extract usage information from the OCR response."""
+        return OCRUsageInfo.model_validate(
+            ocr_response.usage_info, from_attributes=True
+        )
+
+    def _build_metrics(self, pages: list[OCRPage]) -> PagesProcessedMetrics:
+        """Create page-processing cost metrics."""
+        cost_per_page = OCR_MODEL_COSTS.get(self.MODEL_NAME, 0.0)
+        return PagesProcessedMetrics(
+            pages_processed=len(pages),
+            cost_per_page=cost_per_page,
+            total_cost_usd=len(pages) * cost_per_page,
+        )
+
+    def _build_result(
+        self,
+        pdf_filepath: str,
+        ocr_response: OCRResponse,
+        pages: list[OCRPage],
+        usage_info: OCRUsageInfo,
+        metrics: list[PagesProcessedMetrics],
+    ) -> OCRResult:
+        """Assemble the OCRResult object."""
+        return OCRResult(
+            source_file_path=pdf_filepath,
+            processed_at=datetime.now(),
+            model=ocr_response.model,
+            document_annotation=ocr_response.document_annotation,
+            pages=pages,
+            usage_info=usage_info,
+            metrics=metrics,
+        )
+
+    def process(self, pdf_filepath: str, encrypted: bool = False) -> OCRResult:
+        """High-level orchestration for running Mistral OCR on a PDF."""
+
+        file_name, reader = self.validate_and_prepare_file(
+                                pdf_filepath, encrypted
+                            )
+
+        uploaded = self._upload_file(file_name, reader)
+        del reader  # free memory
+
+        try:
+            doc_url = self._get_signed_url(uploaded.id)
+            ocr_resp = self._run_ocr(doc_url)
+        finally:
+            self._delete_file_with_retries(uploaded.id)
+
+        pages = self._convert_pages(ocr_resp)
+        usage_info = self._build_usage_info(ocr_resp)
+        metrics = [self._build_metrics(pages)]
+
+        return self._build_result(
+            pdf_filepath=pdf_filepath,
+            ocr_response=ocr_resp,
+            pages=pages,
+            usage_info=usage_info,
+            metrics=metrics,
+        )

ragbandit/documents/processors/__init__.py

@@ -0,0 +1,16 @@
+"""
+Document processors for enhancing and transforming document content.
+
+This module provides various processors that can be applied to documents
+to extract, transform, or enhance their content.
+"""
+
+from ragbandit.documents.processors.base_processor import BaseProcessor
+from ragbandit.documents.processors.footnotes_processor import FootnoteProcessor  # noqa
+from ragbandit.documents.processors.references_processor import ReferencesProcessor  # noqa
+
+__all__ = [
+    "BaseProcessor",
+    "FootnoteProcessor",
+    "ReferencesProcessor"
+]

ragbandit/documents/processors/base_processor.py

@@ -0,0 +1,88 @@
+import logging
+from abc import ABC, abstractmethod
+from datetime import datetime, timezone
+from ragbandit.schema import OCRResult, ProcessingResult, ProcessedPage
+from ragbandit.utils.token_usage_tracker import TokenUsageTracker
+
+
+class BaseProcessor(ABC):
+    """
+    Base class or mix-in for individual processors.
+    Subclasses override `process()` and, optionally, `extend_response()`.
+    """
+
+    def __init__(self, name: str | None = None, api_key: str | None = None):
+        """
+        Initialize the processor.
+
+        Args:
+            name: Optional name for the processor
+            api_key: API key for LLM services
+        """
+        # Hierarchical names make it easy to filter later:
+        #   pipeline.text_cleaner, pipeline.language_model, …
+        base = "pipeline"
+        self.logger = logging.getLogger(
+            f"{base}.{name or self.__class__.__name__}"
+        )
+        self.api_key = api_key
+
+    @abstractmethod
+    def process(
+        self,
+        document: OCRResult | ProcessingResult,
+        usage_tracker: TokenUsageTracker | None = None,
+    ) -> ProcessingResult:
+        """
+        Do one step of work and return:
+          * a (possibly modified) ProcessingResult
+          * a dict of metadata specific to this processor
+
+        Args:
+            response: The OCR or intermediate processing result to process.
+                This can be either an `OCRResult` (raw OCR output) or
+                a `ProcessingResult` (output of a previous processor).
+            usage_tracker: Optional token usage tracker
+        """
+        raise NotImplementedError
+
+    # ----------------------------------------------------------------------
+    def __str__(self) -> str:
+        """Return a string representation of the processor."""
+        return self.__class__.__name__
+
+    def __repr__(self) -> str:
+        """Return a string representation of the processor."""
+        return f"{self.__class__.__name__}()"
+
+    # ----------------------------------------------------------------------
+    # Utility helpers
+    @staticmethod
+    def ensure_processing_result(
+        document: OCRResult | ProcessingResult,
+        processor_name: str = "bootstrap",
+    ) -> ProcessingResult:
+        """Ensure the incoming `document` is a `ProcessingResult`.
+
+        If an `OCRResult` is supplied (as is the case for the first processor
+        in a pipeline), it will be converted to a shallow `ProcessingResult` so
+        that downstream logic can assume a consistent type.
+        """
+
+        # Always create a fresh ProcessingResult so that timestamps, metrics,
+        # and extracted data do not roll over between processors.
+
+        source_pages = document.pages if hasattr(document, "pages") else []
+
+        pages_processed = [
+            ProcessedPage(**page.model_dump()) for page in source_pages
+        ]
+
+        return ProcessingResult(
+            processor_name=processor_name,
+            processed_at=datetime.now(timezone.utc),
+            pages=pages_processed,
+            processing_trace=[],
+            extracted_data={},
+            metrics=None,
+        )