Parcourir la source

enh: configurable mistral ocr base url

Timothy Jaeryang Baek il y a 3 mois
Parent
commit
415b93c7c3

+ 6 - 0
backend/open_webui/config.py

@@ -2464,6 +2464,12 @@ DOCUMENT_INTELLIGENCE_KEY = PersistentConfig(
     os.getenv("DOCUMENT_INTELLIGENCE_KEY", ""),
 )
 
+MISTRAL_OCR_API_BASE_URL = PersistentConfig(
+    "MISTRAL_OCR_API_BASE_URL",
+    "rag.MISTRAL_OCR_API_BASE_URL",
+    os.getenv("MISTRAL_OCR_API_BASE_URL", "https://api.mistral.ai/v1"),
+)
+
 MISTRAL_OCR_API_KEY = PersistentConfig(
     "MISTRAL_OCR_API_KEY",
     "rag.mistral_ocr_api_key",

+ 2 - 0
backend/open_webui/main.py

@@ -276,6 +276,7 @@ from open_webui.config import (
     DOCLING_PICTURE_DESCRIPTION_API,
     DOCUMENT_INTELLIGENCE_ENDPOINT,
     DOCUMENT_INTELLIGENCE_KEY,
+    MISTRAL_OCR_API_BASE_URL,
     MISTRAL_OCR_API_KEY,
     RAG_TEXT_SPLITTER,
     TIKTOKEN_ENCODING_NAME,
@@ -868,6 +869,7 @@ app.state.config.DOCLING_PICTURE_DESCRIPTION_LOCAL = DOCLING_PICTURE_DESCRIPTION
 app.state.config.DOCLING_PICTURE_DESCRIPTION_API = DOCLING_PICTURE_DESCRIPTION_API
 app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT = DOCUMENT_INTELLIGENCE_ENDPOINT
 app.state.config.DOCUMENT_INTELLIGENCE_KEY = DOCUMENT_INTELLIGENCE_KEY
+app.state.config.MISTRAL_OCR_API_BASE_URL = MISTRAL_OCR_API_BASE_URL
 app.state.config.MISTRAL_OCR_API_KEY = MISTRAL_OCR_API_KEY
 app.state.config.MINERU_API_MODE = MINERU_API_MODE
 app.state.config.MINERU_API_URL = MINERU_API_URL

+ 3 - 1
backend/open_webui/retrieval/loaders/main.py

@@ -384,7 +384,9 @@ class Loader:
             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
+                base_url=self.kwargs.get("MISTRAL_OCR_API_BASE_URL"),
+                api_key=self.kwargs.get("MISTRAL_OCR_API_KEY"),
+                file_path=file_path,
             )
         else:
             if file_ext == "pdf":

+ 12 - 10
backend/open_webui/retrieval/loaders/mistral.py

@@ -30,10 +30,9 @@ class MistralLoader:
     - Enhanced error handling with retryable error classification
     """
 
-    BASE_API_URL = "https://api.mistral.ai/v1"
-
     def __init__(
         self,
+        base_url: str,
         api_key: str,
         file_path: str,
         timeout: int = 300,  # 5 minutes default
@@ -55,6 +54,9 @@ class MistralLoader:
         if not os.path.exists(file_path):
             raise FileNotFoundError(f"File not found at {file_path}")
 
+        self.base_url = (
+            base_url.rstrip("/") if base_url else "https://api.mistral.ai/v1"
+        )
         self.api_key = api_key
         self.file_path = file_path
         self.timeout = timeout
@@ -240,7 +242,7 @@ class MistralLoader:
         in a context manager to minimize memory usage duration.
         """
         log.info("Uploading file to Mistral API")
-        url = f"{self.BASE_API_URL}/files"
+        url = f"{self.base_url}/files"
 
         def upload_request():
             # MEMORY OPTIMIZATION: Use context manager to minimize file handle lifetime
@@ -275,7 +277,7 @@ class MistralLoader:
 
     async def _upload_file_async(self, session: aiohttp.ClientSession) -> str:
         """Async file upload with streaming for better memory efficiency."""
-        url = f"{self.BASE_API_URL}/files"
+        url = f"{self.base_url}/files"
 
         async def upload_request():
             # Create multipart writer for streaming upload
@@ -321,7 +323,7 @@ class MistralLoader:
     def _get_signed_url(self, file_id: str) -> str:
         """Retrieves a temporary signed URL for the uploaded file (sync version)."""
         log.info(f"Getting signed URL for file ID: {file_id}")
-        url = f"{self.BASE_API_URL}/files/{file_id}/url"
+        url = f"{self.base_url}/files/{file_id}/url"
         params = {"expiry": 1}
         signed_url_headers = {**self.headers, "Accept": "application/json"}
 
@@ -346,7 +348,7 @@ class MistralLoader:
         self, session: aiohttp.ClientSession, file_id: str
     ) -> str:
         """Async signed URL retrieval."""
-        url = f"{self.BASE_API_URL}/files/{file_id}/url"
+        url = f"{self.base_url}/files/{file_id}/url"
         params = {"expiry": 1}
 
         headers = {**self.headers, "Accept": "application/json"}
@@ -373,7 +375,7 @@ class MistralLoader:
     def _process_ocr(self, signed_url: str) -> Dict[str, Any]:
         """Sends the signed URL to the OCR endpoint for processing (sync version)."""
         log.info("Processing OCR via Mistral API")
