open-webui/backend/open_webui/retrieval/loaders/mistral.py

import requests
import logging
import os
import sys
from typing import List, Dict, Any

from langchain_core.documents import Document
from open_webui.env import SRC_LOG_LEVELS, GLOBAL_LOG_LEVEL

logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])


class MistralLoader:
    """
    Loads documents by processing them through the Mistral OCR API.
    """

    BASE_API_URL = "https://api.mistral.ai/v1"

    def __init__(self, api_key: str, file_path: str):
        """
        Initializes the loader.

        Args:
            api_key: Your Mistral API key.
            file_path: The local path to the PDF file to process.
        """
        if not api_key:
            raise ValueError("API key cannot be empty.")
        if not os.path.exists(file_path):
            raise FileNotFoundError(f"File not found at {file_path}")

        self.api_key = api_key
        self.file_path = file_path
        self.headers = {"Authorization": f"Bearer {self.api_key}"}

    def _handle_response(self, response: requests.Response) -> Dict[str, Any]:
        """Checks response status and returns JSON content."""
        try:
            response.raise_for_status()  # Raises HTTPError for bad responses (4xx or 5xx)
            # Handle potential empty responses for certain successful requests (e.g., DELETE)
            if response.status_code == 204 or not response.content:
                return {}  # Return empty dict if no content
            return response.json()
        except requests.exceptions.HTTPError as http_err:
            log.error(f"HTTP error occurred: {http_err} - Response: {response.text}")
            raise
        except requests.exceptions.RequestException as req_err:
            log.error(f"Request exception occurred: {req_err}")
            raise
        except ValueError as json_err:  # Includes JSONDecodeError
            log.error(f"JSON decode error: {json_err} - Response: {response.text}")
            raise  # Re-raise after logging

    def _upload_file(self) -> str:
        """Uploads the file to Mistral for OCR processing."""
        log.info("Uploading file to Mistral API")
        url = f"{self.BASE_API_URL}/files"
        file_name = os.path.basename(self.file_path)

        try:
            with open(self.file_path, "rb") as f:
                files = {"file": (file_name, f, "application/pdf")}
                data = {"purpose": "ocr"}

                upload_headers = self.headers.copy()  # Avoid modifying self.headers

                response = requests.post(
                    url, headers=upload_headers, files=files, data=data
                )

            response_data = self._handle_response(response)
            file_id = response_data.get("id")
            if not file_id:
                raise ValueError("File ID not found in upload response.")
            log.info(f"File uploaded successfully. File ID: {file_id}")
            return file_id
        except Exception as e:
            log.error(f"Failed to upload file: {e}")
            raise

    def _get_signed_url(self, file_id: str) -> str:
        """Retrieves a temporary signed URL for the uploaded file."""
        log.info(f"Getting signed URL for file ID: {file_id}")
        url = f"{self.BASE_API_URL}/files/{file_id}/url"
        params = {"expiry": 1}
        signed_url_headers = {**self.headers, "Accept": "application/json"}

        try:
            response = requests.get(url, headers=signed_url_headers, params=params)
            response_data = self._handle_response(response)
            signed_url = response_data.get("url")
            if not signed_url:
                raise ValueError("Signed URL not found in response.")
            log.info("Signed URL received.")
            return signed_url
        except Exception as e:
            log.error(f"Failed to get signed URL: {e}")
            raise

    def _process_ocr(self, signed_url: str) -> Dict[str, Any]:
        """Sends the signed URL to the OCR endpoint for processing."""
        log.info("Processing OCR via Mistral API")
        url = f"{self.BASE_API_URL}/ocr"
        ocr_headers = {
            **self.headers,
            "Content-Type": "application/json",
            "Accept": "application/json",
        }
        payload = {
            "model": "mistral-ocr-latest",
            "document": {
                "type": "document_url",
                "document_url": signed_url,
            },
            "include_image_base64": False,
        }

        try:
            response = requests.post(url, headers=ocr_headers, json=payload)
            ocr_response = self._handle_response(response)
            log.info("OCR processing done.")
            log.debug("OCR response: %s", ocr_response)
            return ocr_response
        except Exception as e:
            log.error(f"Failed during OCR processing: {e}")
            raise

    def _delete_file(self, file_id: str) -> None:
        """Deletes the file from Mistral storage."""
        log.info(f"Deleting uploaded file ID: {file_id}")
        url = f"{self.BASE_API_URL}/files/{file_id}"
        # No specific Accept header needed, default or Authorization is usually sufficient

        try:
            response = requests.delete(url, headers=self.headers)
            delete_response = self._handle_response(
                response
            )  # Check status, ignore response body unless needed
            log.info(
                f"File deleted successfully: {delete_response}"
            )  # Log the response if available
        except Exception as e:
            # Log error but don't necessarily halt execution if deletion fails
            log.error(f"Failed to delete file ID {file_id}: {e}")
            # Depending on requirements, you might choose to raise the error here

    def load(self) -> List[Document]:
        """
        Executes the full OCR workflow: upload, get URL, process OCR, delete file.

        Returns:
            A list of Document objects, one for each page processed.
        """
        file_id = None
        try:
            # 1. Upload file
            file_id = self._upload_file()

            # 2. Get Signed URL
            signed_url = self._get_signed_url(file_id)

            # 3. Process OCR
            ocr_response = self._process_ocr(signed_url)

            # 4. Process results
            pages_data = ocr_response.get("pages")
            if not pages_data:
                log.warning("No pages found in OCR response.")
                return [Document(page_content="No text content found", metadata={})]

            documents = []
            total_pages = len(pages_data)
            for page_data in pages_data:
                page_content = page_data.get("markdown")
                page_index = page_data.get("index")  # API uses 0-based index

                if page_content is not None and page_index is not None:
                    documents.append(
                        Document(
                            page_content=page_content,
                            metadata={
                                "page": page_index,  # 0-based index from API
                                "page_label": page_index
                                + 1,  # 1-based label for convenience
                                "total_pages": total_pages,
                                # Add other relevant metadata from page_data if available/needed
                                # e.g., page_data.get('width'), page_data.get('height')
                            },
                        )
                    )
                else:
                    log.warning(
                        f"Skipping page due to missing 'markdown' or 'index'. Data: {page_data}"
                    )

            if not documents:
                # Case where pages existed but none had valid markdown/index
                log.warning(
                    "OCR response contained pages, but none had valid content/index."
                )
                return [
                    Document(
                        page_content="No text content found in valid pages", metadata={}
                    )
                ]

            return documents

        except Exception as e:
            log.error(f"An error occurred during the loading process: {e}")
            # Return an empty list or a specific error document on failure
            return [Document(page_content=f"Error during processing: {e}", metadata={})]
        finally:
            # 5. Delete file (attempt even if prior steps failed after upload)
            if file_id:
                try:
                    self._delete_file(file_id)
                except Exception as del_e:
                    # Log deletion error, but don't overwrite original error if one occurred
                    log.error(
                        f"Cleanup error: Could not delete file ID {file_id}. Reason: {del_e}"
                    )