Add Mistral OCR integration and configuration support

backend/open_webui/retrieval/loaders/main.py

@@ -20,6 +20,9 @@ from langchain_community.document_loaders import (
     YoutubeLoader,
 )
 from langchain_core.documents import Document
+
+from mistralai import Mistral
+
 from open_webui.env import SRC_LOG_LEVELS, GLOBAL_LOG_LEVEL
 
 logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL)
@@ -163,6 +166,53 @@ class DoclingLoader:
             raise Exception(f"Error calling Docling: {error_msg}")
 
 
+class MistralLoader:
+    def __init__(self, api_key: str, file_path: str):
+        self.api_key = api_key
+        self.file_path = file_path
+        self.client = Mistral(api_key=api_key)
+
+    def load(self) -> list[Document]:
+        log.info("Uploading file to Mistral OCR")
+        uploaded_pdf = self.client.files.upload(
+            file={
+                "file_name": self.file_path.split("/")[-1],
+                "content": open(self.file_path, "rb"),
+            },
+            purpose="ocr",
+        )
+        log.info("File uploaded to Mistral OCR, getting signed URL")
+        signed_url = self.client.files.get_signed_url(file_id=uploaded_pdf.id)
+        log.info("Signed URL received, processing OCR")
+        ocr_response = self.client.ocr.process(
+            model="mistral-ocr-latest",
+            document={
+                "type": "document_url",
+                "document_url": signed_url.url,
+            },
+        )
+        log.info("OCR processing done, deleting uploaded file")
+        deleted_pdf = self.client.files.delete(file_id=uploaded_pdf.id)
+        log.info("Uploaded file deleted")
+        log.debug("OCR response: %s", ocr_response)
+        if not hasattr(ocr_response, "pages") or not ocr_response.pages:
+            log.error("No pages found in OCR response")
+            return [Document(page_content="No text content found", metadata={})]
+
+        return [
+            Document(
+                page_content=page.markdown,
+                metadata={
+                    "page": page.index,
+                    "page_label": page.index + 1,
+                    "total_pages": len(ocr_response.pages),
+                },
+            )
+            for page in ocr_response.pages
+            if hasattr(page, "markdown") and hasattr(page, "index")
+        ]
+
+
 class Loader:
     def __init__(self, engine: str = "", **kwargs):
         self.engine = engine
@@ -222,6 +272,15 @@ class Loader:
                 api_endpoint=self.kwargs.get("DOCUMENT_INTELLIGENCE_ENDPOINT"),
                 api_key=self.kwargs.get("DOCUMENT_INTELLIGENCE_KEY"),
             )
+        elif (
+            self.engine == "mistral_ocr"
+            and self.kwargs.get("MISTRAL_OCR_API_KEY") != ""
+            and file_ext
+            in ["pdf"]  # Mistral OCR currently only supports PDF and images
+        ):
+            loader = MistralLoader(
+                api_key=self.kwargs.get("MISTRAL_OCR_API_KEY"), file_path=file_path
+            )
         else:
             if file_ext == "pdf":
                 loader = PyPDFLoader(