refac
This commit is contained in:
1 parent
09a1c3a948
commit
89709f5f80
6 files changed
+50
No files matched your search
@@ -912,6 +912,8 @@ MISTRAL_OCR_API_BASE_URL = os.getenv('MISTRAL_OCR_API_BASE_URL', 'https://api.mi
|
|||||||
|
|
||||||
MISTRAL_OCR_API_KEY = os.getenv('MISTRAL_OCR_API_KEY', '')
|
MISTRAL_OCR_API_KEY = os.getenv('MISTRAL_OCR_API_KEY', '')
|
||||||
|
|
||||||
|
MISTRAL_OCR_USE_BASE64 = os.getenv('MISTRAL_OCR_USE_BASE64', 'False').lower() == 'true'
|
||||||
|
|
||||||
PADDLEOCR_VL_BASE_URL = os.getenv('PADDLEOCR_VL_BASE_URL', 'http://localhost:8080')
|
PADDLEOCR_VL_BASE_URL = os.getenv('PADDLEOCR_VL_BASE_URL', 'http://localhost:8080')
|
||||||
|
|
||||||
PADDLEOCR_VL_TOKEN = os.getenv('PADDLEOCR_VL_TOKEN', '')
|
PADDLEOCR_VL_TOKEN = os.getenv('PADDLEOCR_VL_TOKEN', '')
|
||||||
@@ -2786,6 +2788,7 @@ DEFAULT_CONFIG = {
|
|||||||
'rag.document_intelligence_model': DOCUMENT_INTELLIGENCE_MODEL,
|
'rag.document_intelligence_model': DOCUMENT_INTELLIGENCE_MODEL,
|
||||||
'rag.mistral_ocr_api_base_url': MISTRAL_OCR_API_BASE_URL,
|
'rag.mistral_ocr_api_base_url': MISTRAL_OCR_API_BASE_URL,
|
||||||
'rag.mistral_ocr_api_key': MISTRAL_OCR_API_KEY,
|
'rag.mistral_ocr_api_key': MISTRAL_OCR_API_KEY,
|
||||||
|
'rag.mistral_ocr_use_base64': MISTRAL_OCR_USE_BASE64,
|
||||||
'rag.paddleocr_vl_base_url': PADDLEOCR_VL_BASE_URL,
|
'rag.paddleocr_vl_base_url': PADDLEOCR_VL_BASE_URL,
|
||||||
'rag.paddleocr_vl_token': PADDLEOCR_VL_TOKEN,
|
'rag.paddleocr_vl_token': PADDLEOCR_VL_TOKEN,
|
||||||
'rag.bypass_embedding_and_retrieval': BYPASS_EMBEDDING_AND_RETRIEVAL,
|
'rag.bypass_embedding_and_retrieval': BYPASS_EMBEDDING_AND_RETRIEVAL,
|
||||||
|
|||||||
@@ -541,6 +541,7 @@ class Loader:
|
|||||||
base_url=self.kwargs.get('MISTRAL_OCR_API_BASE_URL'),
|
base_url=self.kwargs.get('MISTRAL_OCR_API_BASE_URL'),
|
||||||
api_key=self.kwargs.get('MISTRAL_OCR_API_KEY'),
|
api_key=self.kwargs.get('MISTRAL_OCR_API_KEY'),
|
||||||
file_path=file_path,
|
file_path=file_path,
|
||||||
|
use_base64=self.kwargs.get('MISTRAL_OCR_USE_BASE64', False),
|
||||||
)
|
)
|
||||||
elif self.engine == 'paddleocr_vl' and self.kwargs.get('PADDLEOCR_VL_TOKEN') != '':
|
elif self.engine == 'paddleocr_vl' and self.kwargs.get('PADDLEOCR_VL_TOKEN') != '':
|
||||||
loader = PaddleOCRVLLoader(
|
loader = PaddleOCRVLLoader(
|
||||||
|
|||||||
@@ -1,4 +1,5 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
|
import base64
|
||||||
import logging
|
import logging
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
@@ -37,6 +38,7 @@ class MistralLoader:
|
|||||||
timeout: int = 300, # 5 minutes default
|
timeout: int = 300, # 5 minutes default
|
||||||
max_retries: int = 3,
|
max_retries: int = 3,
|
||||||
enable_debug_logging: bool = False,
|
enable_debug_logging: bool = False,
|
||||||
|
use_base64: bool = False,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Initializes the loader with enhanced features.
|
Initializes the loader with enhanced features.
|
||||||
@@ -47,6 +49,7 @@ class MistralLoader:
|
|||||||
timeout: Request timeout in seconds.
|
timeout: Request timeout in seconds.
|
||||||
max_retries: Maximum number of retry attempts.
|
max_retries: Maximum number of retry attempts.
|
||||||
enable_debug_logging: Enable detailed debug logs.
|
enable_debug_logging: Enable detailed debug logs.
|
||||||
|
use_base64: Send the document as a data URL instead of uploading it first.
|
||||||
"""
|
"""
|
||||||
if not api_key:
|
if not api_key:
|
||||||
raise ValueError('API key cannot be empty.')
|
raise ValueError('API key cannot be empty.')
|
||||||
@@ -59,6 +62,7 @@ class MistralLoader:
|
|||||||
self.timeout = timeout
|
self.timeout = timeout
|
||||||
self.max_retries = max_retries
|
self.max_retries = max_retries
|
||||||
self.debug = enable_debug_logging
|
self.debug = enable_debug_logging
|
||||||
|
self.use_base64 = use_base64
|
||||||
|
|
||||||
# PERFORMANCE OPTIMIZATION: Differentiated timeouts for different operations
|
# PERFORMANCE OPTIMIZATION: Differentiated timeouts for different operations
|
||||||
# This prevents long-running OCR operations from affecting quick operations
|
# This prevents long-running OCR operations from affecting quick operations
|
||||||
@@ -416,6 +420,11 @@ class MistralLoader:
|
|||||||
|
|
||||||
return await self._retry_request_async(ocr_request)
|
return await self._retry_request_async(ocr_request)
|
||||||
|
|
||||||
|
def _get_file_data_url(self) -> str:
|
||||||
|
with open(self.file_path, 'rb') as f:
|
||||||
|
encoded_file = base64.b64encode(f.read()).decode('utf-8')
|
||||||
|
return f'data:application/pdf;base64,{encoded_file}'
|
||||||
|
|
||||||
def _delete_file(self, file_id: str) -> None:
|
def _delete_file(self, file_id: str) -> None:
|
||||||
"""Deletes the file from Mistral storage (sync version)."""
|
"""Deletes the file from Mistral storage (sync version)."""
|
||||||
log.info(f'Deleting uploaded file ID: {file_id}')
|
log.info(f'Deleting uploaded file ID: {file_id}')
|
||||||
@@ -565,6 +574,12 @@ class MistralLoader:
|
|||||||
start_time = time.time()
|
start_time = time.time()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
|
if self.use_base64:
|
||||||
|
documents = self._process_results(self._process_ocr(self._get_file_data_url()))
|
||||||
|
total_time = time.time() - start_time
|
||||||
|
log.info(f'Sync OCR workflow completed in {total_time:.2f}s, produced {len(documents)} documents')
|
||||||
|
return documents
|
||||||
|
|
||||||
# 1. Upload file
|
# 1. Upload file
|
||||||
file_id = self._upload_file()
|
file_id = self._upload_file()
|
||||||
|
|
||||||
@@ -616,6 +631,14 @@ class MistralLoader:
|
|||||||
|
|
||||||
try:
|
try:
|
||||||
async with self._get_session() as session:
|
async with self._get_session() as session:
|
||||||
|
if self.use_base64:
|
||||||
|
documents = self._process_results(await self._process_ocr_async(session, self._get_file_data_url()))
|
||||||
|
total_time = time.time() - start_time
|
||||||
|
log.info(
|
||||||
|
f'Async OCR workflow completed in {total_time:.2f}s, produced {len(documents)} documents'
|
||||||
|
)
|
||||||
|
return documents
|
||||||
|
|
||||||
# 1. Upload file with streaming
|
# 1. Upload file with streaming
|
||||||
file_id = await self._upload_file_async(session)
|
file_id = await self._upload_file_async(session)
|
||||||
|
|
||||||
|
|||||||
@@ -97,6 +97,7 @@ LOADER_CONFIG_KEYS = {
|
|||||||
'DOCUMENT_INTELLIGENCE_MODEL': 'rag.document_intelligence_model',
|
'DOCUMENT_INTELLIGENCE_MODEL': 'rag.document_intelligence_model',
|
||||||
'MISTRAL_OCR_API_BASE_URL': 'rag.mistral_ocr_api_base_url',
|
'MISTRAL_OCR_API_BASE_URL': 'rag.mistral_ocr_api_base_url',
|
||||||
'MISTRAL_OCR_API_KEY': 'rag.mistral_ocr_api_key',
|
'MISTRAL_OCR_API_KEY': 'rag.mistral_ocr_api_key',
|
||||||
|
'MISTRAL_OCR_USE_BASE64': 'rag.mistral_ocr_use_base64',
|
||||||
'PADDLEOCR_VL_BASE_URL': 'rag.paddleocr_vl_base_url',
|
'PADDLEOCR_VL_BASE_URL': 'rag.paddleocr_vl_base_url',
|
||||||
'PADDLEOCR_VL_TOKEN': 'rag.paddleocr_vl_token',
|
'PADDLEOCR_VL_TOKEN': 'rag.paddleocr_vl_token',
|
||||||
'MINERU_API_MODE': 'rag.mineru_api_mode',
|
'MINERU_API_MODE': 'rag.mineru_api_mode',
|
||||||
|
|||||||
@@ -325,6 +325,7 @@ RETRIEVAL_CONFIG_KEYS = {
|
|||||||
'MICROSOFT_WEB_IQ_LANGUAGE': 'web.search.microsoft_web_iq_language',
|
'MICROSOFT_WEB_IQ_LANGUAGE': 'web.search.microsoft_web_iq_language',
|
||||||
'MISTRAL_OCR_API_BASE_URL': 'rag.mistral_ocr_api_base_url',
|
'MISTRAL_OCR_API_BASE_URL': 'rag.mistral_ocr_api_base_url',
|
||||||
'MISTRAL_OCR_API_KEY': 'rag.mistral_ocr_api_key',
|
'MISTRAL_OCR_API_KEY': 'rag.mistral_ocr_api_key',
|
||||||
|
'MISTRAL_OCR_USE_BASE64': 'rag.mistral_ocr_use_base64',
|
||||||
'MOJEEK_SEARCH_API_KEY': 'web.search.mojeek_search_api_key',
|
'MOJEEK_SEARCH_API_KEY': 'web.search.mojeek_search_api_key',
|
||||||
'OLLAMA_CLOUD_WEB_SEARCH_API_KEY': 'web.search.ollama_cloud_api_key',
|
'OLLAMA_CLOUD_WEB_SEARCH_API_KEY': 'web.search.ollama_cloud_api_key',
|
||||||
'PADDLEOCR_VL_BASE_URL': 'rag.paddleocr_vl_base_url',
|
'PADDLEOCR_VL_BASE_URL': 'rag.paddleocr_vl_base_url',
|
||||||
@@ -650,6 +651,7 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
|
|||||||
'DOCUMENT_INTELLIGENCE_MODEL': config.DOCUMENT_INTELLIGENCE_MODEL,
|
'DOCUMENT_INTELLIGENCE_MODEL': config.DOCUMENT_INTELLIGENCE_MODEL,
|
||||||
'MISTRAL_OCR_API_BASE_URL': config.MISTRAL_OCR_API_BASE_URL,
|
'MISTRAL_OCR_API_BASE_URL': config.MISTRAL_OCR_API_BASE_URL,
|
||||||
'MISTRAL_OCR_API_KEY': config.MISTRAL_OCR_API_KEY,
|
'MISTRAL_OCR_API_KEY': config.MISTRAL_OCR_API_KEY,
|
||||||
|
'MISTRAL_OCR_USE_BASE64': config.MISTRAL_OCR_USE_BASE64,
|
||||||
'PADDLEOCR_VL_BASE_URL': config.PADDLEOCR_VL_BASE_URL,
|
'PADDLEOCR_VL_BASE_URL': config.PADDLEOCR_VL_BASE_URL,
|
||||||
'PADDLEOCR_VL_TOKEN': config.PADDLEOCR_VL_TOKEN,
|
'PADDLEOCR_VL_TOKEN': config.PADDLEOCR_VL_TOKEN,
|
||||||
# MinerU settings
|
# MinerU settings
|
||||||
@@ -881,6 +883,7 @@ class ConfigForm(BaseModel):
|
|||||||
DOCUMENT_INTELLIGENCE_MODEL: str | None = None
|
DOCUMENT_INTELLIGENCE_MODEL: str | None = None
|
||||||
MISTRAL_OCR_API_BASE_URL: str | None = None
|
MISTRAL_OCR_API_BASE_URL: str | None = None
|
||||||
MISTRAL_OCR_API_KEY: str | None = None
|
MISTRAL_OCR_API_KEY: str | None = None
|
||||||
|
MISTRAL_OCR_USE_BASE64: bool | None = None
|
||||||
PADDLEOCR_VL_BASE_URL: str | None = None
|
PADDLEOCR_VL_BASE_URL: str | None = None
|
||||||
PADDLEOCR_VL_TOKEN: str | None = None
|
PADDLEOCR_VL_TOKEN: str | None = None
|
||||||
|
|
||||||
@@ -1074,6 +1077,11 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
|
|||||||
config.MISTRAL_OCR_API_KEY = (
|
config.MISTRAL_OCR_API_KEY = (
|
||||||
form_data.MISTRAL_OCR_API_KEY if form_data.MISTRAL_OCR_API_KEY is not None else config.MISTRAL_OCR_API_KEY
|
form_data.MISTRAL_OCR_API_KEY if form_data.MISTRAL_OCR_API_KEY is not None else config.MISTRAL_OCR_API_KEY
|
||||||
)
|
)
|
||||||
|
config.MISTRAL_OCR_USE_BASE64 = (
|
||||||
|
form_data.MISTRAL_OCR_USE_BASE64
|
||||||
|
if form_data.MISTRAL_OCR_USE_BASE64 is not None
|
||||||
|
else config.MISTRAL_OCR_USE_BASE64
|
||||||
|
)
|
||||||
config.PADDLEOCR_VL_BASE_URL = (
|
config.PADDLEOCR_VL_BASE_URL = (
|
||||||
form_data.PADDLEOCR_VL_BASE_URL if form_data.PADDLEOCR_VL_BASE_URL is not None else config.PADDLEOCR_VL_BASE_URL
|
form_data.PADDLEOCR_VL_BASE_URL if form_data.PADDLEOCR_VL_BASE_URL is not None else config.PADDLEOCR_VL_BASE_URL
|
||||||
)
|
)
|
||||||
@@ -1342,6 +1350,7 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
|
|||||||
'DOCUMENT_INTELLIGENCE_MODEL': config.DOCUMENT_INTELLIGENCE_MODEL,
|
'DOCUMENT_INTELLIGENCE_MODEL': config.DOCUMENT_INTELLIGENCE_MODEL,
|
||||||
'MISTRAL_OCR_API_BASE_URL': config.MISTRAL_OCR_API_BASE_URL,
|
'MISTRAL_OCR_API_BASE_URL': config.MISTRAL_OCR_API_BASE_URL,
|
||||||
'MISTRAL_OCR_API_KEY': config.MISTRAL_OCR_API_KEY,
|
'MISTRAL_OCR_API_KEY': config.MISTRAL_OCR_API_KEY,
|
||||||
|
'MISTRAL_OCR_USE_BASE64': config.MISTRAL_OCR_USE_BASE64,
|
||||||
'PADDLEOCR_VL_BASE_URL': config.PADDLEOCR_VL_BASE_URL,
|
'PADDLEOCR_VL_BASE_URL': config.PADDLEOCR_VL_BASE_URL,
|
||||||
'PADDLEOCR_VL_TOKEN': config.PADDLEOCR_VL_TOKEN,
|
'PADDLEOCR_VL_TOKEN': config.PADDLEOCR_VL_TOKEN,
|
||||||
# MinerU settings
|
# MinerU settings
|
||||||
|
|||||||
@@ -760,6 +760,19 @@
|
|||||||
bind:value={RAGConfig.MISTRAL_OCR_API_KEY}
|
bind:value={RAGConfig.MISTRAL_OCR_API_KEY}
|
||||||
/>
|
/>
|
||||||
</div>
|
</div>
|
||||||
|
<div class="flex justify-between w-full mt-2 pr-2">
|
||||||
|
<div class="self-center text-xs font-medium">
|
||||||
|
<Tooltip
|
||||||
|
content={$i18n.t('Send the PDF as a base64 data URL instead of uploading it first.')}
|
||||||
|
placement="top-start"
|
||||||
|
>
|
||||||
|
{$i18n.t('Use Base64')}
|
||||||
|
</Tooltip>
|
||||||
|
</div>
|
||||||
|
<div class="flex items-center">
|
||||||
|
<Switch bind:state={RAGConfig.MISTRAL_OCR_USE_BASE64} />
|
||||||
|
</div>
|
||||||
|
</div>
|
||||||
{:else if RAGConfig.CONTENT_EXTRACTION_ENGINE === 'paddleocr_vl'}
|
{:else if RAGConfig.CONTENT_EXTRACTION_ENGINE === 'paddleocr_vl'}
|
||||||
<div class="my-0.5 flex gap-2 pr-2">
|
<div class="my-0.5 flex gap-2 pr-2">
|
||||||
<input
|
<input
|
||||||
|
|||||||
Reference in new issue
Block a user