-        url = f"{self.BASE_API_URL}/ocr"
+        url = f"{self.base_url}/ocr"
         ocr_headers = {
             **self.headers,
             "Content-Type": "application/json",
@@ -407,7 +409,7 @@ class MistralLoader:
         self, session: aiohttp.ClientSession, signed_url: str
     ) -> Dict[str, Any]:
         """Async OCR processing with timing metrics."""
-        url = f"{self.BASE_API_URL}/ocr"
+        url = f"{self.base_url}/ocr"
 
         headers = {
             **self.headers,
@@ -446,7 +448,7 @@ class MistralLoader:
     def _delete_file(self, file_id: str) -> None:
         """Deletes the file from Mistral storage (sync version)."""
         log.info(f"Deleting uploaded file ID: {file_id}")
-        url = f"{self.BASE_API_URL}/files/{file_id}"
+        url = f"{self.base_url}/files/{file_id}"
 
         try:
             response = requests.delete(
@@ -467,7 +469,7 @@ class MistralLoader:
             async def delete_request():
                 self._debug_log(f"Deleting file ID: {file_id}")
                 async with session.delete(
-                    url=f"{self.BASE_API_URL}/files/{file_id}",
+                    url=f"{self.base_url}/files/{file_id}",
                     headers=self.headers,
                     timeout=aiohttp.ClientTimeout(
                         total=self.cleanup_timeout

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

@@ -465,6 +465,7 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
         "DOCLING_PICTURE_DESCRIPTION_API": request.app.state.config.DOCLING_PICTURE_DESCRIPTION_API,
         "DOCUMENT_INTELLIGENCE_ENDPOINT": request.app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT,
         "DOCUMENT_INTELLIGENCE_KEY": request.app.state.config.DOCUMENT_INTELLIGENCE_KEY,
+        "MISTRAL_OCR_API_BASE_URL": request.app.state.config.MISTRAL_OCR_API_BASE_URL,
         "MISTRAL_OCR_API_KEY": request.app.state.config.MISTRAL_OCR_API_KEY,
         # MinerU settings
         "MINERU_API_MODE": request.app.state.config.MINERU_API_MODE,
@@ -650,6 +651,7 @@ class ConfigForm(BaseModel):
     DOCLING_PICTURE_DESCRIPTION_API: Optional[dict] = None
     DOCUMENT_INTELLIGENCE_ENDPOINT: Optional[str] = None
     DOCUMENT_INTELLIGENCE_KEY: Optional[str] = None
+    MISTRAL_OCR_API_BASE_URL: Optional[str] = None
     MISTRAL_OCR_API_KEY: Optional[str] = None
 
     # MinerU settings
@@ -891,6 +893,12 @@ async def update_rag_config(
         if form_data.DOCUMENT_INTELLIGENCE_KEY is not None
         else request.app.state.config.DOCUMENT_INTELLIGENCE_KEY
     )
+
+    request.app.state.config.MISTRAL_OCR_API_BASE_URL = (
+        form_data.MISTRAL_OCR_API_BASE_URL
+        if form_data.MISTRAL_OCR_API_BASE_URL is not None
+        else request.app.state.config.MISTRAL_OCR_API_BASE_URL
+    )
     request.app.state.config.MISTRAL_OCR_API_KEY = (
         form_data.MISTRAL_OCR_API_KEY
         if form_data.MISTRAL_OCR_API_KEY is not None
@@ -1182,6 +1190,7 @@ async def update_rag_config(
         "DOCLING_PICTURE_DESCRIPTION_API": request.app.state.config.DOCLING_PICTURE_DESCRIPTION_API,
         "DOCUMENT_INTELLIGENCE_ENDPOINT": request.app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT,
         "DOCUMENT_INTELLIGENCE_KEY": request.app.state.config.DOCUMENT_INTELLIGENCE_KEY,
+        "MISTRAL_OCR_API_BASE_URL": request.app.state.config.MISTRAL_OCR_API_BASE_URL,
         "MISTRAL_OCR_API_KEY": request.app.state.config.MISTRAL_OCR_API_KEY,
         # MinerU settings
         "MINERU_API_MODE": request.app.state.config.MINERU_API_MODE,
@@ -1597,6 +1606,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_BASE_URL=request.app.state.config.MISTRAL_OCR_API_BASE_URL,
                         MISTRAL_OCR_API_KEY=request.app.state.config.MISTRAL_OCR_API_KEY,
                         MINERU_API_MODE=request.app.state.config.MINERU_API_MODE,
                         MINERU_API_URL=request.app.state.config.MINERU_API_URL,

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

@@ -766,6 +766,11 @@
 							</div>
 						{:else if RAGConfig.CONTENT_EXTRACTION_ENGINE === 'mistral_ocr'}
 							<div class="my-0.5 flex gap-2 pr-2">
+								<input
+									class="flex-1 w-full text-sm bg-transparent outline-hidden"
+									placeholder={$i18n.t('Enter Mistral API Base URL')}
+									bind:value={RAGConfig.MISTRAL_OCR_API_BASE_URL}
+								/>
 								<SensitiveInput
 									placeholder={$i18n.t('Enter Mistral API Key')}
 									bind:value={RAGConfig.MISTRAL_OCR_API_KEY}