scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-10-09 03:18:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

feat: add LiteLLM/Azure compatibility for Mistral OCR

Browse source
This commit is contained in:
KevinR 2025-12-22 19:48:06 +01:00
parent 117b1bfa65
commit 3ca208ce4b
6 changed files with 184 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
backend/open_webui/config.py
View file

@ -2619,6 +2619,18 @@ MISTRAL_OCR_API_KEY = PersistentConfig(
os.getenv("MISTRAL_OCR_API_KEY", ""),
)
MISTRAL_OCR_USE_BASE64 = PersistentConfig(
"MISTRAL_OCR_USE_BASE64",
"rag.mistral_ocr_use_base64",
os.environ.get("MISTRAL_OCR_USE_BASE64", "False").lower() == "true",
)
MISTRAL_OCR_MODEL = PersistentConfig(
"MISTRAL_OCR_MODEL",
"rag.mistral_ocr_model",
os.getenv("MISTRAL_OCR_MODEL", "mistral-ocr-latest"),
)
BYPASS_EMBEDDING_AND_RETRIEVAL = PersistentConfig(
"BYPASS_EMBEDDING_AND_RETRIEVAL",
"rag.bypass_embedding_and_retrieval",

4
backend/open_webui/main.py
View file

@ -279,6 +279,8 @@ from open_webui.config import (
DOCUMENT_INTELLIGENCE_MODEL,
MISTRAL_OCR_API_BASE_URL,
MISTRAL_OCR_API_KEY,
MISTRAL_OCR_USE_BASE64,
MISTRAL_OCR_MODEL,
RAG_TEXT_SPLITTER,
TIKTOKEN_ENCODING_NAME,
PDF_EXTRACT_IMAGES,
@ -879,6 +881,8 @@ app.state.config.DOCUMENT_INTELLIGENCE_KEY = DOCUMENT_INTELLIGENCE_KEY
app.state.config.DOCUMENT_INTELLIGENCE_MODEL = DOCUMENT_INTELLIGENCE_MODEL
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.MISTRAL_OCR_USE_BASE64 = MISTRAL_OCR_USE_BASE64
app.state.config.MISTRAL_OCR_MODEL = MISTRAL_OCR_MODEL
app.state.config.MINERU_API_MODE = MINERU_API_MODE
app.state.config.MINERU_API_URL = MINERU_API_URL
app.state.config.MINERU_API_KEY = MINERU_API_KEY

2
backend/open_webui/retrieval/loaders/main.py
View file

@ -357,6 +357,8 @@ class Loader:
base_url=self.kwargs.get("MISTRAL_OCR_API_BASE_URL"),
api_key=self.kwargs.get("MISTRAL_OCR_API_KEY"),
file_path=file_path,
use_base64=self.kwargs.get("MISTRAL_OCR_USE_BASE64", False),
model=self.kwargs.get("MISTRAL_OCR_MODEL", "mistral-ocr-latest"),
)
else:
if file_ext == "pdf":

133
backend/open_webui/retrieval/loaders/mistral.py
View file

@ -1,6 +1,7 @@
import requests
import aiohttp
import asyncio
import base64
import logging
import os
import sys
@ -37,6 +38,8 @@ class MistralLoader:
timeout: int = 300, # 5 minutes default
max_retries: int = 3,
enable_debug_logging: bool = False,
use_base64: bool = False,
model: str = "mistral-ocr-latest",
):
"""
Initializes the loader with enhanced features.
@ -47,6 +50,8 @@ class MistralLoader:
timeout: Request timeout in seconds.
max_retries: Maximum number of retry attempts.
enable_debug_logging: Enable detailed debug logs.
use_base64: Use base64 encoding instead of file upload.
model: The OCR model to use (default: mistral-ocr-latest).
"""
if not api_key:
raise ValueError("API key cannot be empty.")
@ -61,6 +66,8 @@ class MistralLoader:
self.timeout = timeout
self.max_retries = max_retries
self.debug = enable_debug_logging
self.use_base64 = use_base64
self.model = model if model else "mistral-ocr-latest"
# PERFORMANCE OPTIMIZATION: Differentiated timeouts for different operations
# This prevents long-running OCR operations from affecting quick operations
@ -444,6 +451,96 @@ class MistralLoader:
return await self._retry_request_async(ocr_request)
def _process_ocr_base64(self) -> Dict[str, Any]:
"""Process OCR using base64 encoded document (for some LLM Proxys compatibility).
Uses data URI format: data:application/pdf;base64,<base64_data>
"""
log.info(f"Processing OCR via base64 encoding with model: {self.model}")
url = f"{self.base_url}/ocr"
ocr_headers = {
**self.headers,
"Content-Type": "application/json",
"Accept": "application/json",
}
with open(self.file_path, "rb") as f:
file_content = base64.standard_b64encode(f.read()).decode("utf-8")
# Use data URI format with base64
data_uri = f"data:application/pdf;base64,{file_content}"
payload = {
"model": self.model,
"document": {
"type": "document_url",
"document_url": data_uri,
},
}
def ocr_request():
response = requests.post(
url, headers=ocr_headers, json=payload, timeout=self.ocr_timeout
)
return self._handle_response(response)
try:
ocr_response = self._retry_request_sync(ocr_request)
log.info("OCR processing done (base64 encoding).")
self._debug_log("OCR response: %s", ocr_response)
return ocr_response
except Exception as e:
log.error(f"Failed during OCR processing (base64 encoding): {e}")
raise
async def _process_ocr_base64_async(
self, session: aiohttp.ClientSession
) -> Dict[str, Any]:
"""Async OCR processing using base64 encoding (for some LLM Proxys compatibility).
Uses data URI format: data:application/pdf;base64,<base64_data>
"""
url = f"{self.base_url}/ocr"
headers = {
**self.headers,
"Content-Type": "application/json",
"Accept": "application/json",
}
with open(self.file_path, "rb") as f:
file_content = base64.standard_b64encode(f.read()).decode("utf-8")
# Use data URI format with base64
data_uri = f"data:application/pdf;base64,{file_content}"
payload = {
"model": self.model,
"document": {
"type": "document_url",
"document_url": data_uri,
},
}
async def ocr_request():
log.info(f"Starting OCR processing via base64 with model: {self.model}")
start_time = time.time()
async with session.post(
url,
json=payload,
headers=headers,
timeout=aiohttp.ClientTimeout(total=self.ocr_timeout),
) as response:
ocr_response = await self._handle_response_async(response)
processing_time = time.time() - start_time
log.info(f"OCR processing completed in {processing_time:.2f}s (base64 encoding)")
return ocr_response
return await self._retry_request_async(ocr_request)
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}")
@ -593,6 +690,7 @@ class MistralLoader:
def load(self) -> List[Document]:
"""
Executes the full OCR workflow: upload, get URL, process OCR, delete file.
Uses base64 encoding, or file upload for standard Mistral API.
Synchronous version for backward compatibility.
Returns:
@ -602,14 +700,19 @@ class MistralLoader:
start_time = time.time()
try:
# 1. Upload file
file_id = self._upload_file()
if self.use_base64:
# Base64 encoding: Direct OCR call without file upload
ocr_response = self._process_ocr_base64()
else:
# Standard Mistral API flow
# 1. Upload file
file_id = self._upload_file()
# 2. Get Signed URL
signed_url = self._get_signed_url(file_id)
# 2. Get Signed URL
signed_url = self._get_signed_url(file_id)
# 3. Process OCR
ocr_response = self._process_ocr(signed_url)
# 3. Process OCR
ocr_response = self._process_ocr(signed_url)
# 4. Process results
documents = self._process_results(ocr_response)
@ -650,6 +753,7 @@ class MistralLoader:
async def load_async(self) -> List[Document]:
"""
Asynchronous OCR workflow execution with optimized performance.
Uses base64 encoding, or file upload for standard Mistral API.
Returns:
A list of Document objects, one for each page processed.
@ -659,14 +763,19 @@ class MistralLoader:
try:
async with self._get_session() as session:
# 1. Upload file with streaming
file_id = await self._upload_file_async(session)
if self.use_base64:
# Base64 encoding: Direct OCR call without file upload
ocr_response = await self._process_ocr_base64_async(session)
else:
# Standard Mistral API flow
# 1. Upload file with streaming
file_id = await self._upload_file_async(session)
# 2. Get signed URL
signed_url = await self._get_signed_url_async(session, file_id)
# 2. Get signed URL
signed_url = await self._get_signed_url_async(session, file_id)
# 3. Process OCR
ocr_response = await self._process_ocr_async(session, signed_url)
# 3. Process OCR
ocr_response = await self._process_ocr_async(session, signed_url)
# 4. Process results
documents = self._process_results(ocr_response)

