1
0
Эх сурвалжийг харах

Add Mistral OCR integration and configuration support

Patrick Wachter 1 сар өмнө
parent
commit
1ac6879268

+ 5 - 0
backend/open_webui/config.py

@@ -1727,6 +1727,11 @@ DOCUMENT_INTELLIGENCE_KEY = PersistentConfig(
     os.getenv("DOCUMENT_INTELLIGENCE_KEY", ""),
 )
 
+MISTRAL_OCR_API_KEY = PersistentConfig(
+    "MISTRAL_OCR_API_KEY",
+    "rag.mistral_ocr_api_key",
+    os.getenv("MISTRAL_OCR_API_KEY", ""),
+)
 
 BYPASS_EMBEDDING_AND_RETRIEVAL = PersistentConfig(
     "BYPASS_EMBEDDING_AND_RETRIEVAL",

+ 2 - 0
backend/open_webui/main.py

@@ -191,6 +191,7 @@ from open_webui.config import (
     DOCLING_SERVER_URL,
     DOCUMENT_INTELLIGENCE_ENDPOINT,
     DOCUMENT_INTELLIGENCE_KEY,
+    MISTRAL_OCR_API_KEY,
     RAG_TOP_K,
     RAG_TOP_K_RERANKER,
     RAG_TEXT_SPLITTER,
@@ -582,6 +583,7 @@ app.state.config.TIKA_SERVER_URL = TIKA_SERVER_URL
 app.state.config.DOCLING_SERVER_URL = DOCLING_SERVER_URL
 app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT = DOCUMENT_INTELLIGENCE_ENDPOINT
 app.state.config.DOCUMENT_INTELLIGENCE_KEY = DOCUMENT_INTELLIGENCE_KEY
+app.state.config.MISTRAL_OCR_API_KEY = MISTRAL_OCR_API_KEY
 
 app.state.config.TEXT_SPLITTER = RAG_TEXT_SPLITTER
 app.state.config.TIKTOKEN_ENCODING_NAME = TIKTOKEN_ENCODING_NAME

+ 59 - 0
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(

+ 16 - 0
backend/open_webui/routers/retrieval.py

@@ -364,6 +364,9 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
                 "endpoint": request.app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT,
                 "key": request.app.state.config.DOCUMENT_INTELLIGENCE_KEY,
             },
+            "mistral_ocr_config": {
+                "api_key": request.app.state.config.MISTRAL_OCR_API_KEY,
+            },
         },
         "chunk": {
             "text_splitter": request.app.state.config.TEXT_SPLITTER,
@@ -427,11 +430,16 @@ class DocumentIntelligenceConfigForm(BaseModel):
     key: str
 
 
+class MistralOCRConfigForm(BaseModel):
+    api_key: str
+
+
 class ContentExtractionConfig(BaseModel):
     engine: str = ""
     tika_server_url: Optional[str] = None
     docling_server_url: Optional[str] = None
     document_intelligence_config: Optional[DocumentIntelligenceConfigForm] = None
+    mistral_ocr_config: Optional[MistralOCRConfigForm] = None
 
 
 class ChunkParamUpdateForm(BaseModel):
@@ -553,6 +561,10 @@ async def update_rag_config(
             request.app.state.config.DOCUMENT_INTELLIGENCE_KEY = (
                 form_data.content_extraction.document_intelligence_config.key
             )
+        if form_data.content_extraction.mistral_ocr_config is not None:
+            request.app.state.config.MISTRAL_OCR_API_KEY = (
+                form_data.content_extraction.mistral_ocr_config.api_key
+            )
 
     if form_data.chunk is not None:
         request.app.state.config.TEXT_SPLITTER = form_data.chunk.text_splitter
@@ -659,6 +671,9 @@ async def update_rag_config(
                 "endpoint": request.app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT,
                 "key": request.app.state.config.DOCUMENT_INTELLIGENCE_KEY,
             },
+            "mistral_ocr_config": {
+                "api_key": request.app.state.config.MISTRAL_OCR_API_KEY,
+            },
         },
         "chunk": {
             "text_splitter": request.app.state.config.TEXT_SPLITTER,
@@ -1007,6 +1022,7 @@ def process_file(
                     PDF_EXTRACT_IMAGES=request.app.state.config.PDF_EXTRACT_IMAGES,
                     DOCUMENT_INTELLIGENCE_ENDPOINT=request.app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT,
                     DOCUMENT_INTELLIGENCE_KEY=request.app.state.config.DOCUMENT_INTELLIGENCE_KEY,
+                    MISTRAL_OCR_API_KEY=request.app.state.config.MISTRAL_OCR_API_KEY,
                 )
                 docs = loader.load(
                     file.filename, file.meta.get("content_type"), file_path

+ 1 - 0
backend/requirements.txt

@@ -77,6 +77,7 @@ psutil
 sentencepiece
 soundfile==0.13.1
 azure-ai-documentintelligence==1.0.0
+mistralai==1.6.0
 
 pillow==11.1.0
 opencv-python-headless==4.11.0.86

+ 22 - 5
src/lib/components/admin/Settings/Documents.svelte

@@ -54,6 +54,8 @@
 	let documentIntelligenceEndpoint = '';
 	let documentIntelligenceKey = '';
 	let showDocumentIntelligenceConfig = false;
+	let mistralApiKey = '';
+	let showMistralOcrConfig = false;
 
 	let textSplitter = '';
 	let chunkSize = 0;
@@ -189,6 +191,10 @@
 			toast.error($i18n.t('Document Intelligence endpoint and key required.'));
 			return;
 		}
+		if (contentExtractionEngine === 'mistral_ocr' && mistralApiKey === '') {
+			toast.error($i18n.t('Mistral OCR API Key required.'));
+			return;
+		}
 
 		if (!BYPASS_EMBEDDING_AND_RETRIEVAL) {
 			await embeddingModelUpdateHandler();
@@ -220,6 +226,9 @@
 				document_intelligence_config: {
 					key: documentIntelligenceKey,
 					endpoint: documentIntelligenceEndpoint
+				},
+				mistral_ocr_config: {
+					api_key: mistralApiKey
 				}
 			}
 		});
@@ -284,6 +293,8 @@
 			documentIntelligenceEndpoint = res.content_extraction.document_intelligence_config.endpoint;
 			documentIntelligenceKey = res.content_extraction.document_intelligence_config.key;
 			showDocumentIntelligenceConfig = contentExtractionEngine === 'document_intelligence';
+			mistralApiKey = res.content_extraction.mistral_ocr_config.api_key;
+			showMistralOcrConfig = contentExtractionEngine === 'mistral_ocr';
 
 			fileMaxSize = res?.file.max_size ?? '';
 			fileMaxCount = res?.file.max_count ?? '';
@@ -335,21 +346,21 @@
 
 				<hr class=" border-gray-100 dark:border-gray-850 my-2" />
 
-				<div class="  mb-2.5 flex flex-col w-full justify-between">
+				<div class="mb-2.5 flex flex-col w-full justify-between">
 					<div class="flex w-full justify-between">
-						<div class=" self-center text-xs font-medium">
+						<div class="self-center text-xs font-medium">
 							{$i18n.t('Content Extraction Engine')}
 						</div>
-
 						<div class="">
 							<select
 								class="dark:bg-gray-900 w-fit pr-8 rounded-sm px-2 text-xs bg-transparent outline-hidden text-right"
 								bind:value={contentExtractionEngine}
 							>
-								<option value="">{$i18n.t('Default')} </option>
+								<option value="">{$i18n.t('Default')}</option>
 								<option value="tika">{$i18n.t('Tika')}</option>
 								<option value="docling">{$i18n.t('Docling')}</option>
 								<option value="document_intelligence">{$i18n.t('Document Intelligence')}</option>
+								<option value="mistral_ocr">{$i18n.t('Mistral OCR')}</option>
 							</select>
 						</div>
 					</div>
@@ -378,12 +389,18 @@
 								placeholder={$i18n.t('Enter Document Intelligence Endpoint')}
 								bind:value={documentIntelligenceEndpoint}
 							/>
-
 							<SensitiveInput
 								placeholder={$i18n.t('Enter Document Intelligence Key')}
 								bind:value={documentIntelligenceKey}
 							/>
 						</div>
+					{:else if contentExtractionEngine === 'mistral_ocr'}
+						<div class="my-0.5 flex gap-2 pr-2">
+							<SensitiveInput
+								placeholder={$i18n.t('Enter Mistral API Key')}
+								bind:value={mistralApiKey}
+							/>
+						</div>
 					{/if}
 				</div>