UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-05-03 04:21:46 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

Add Mistral OCR integration and configuration support

Browse Source
This commit is contained in:
Patrick Wachter 2025-03-22 13:44:50 +01:00
parent 04799f1f95
commit 1ac6879268
No known key found for this signature in database
6 changed files with 105 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
backend/open_webui/config.py
View File

@ -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
backend/open_webui/main.py
View File

@ -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
backend/open_webui/retrieval/loaders/main.py
View File

@ -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
backend/open_webui/routers/retrieval.py
View File

@ -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
backend/requirements.txt
View File

@ -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

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

@ -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>