Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-03-17 17:58:01 -05:00
parent fcf7208352
commit de3317e26b
220 changed files with 17200 additions and 22836 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/retrieval/loaders/mistral.py
+139 -191
View File
@@ -49,13 +49,11 @@ class MistralLoader:
enable_debug_logging: Enable detailed debug logs.
"""
if not api_key:
raise ValueError("API key cannot be empty.")
raise ValueError('API key cannot be empty.')
if not os.path.exists(file_path):
raise FileNotFoundError(f"File not found at {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.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
@@ -65,18 +63,10 @@ class MistralLoader:
# PERFORMANCE OPTIMIZATION: Differentiated timeouts for different operations
# This prevents long-running OCR operations from affecting quick operations
# and improves user experience by failing fast on operations that should be quick
self.upload_timeout = min(
timeout, 120
) # Cap upload at 2 minutes - prevents hanging on large files
self.url_timeout = (
30 # URL requests should be fast - fail quickly if API is slow
)
self.ocr_timeout = (
timeout # OCR can take the full timeout - this is the heavy operation
)
self.cleanup_timeout = (
30 # Cleanup should be quick - don't hang on file deletion
)
self.upload_timeout = min(timeout, 120) # Cap upload at 2 minutes - prevents hanging on large files
self.url_timeout = 30 # URL requests should be fast - fail quickly if API is slow
self.ocr_timeout = timeout # OCR can take the full timeout - this is the heavy operation
self.cleanup_timeout = 30 # Cleanup should be quick - don't hang on file deletion
# PERFORMANCE OPTIMIZATION: Pre-compute file info to avoid repeated filesystem calls
# This avoids multiple os.path.basename() and os.path.getsize() calls during processing
@@ -85,8 +75,8 @@ class MistralLoader:
# ENHANCEMENT: Added User-Agent for better API tracking and debugging
self.headers = {
"Authorization": f"Bearer {self.api_key}",
"User-Agent": "OpenWebUI-MistralLoader/2.0", # Helps API provider track usage
'Authorization': f'Bearer {self.api_key}',
'User-Agent': 'OpenWebUI-MistralLoader/2.0', # Helps API provider track usage
}
def _debug_log(self, message: str, *args) -> None:
@@ -108,43 +98,39 @@ class MistralLoader:
return {} # Return empty dict if no content
return response.json()
except requests.exceptions.HTTPError as http_err:
log.error(f"HTTP error occurred: {http_err} - Response: {response.text}")
log.error(f'HTTP error occurred: {http_err} - Response: {response.text}')
raise
except requests.exceptions.RequestException as req_err:
log.error(f"Request exception occurred: {req_err}")
log.error(f'Request exception occurred: {req_err}')
raise
except ValueError as json_err: # Includes JSONDecodeError
log.error(f"JSON decode error: {json_err} - Response: {response.text}")
log.error(f'JSON decode error: {json_err} - Response: {response.text}')
raise # Re-raise after logging
async def _handle_response_async(
self, response: aiohttp.ClientResponse
) -> Dict[str, Any]:
async def _handle_response_async(self, response: aiohttp.ClientResponse) -> Dict[str, Any]:
"""Async version of response handling with better error info."""
try:
response.raise_for_status()
# Check content type
content_type = response.headers.get("content-type", "")
if "application/json" not in content_type:
content_type = response.headers.get('content-type', '')
if 'application/json' not in content_type:
if response.status == 204:
return {}
text = await response.text()
raise ValueError(
f"Unexpected content type: {content_type}, body: {text[:200]}..."
)
raise ValueError(f'Unexpected content type: {content_type}, body: {text[:200]}...')
return await response.json()
except aiohttp.ClientResponseError as e:
error_text = await response.text() if response else "No response"
log.error(f"HTTP {e.status}: {e.message} - Response: {error_text[:500]}")
error_text = await response.text() if response else 'No response'
log.error(f'HTTP {e.status}: {e.message} - Response: {error_text[:500]}')
raise
except aiohttp.ClientError as e:
log.error(f"Client error: {e}")
log.error(f'Client error: {e}')
raise
except Exception as e:
log.error(f"Unexpected error processing response: {e}")
log.error(f'Unexpected error processing response: {e}')
raise
def _is_retryable_error(self, error: Exception) -> bool:
@@ -172,13 +158,11 @@ class MistralLoader:
return True # Timeouts might resolve on retry
if isinstance(error, requests.exceptions.HTTPError):
# Only retry on server errors (5xx) or rate limits (429)
if hasattr(error, "response") and error.response is not None:
if hasattr(error, 'response') and error.response is not None:
status_code = error.response.status_code
return status_code >= 500 or status_code == 429
return False
if isinstance(
error, (aiohttp.ClientConnectionError, aiohttp.ServerTimeoutError)
):
if isinstance(error, (aiohttp.ClientConnectionError, aiohttp.ServerTimeoutError)):
return True # Async network/timeout errors are retryable
if isinstance(error, aiohttp.ClientResponseError):
return error.status >= 500 or error.status == 429
@@ -204,8 +188,7 @@ class MistralLoader:
# Prevents overwhelming the server while ensuring reasonable retry delays
wait_time = min((2**attempt) + 0.5, 30) # Cap at 30 seconds
log.warning(
f"Retryable error (attempt {attempt + 1}/{self.max_retries}): {e}. "
f"Retrying in {wait_time}s..."
f'Retryable error (attempt {attempt + 1}/{self.max_retries}): {e}. Retrying in {wait_time}s...'
)
time.sleep(wait_time)
@@ -226,8 +209,7 @@ class MistralLoader:
# PERFORMANCE OPTIMIZATION: Non-blocking exponential backoff
wait_time = min((2**attempt) + 0.5, 30) # Cap at 30 seconds
log.warning(
f"Retryable error (attempt {attempt + 1}/{self.max_retries}): {e}. "
f"Retrying in {wait_time}s..."
f'Retryable error (attempt {attempt + 1}/{self.max_retries}): {e}. Retrying in {wait_time}s...'
)
await asyncio.sleep(wait_time) # Non-blocking wait
@@ -240,15 +222,15 @@ class MistralLoader:
Although streaming is not enabled for this endpoint, the file is opened
in a context manager to minimize memory usage duration.
"""
log.info("Uploading file to Mistral API")
url = f"{self.base_url}/files"
log.info('Uploading file to Mistral API')
url = f'{self.base_url}/files'
def upload_request():
# MEMORY OPTIMIZATION: Use context manager to minimize file handle lifetime
# This ensures the file is closed immediately after reading, reducing memory usage
with open(self.file_path, "rb") as f:
files = {"file": (self.file_name, f, "application/pdf")}
data = {"purpose": "ocr"}
with open(self.file_path, 'rb') as f:
files = {'file': (self.file_name, f, 'application/pdf')}
data = {'purpose': 'ocr'}
# NOTE: stream=False is required for this endpoint
# The Mistral API doesn't support chunked uploads for this endpoint
@@ -265,42 +247,38 @@ class MistralLoader:
try:
response_data = self._retry_request_sync(upload_request)
file_id = response_data.get("id")
file_id = response_data.get('id')
if not file_id:
raise ValueError("File ID not found in upload response.")
log.info(f"File uploaded successfully. File ID: {file_id}")
raise ValueError('File ID not found in upload response.')
log.info(f'File uploaded successfully. File ID: {file_id}')
return file_id
except Exception as e:
log.error(f"Failed to upload file: {e}")
log.error(f'Failed to upload file: {e}')
raise
async def _upload_file_async(self, session: aiohttp.ClientSession) -> str:
"""Async file upload with streaming for better memory efficiency."""
url = f"{self.base_url}/files"
url = f'{self.base_url}/files'
async def upload_request():
# Create multipart writer for streaming upload
writer = aiohttp.MultipartWriter("form-data")
writer = aiohttp.MultipartWriter('form-data')
# Add purpose field
purpose_part = writer.append("ocr")
purpose_part.set_content_disposition("form-data", name="purpose")
purpose_part = writer.append('ocr')
purpose_part.set_content_disposition('form-data', name='purpose')
# Add file part with streaming
file_part = writer.append_payload(
aiohttp.streams.FilePayload(
self.file_path,
filename=self.file_name,
content_type="application/pdf",
content_type='application/pdf',
)
)
file_part.set_content_disposition(
"form-data", name="file", filename=self.file_name
)
file_part.set_content_disposition('form-data', name='file', filename=self.file_name)
self._debug_log(
f"Uploading file: {self.file_name} ({self.file_size:,} bytes)"
)
self._debug_log(f'Uploading file: {self.file_name} ({self.file_size:,} bytes)')
async with session.post(
url,
@@ -312,48 +290,44 @@ class MistralLoader:
response_data = await self._retry_request_async(upload_request)
file_id = response_data.get("id")
file_id = response_data.get('id')
if not file_id:
raise ValueError("File ID not found in upload response.")
raise ValueError('File ID not found in upload response.')
log.info(f"File uploaded successfully. File ID: {file_id}")
log.info(f'File uploaded successfully. File ID: {file_id}')
return file_id
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_url}/files/{file_id}/url"
params = {"expiry": 1}
signed_url_headers = {**self.headers, "Accept": "application/json"}
log.info(f'Getting signed URL for file ID: {file_id}')
url = f'{self.base_url}/files/{file_id}/url'
params = {'expiry': 1}
signed_url_headers = {**self.headers, 'Accept': 'application/json'}
def url_request():
response = requests.get(
url, headers=signed_url_headers, params=params, timeout=self.url_timeout
)
response = requests.get(url, headers=signed_url_headers, params=params, timeout=self.url_timeout)
return self._handle_response(response)
try:
response_data = self._retry_request_sync(url_request)
signed_url = response_data.get("url")
signed_url = response_data.get('url')
if not signed_url:
raise ValueError("Signed URL not found in response.")
log.info("Signed URL received.")
raise ValueError('Signed URL not found in response.')
log.info('Signed URL received.')
return signed_url
except Exception as e:
log.error(f"Failed to get signed URL: {e}")
log.error(f'Failed to get signed URL: {e}')
raise
async def _get_signed_url_async(
self, session: aiohttp.ClientSession, file_id: str
) -> str:
async def _get_signed_url_async(self, session: aiohttp.ClientSession, file_id: str) -> str:
"""Async signed URL retrieval."""
url = f"{self.base_url}/files/{file_id}/url"
params = {"expiry": 1}
url = f'{self.base_url}/files/{file_id}/url'
params = {'expiry': 1}
headers = {**self.headers, "Accept": "application/json"}
headers = {**self.headers, 'Accept': 'application/json'}
async def url_request():
self._debug_log(f"Getting signed URL for file ID: {file_id}")
self._debug_log(f'Getting signed URL for file ID: {file_id}')
async with session.get(
url,
headers=headers,
@@ -364,69 +338,65 @@ class MistralLoader:
response_data = await self._retry_request_async(url_request)
signed_url = response_data.get("url")
signed_url = response_data.get('url')
if not signed_url:
raise ValueError("Signed URL not found in response.")
raise ValueError('Signed URL not found in response.')
self._debug_log("Signed URL received successfully")
self._debug_log('Signed URL received successfully')
return signed_url
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_url}/ocr"
log.info('Processing OCR via Mistral API')
url = f'{self.base_url}/ocr'
ocr_headers = {
**self.headers,
"Content-Type": "application/json",
"Accept": "application/json",
'Content-Type': 'application/json',
'Accept': 'application/json',
}
payload = {
"model": "mistral-ocr-latest",
"document": {
"type": "document_url",
"document_url": signed_url,
'model': 'mistral-ocr-latest',
'document': {
'type': 'document_url',
'document_url': signed_url,
},
"include_image_base64": False,
'include_image_base64': False,
}
def ocr_request():
response = requests.post(
url, headers=ocr_headers, json=payload, timeout=self.ocr_timeout
)
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.")
self._debug_log("OCR response: %s", ocr_response)
log.info('OCR processing done.')
self._debug_log('OCR response: %s', ocr_response)
return ocr_response
except Exception as e:
log.error(f"Failed during OCR processing: {e}")
log.error(f'Failed during OCR processing: {e}')
raise
async def _process_ocr_async(
self, session: aiohttp.ClientSession, signed_url: str
) -> Dict[str, Any]:
async def _process_ocr_async(self, session: aiohttp.ClientSession, signed_url: str) -> Dict[str, Any]:
"""Async OCR processing with timing metrics."""
url = f"{self.base_url}/ocr"
url = f'{self.base_url}/ocr'
headers = {
**self.headers,
"Content-Type": "application/json",
"Accept": "application/json",
'Content-Type': 'application/json',
'Accept': 'application/json',
}
payload = {
"model": "mistral-ocr-latest",
"document": {
"type": "document_url",
"document_url": signed_url,
'model': 'mistral-ocr-latest',
'document': {
'type': 'document_url',
'document_url': signed_url,
},
"include_image_base64": False,
'include_image_base64': False,
}
async def ocr_request():
log.info("Starting OCR processing via Mistral API")
log.info('Starting OCR processing via Mistral API')
start_time = time.time()
async with session.post(
@@ -438,7 +408,7 @@ class MistralLoader:
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")
log.info(f'OCR processing completed in {processing_time:.2f}s')
return ocr_response
@@ -446,42 +416,36 @@ 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_url}/files/{file_id}"
log.info(f'Deleting uploaded file ID: {file_id}')
url = f'{self.base_url}/files/{file_id}'
try:
response = requests.delete(
url, headers=self.headers, timeout=self.cleanup_timeout
)
response = requests.delete(url, headers=self.headers, timeout=self.cleanup_timeout)
delete_response = self._handle_response(response)
log.info(f"File deleted successfully: {delete_response}")
log.info(f'File deleted successfully: {delete_response}')
except Exception as e:
# Log error but don't necessarily halt execution if deletion fails
log.error(f"Failed to delete file ID {file_id}: {e}")
log.error(f'Failed to delete file ID {file_id}: {e}')
async def _delete_file_async(
self, session: aiohttp.ClientSession, file_id: str
) -> None:
async def _delete_file_async(self, session: aiohttp.ClientSession, file_id: str) -> None:
"""Async file deletion with error tolerance."""
try:
async def delete_request():
self._debug_log(f"Deleting file ID: {file_id}")
self._debug_log(f'Deleting file ID: {file_id}')
async with session.delete(
url=f"{self.base_url}/files/{file_id}",
url=f'{self.base_url}/files/{file_id}',
headers=self.headers,
timeout=aiohttp.ClientTimeout(
total=self.cleanup_timeout
), # Shorter timeout for cleanup
timeout=aiohttp.ClientTimeout(total=self.cleanup_timeout), # Shorter timeout for cleanup
) as response:
return await self._handle_response_async(response)
await self._retry_request_async(delete_request)
self._debug_log(f"File {file_id} deleted successfully")
self._debug_log(f'File {file_id} deleted successfully')
except Exception as e:
# Don't fail the entire process if cleanup fails
log.warning(f"Failed to delete file ID {file_id}: {e}")
log.warning(f'Failed to delete file ID {file_id}: {e}')
@asynccontextmanager
async def _get_session(self):
@@ -506,7 +470,7 @@ class MistralLoader:
async with aiohttp.ClientSession(
connector=connector,
timeout=timeout,
headers={"User-Agent": "OpenWebUI-MistralLoader/2.0"},
headers={'User-Agent': 'OpenWebUI-MistralLoader/2.0'},
raise_for_status=False, # We handle status codes manually
trust_env=True,
) as session:
@@ -514,13 +478,13 @@ class MistralLoader:
def _process_results(self, ocr_response: Dict[str, Any]) -> List[Document]:
"""Process OCR results into Document objects with enhanced metadata and memory efficiency."""
pages_data = ocr_response.get("pages")
pages_data = ocr_response.get('pages')
if not pages_data:
log.warning("No pages found in OCR response.")
log.warning('No pages found in OCR response.')
return [
Document(
page_content="No text content found",
metadata={"error": "no_pages", "file_name": self.file_name},
page_content='No text content found',
metadata={'error': 'no_pages', 'file_name': self.file_name},
)
]
@@ -530,8 +494,8 @@ class MistralLoader:
# Process pages in a memory-efficient way
for page_data in pages_data:
page_content = page_data.get("markdown")
page_index = page_data.get("index") # API uses 0-based index
page_content = page_data.get('markdown')
page_index = page_data.get('index') # API uses 0-based index
if page_content is None or page_index is None:
skipped_pages += 1
@@ -548,7 +512,7 @@ class MistralLoader:
if not cleaned_content:
skipped_pages += 1
self._debug_log(f"Skipping empty page {page_index}")
self._debug_log(f'Skipping empty page {page_index}')
continue
# Create document with optimized metadata
@@ -556,34 +520,30 @@ class MistralLoader:
Document(
page_content=cleaned_content,
metadata={
"page": page_index, # 0-based index from API
"page_label": page_index + 1, # 1-based label for convenience
"total_pages": total_pages,
"file_name": self.file_name,
"file_size": self.file_size,
"processing_engine": "mistral-ocr",
"content_length": len(cleaned_content),
'page': page_index, # 0-based index from API
'page_label': page_index + 1, # 1-based label for convenience
'total_pages': total_pages,
'file_name': self.file_name,
'file_size': self.file_size,
'processing_engine': 'mistral-ocr',
'content_length': len(cleaned_content),
},
)
)
if skipped_pages > 0:
log.info(
f"Processed {len(documents)} pages, skipped {skipped_pages} empty/invalid pages"
)
log.info(f'Processed {len(documents)} pages, skipped {skipped_pages} empty/invalid pages')
if not documents:
# Case where pages existed but none had valid markdown/index
log.warning(
"OCR response contained pages, but none had valid content/index."
)
log.warning('OCR response contained pages, but none had valid content/index.')
return [
Document(
page_content="No valid text content found in document",
page_content='No valid text content found in document',
metadata={
"error": "no_valid_pages",
"total_pages": total_pages,
"file_name": self.file_name,
'error': 'no_valid_pages',
'total_pages': total_pages,
'file_name': self.file_name,
},
)
]
@@ -615,24 +575,20 @@ class MistralLoader:
documents = self._process_results(ocr_response)
total_time = time.time() - start_time
log.info(
f"Sync OCR workflow completed in {total_time:.2f}s, produced {len(documents)} documents"
)
log.info(f'Sync OCR workflow completed in {total_time:.2f}s, produced {len(documents)} documents')
return documents
except Exception as e:
total_time = time.time() - start_time
log.error(
f"An error occurred during the loading process after {total_time:.2f}s: {e}"
)
log.error(f'An error occurred during the loading process after {total_time:.2f}s: {e}')
# Return an error document on failure
return [
Document(
page_content=f"Error during processing: {e}",
page_content=f'Error during processing: {e}',
metadata={
"error": "processing_failed",
"file_name": self.file_name,
'error': 'processing_failed',
'file_name': self.file_name,
},
)
]
@@ -643,9 +599,7 @@ class MistralLoader:
self._delete_file(file_id)
except Exception as del_e:
# Log deletion error, but don't overwrite original error if one occurred
log.error(
f"Cleanup error: Could not delete file ID {file_id}. Reason: {del_e}"
)
log.error(f'Cleanup error: Could not delete file ID {file_id}. Reason: {del_e}')
async def load_async(self) -> List[Document]:
"""
@@ -672,21 +626,19 @@ class MistralLoader:
documents = self._process_results(ocr_response)
total_time = time.time() - start_time
log.info(
f"Async OCR workflow completed in {total_time:.2f}s, produced {len(documents)} documents"
)
log.info(f'Async OCR workflow completed in {total_time:.2f}s, produced {len(documents)} documents')
return documents
except Exception as e:
total_time = time.time() - start_time
log.error(f"Async OCR workflow failed after {total_time:.2f}s: {e}")
log.error(f'Async OCR workflow failed after {total_time:.2f}s: {e}')
return [
Document(
page_content=f"Error during OCR processing: {e}",
page_content=f'Error during OCR processing: {e}',
metadata={
"error": "processing_failed",
"file_name": self.file_name,
'error': 'processing_failed',
'file_name': self.file_name,
},
)
]
@@ -697,11 +649,11 @@ class MistralLoader:
async with self._get_session() as session:
await self._delete_file_async(session, file_id)
except Exception as cleanup_error:
log.error(f"Cleanup failed for file ID {file_id}: {cleanup_error}")
log.error(f'Cleanup failed for file ID {file_id}: {cleanup_error}')
@staticmethod
async def load_multiple_async(
loaders: List["MistralLoader"],
loaders: List['MistralLoader'],
max_concurrent: int = 5, # Limit concurrent requests
) -> List[List[Document]]:
"""
@@ -717,15 +669,13 @@ class MistralLoader:
if not loaders:
return []
log.info(
f"Starting concurrent processing of {len(loaders)} files with max {max_concurrent} concurrent"
)
log.info(f'Starting concurrent processing of {len(loaders)} files with max {max_concurrent} concurrent')
start_time = time.time()
# Use semaphore to control concurrency
semaphore = asyncio.Semaphore(max_concurrent)
async def process_with_semaphore(loader: "MistralLoader") -> List[Document]:
async def process_with_semaphore(loader: 'MistralLoader') -> List[Document]:
async with semaphore:
return await loader.load_async()
@@ -737,14 +687,14 @@ class MistralLoader:
processed_results = []
for i, result in enumerate(results):
if isinstance(result, Exception):
log.error(f"File {i} failed: {result}")
log.error(f'File {i} failed: {result}')
processed_results.append(
[
Document(
page_content=f"Error processing file: {result}",
page_content=f'Error processing file: {result}',
metadata={
"error": "batch_processing_failed",
"file_index": i,
'error': 'batch_processing_failed',
'file_index': i,
},
)
]
@@ -755,15 +705,13 @@ class MistralLoader:
# MONITORING: Log comprehensive batch processing statistics
total_time = time.time() - start_time
total_docs = sum(len(docs) for docs in processed_results)
success_count = sum(
1 for result in results if not isinstance(result, Exception)
)
success_count = sum(1 for result in results if not isinstance(result, Exception))
failure_count = len(results) - success_count
log.info(
f"Batch processing completed in {total_time:.2f}s: "
f"{success_count} files succeeded, {failure_count} files failed, "
f"produced {total_docs} total documents"
f'Batch processing completed in {total_time:.2f}s: '
f'{success_count} files succeeded, {failure_count} files failed, '
f'produced {total_docs} total documents'
)
return processed_results