18
backend/open_webui/routers/retrieval.py
View file

@ -480,6 +480,8 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
"DOCUMENT_INTELLIGENCE_MODEL": request.app.state.config.DOCUMENT_INTELLIGENCE_MODEL,
"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,
"MISTRAL_OCR_USE_BASE64": request.app.state.config.MISTRAL_OCR_USE_BASE64,
"MISTRAL_OCR_MODEL": request.app.state.config.MISTRAL_OCR_MODEL,
# MinerU settings
"MINERU_API_MODE": request.app.state.config.MINERU_API_MODE,
"MINERU_API_URL": request.app.state.config.MINERU_API_URL,
@ -666,6 +668,8 @@ class ConfigForm(BaseModel):
DOCUMENT_INTELLIGENCE_MODEL: Optional[str] = None
MISTRAL_OCR_API_BASE_URL: Optional[str] = None
MISTRAL_OCR_API_KEY: Optional[str] = None
MISTRAL_OCR_USE_BASE64: Optional[bool] = None
MISTRAL_OCR_MODEL: Optional[str] = None
# MinerU settings
MINERU_API_MODE: Optional[str] = None
@ -877,6 +881,16 @@ async def update_rag_config(
if form_data.MISTRAL_OCR_API_KEY is not None
else request.app.state.config.MISTRAL_OCR_API_KEY
)
request.app.state.config.MISTRAL_OCR_USE_BASE64 = (
form_data.MISTRAL_OCR_USE_BASE64
if form_data.MISTRAL_OCR_USE_BASE64 is not None
else request.app.state.config.MISTRAL_OCR_USE_BASE64
)
request.app.state.config.MISTRAL_OCR_MODEL = (
form_data.MISTRAL_OCR_MODEL
if form_data.MISTRAL_OCR_MODEL is not None
else request.app.state.config.MISTRAL_OCR_MODEL
)
# MinerU settings
request.app.state.config.MINERU_API_MODE = (
@ -1173,6 +1187,8 @@ async def update_rag_config(
"DOCUMENT_INTELLIGENCE_MODEL": request.app.state.config.DOCUMENT_INTELLIGENCE_MODEL,
"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,
"MISTRAL_OCR_USE_BASE64": request.app.state.config.MISTRAL_OCR_USE_BASE64,
"MISTRAL_OCR_MODEL": request.app.state.config.MISTRAL_OCR_MODEL,
# MinerU settings
"MINERU_API_MODE": request.app.state.config.MINERU_API_MODE,
"MINERU_API_URL": request.app.state.config.MINERU_API_URL,
@ -1591,6 +1607,8 @@ def process_file(
DOCUMENT_INTELLIGENCE_MODEL=request.app.state.config.DOCUMENT_INTELLIGENCE_MODEL,
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,
MISTRAL_OCR_USE_BASE64=request.app.state.config.MISTRAL_OCR_USE_BASE64,
MISTRAL_OCR_MODEL=request.app.state.config.MISTRAL_OCR_MODEL,
MINERU_API_MODE=request.app.state.config.MINERU_API_MODE,
MINERU_API_URL=request.app.state.config.MINERU_API_URL,
MINERU_API_KEY=request.app.state.config.MINERU_API_KEY,

27
src/lib/components/admin/Settings/Documents.svelte
View file

@ -623,6 +623,33 @@
bind:value={RAGConfig.MISTRAL_OCR_API_KEY}
/>
</div>
<div class="my-0.5 flex flex-col w-full">
<div class=" mb-1 text-xs font-medium">
{$i18n.t('OCR Model')}
</div>
<div class="flex w-full">
<div class="flex-1 mr-2">
<input
class="flex-1 w-full text-sm bg-transparent outline-hidden"
placeholder={$i18n.t('Enter OCR Model')}
bind:value={RAGConfig.MISTRAL_OCR_MODEL}
/>
</div>
</div>
</div>
<div class="flex justify-between w-full mt-2">
<div class="self-center text-xs font-medium">
<Tooltip
content={$i18n.t('Uses data URI base64 format instead of file upload.')}
placement="top-start"
>
{$i18n.t('Use Base64 Encoding')}
</Tooltip>
</div>
<div class="flex items-center">
<Switch bind:state={RAGConfig.MISTRAL_OCR_USE_BASE64} />
</div>
</div>
{:else if RAGConfig.CONTENT_EXTRACTION_ENGINE === 'mineru'}
<!-- API Mode Selection -->
<div class="flex w-full mt-2">