Coverage for open_webui/retrieval/loaders/mistral.py: 11%

1import base64 

2import logging 

3import os 

4import sys 

5import time 

6from typing import Any, Dict, List, Optional 

7 

8import requests 

9from langchain_core.documents import Document 

10from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS, GLOBAL_LOG_LEVEL 

11from open_webui.utils.headers import include_user_info_headers 

12 

13logging.basicConfig(stream=sys.stdout, level=GLOBAL_LOG_LEVEL) 

14log = logging.getLogger(__name__) 

15 

16 

17class MistralLoader: 

18 """ 

19 Enhanced Mistral OCR loader. 

20 Loads documents by processing them through the Mistral OCR API. 

21 

22 Performance Optimizations: 

23 - Differentiated timeouts for different operations 

24 - Intelligent retry logic with exponential backoff 

25 - Enhanced error handling with retryable error classification 

26 """ 

27 

28 def __init__( 

29 self, 

30 base_url: str, 

31 api_key: str, 

32 file_path: str, 

33 timeout: int = 300, # 5 minutes default 

34 max_retries: int = 3, 

35 enable_debug_logging: bool = False, 

36 use_base64: bool = False, 

37 user: Optional[Any] = None, 

38 ): 

39 """ 

40 Initializes the loader with enhanced features. 

41 

42 Args: 

43 api_key: Your Mistral API key. 

44 file_path: The local path to the PDF file to process. 

45 timeout: Request timeout in seconds. 

46 max_retries: Maximum number of retry attempts. 

47 enable_debug_logging: Enable detailed debug logs. 

48 use_base64: Send the document as a data URL instead of uploading it first. 

49 user: The requesting user, forwarded to Mistral via user-info headers 

50 when ENABLE_FORWARD_USER_INFO_HEADERS is enabled. 

51 """ 

52 if not api_key: 

53 raise ValueError('API key cannot be empty.') 

54 if not os.path.exists(file_path): 

55 raise FileNotFoundError(f'File not found at {file_path}') 

56 

57 self.base_url = base_url.rstrip('/') if base_url else 'https://api.mistral.ai/v1' 

58 self.api_key = api_key 

59 self.file_path = file_path 

60 self.max_retries = max_retries 

61 self.debug = enable_debug_logging 

62 self.use_base64 = use_base64 

63 self.user = user 

64 

65 # PERFORMANCE OPTIMIZATION: Differentiated timeouts for different operations 

66 # This prevents long-running OCR operations from affecting quick operations 

67 # and improves user experience by failing fast on operations that should be quick 

68 self.upload_timeout = min(timeout, 120) # Cap upload at 2 minutes - prevents hanging on large files 

69 self.url_timeout = 30 # URL requests should be fast - fail quickly if API is slow 

70 self.ocr_timeout = timeout # OCR can take the full timeout - this is the heavy operation 

71 self.cleanup_timeout = 30 # Cleanup should be quick - don't hang on file deletion 

72 

73 # PERFORMANCE OPTIMIZATION: Pre-compute file info to avoid repeated filesystem calls 

74 # This avoids multiple os.path.basename() and os.path.getsize() calls during processing 

75 self.file_name = os.path.basename(file_path) 

76 self.file_size = os.path.getsize(file_path) 

77 

78 # ENHANCEMENT: Added User-Agent for better API tracking and debugging 

79 self.headers = { 

80 'Authorization': f'Bearer {self.api_key}', 

81 'User-Agent': 'OpenWebUI-MistralLoader/2.0', # Helps API provider track usage 

82 } 

83 if self.user is not None and ENABLE_FORWARD_USER_INFO_HEADERS: 

84 self.headers = include_user_info_headers(self.headers, self.user) 

85 

86 def _debug_log(self, message: str, *args) -> None: 

87 """ 

88 PERFORMANCE OPTIMIZATION: Conditional debug logging for performance. 

89 

90 Only processes debug messages when debug mode is enabled, avoiding 

91 string formatting overhead in production environments. 

92 """ 

93 if self.debug: 

94 log.debug(message, *args) 

95 

96 def _handle_response(self, response: requests.Response) -> Dict[str, Any]: 

97 """Checks response status and returns JSON content.""" 

98 try: 

99 response.raise_for_status() # Raises HTTPError for bad responses (4xx or 5xx) 

100 # Handle potential empty responses for certain successful requests (e.g., DELETE) 

101 if response.status_code == 204 or not response.content: 

102 return {} # Return empty dict if no content 

103 return response.json() 

104 except requests.exceptions.HTTPError as http_err: 

105 log.error(f'HTTP error occurred: {http_err} - Response: {response.text}') 

106 raise 

107 except requests.exceptions.RequestException as req_err: 

108 log.error(f'Request exception occurred: {req_err}') 

109 raise 

110 except ValueError as json_err: # Includes JSONDecodeError 

111 log.error(f'JSON decode error: {json_err} - Response: {response.text}') 

112 raise # Re-raise after logging 

113 

114 def _is_retryable_error(self, error: Exception) -> bool: 

115 """ 

116 ENHANCEMENT: Intelligent error classification for retry logic. 

117 

118 Determines if an error is retryable based on its type and status code. 

119 This prevents wasting time retrying errors that will never succeed 

120 (like authentication errors) while ensuring transient errors are retried. 

121 

122 Retryable errors: 

123 - Network connection errors (temporary network issues) 

124 - Timeouts (server might be temporarily overloaded) 

125 - Server errors (5xx status codes - server-side issues) 

126 - Rate limiting (429 status - temporary throttling) 

127 

128 Non-retryable errors: 

129 - Authentication errors (401, 403 - won't fix with retry) 

130 - Bad request errors (400 - malformed request) 

131 - Not found errors (404 - resource doesn't exist) 

132 """ 

133 if isinstance(error, requests.exceptions.ConnectionError): 

134 return True # Network issues are usually temporary 

135 if isinstance(error, requests.exceptions.Timeout): 

136 return True # Timeouts might resolve on retry 

137 if isinstance(error, requests.exceptions.HTTPError): 

138 # Only retry on server errors (5xx) or rate limits (429) 

139 if hasattr(error, 'response') and error.response is not None: 

140 status_code = error.response.status_code 

141 return status_code >= 500 or status_code == 429 

142 return False 

143 return False # All other errors are non-retryable 

144 

145 def _retry_request_sync(self, request_func, *args, **kwargs): 

146 """ 

147 ENHANCEMENT: Synchronous retry logic with intelligent error classification. 

148 

149 Uses exponential backoff with jitter to avoid thundering herd problems. 

150 The wait time increases exponentially but is capped at 30 seconds to 

151 prevent excessive delays. Only retries errors that are likely to succeed 

152 on subsequent attempts. 

153 """ 

154 for attempt in range(self.max_retries): 

155 try: 

156 return request_func(*args, **kwargs) 

157 except Exception as e: 

158 if attempt == self.max_retries - 1 or not self._is_retryable_error(e): 

159 raise 

160 

161 # PERFORMANCE OPTIMIZATION: Exponential backoff with cap 

162 # Prevents overwhelming the server while ensuring reasonable retry delays 

163 wait_time = min((2**attempt) + 0.5, 30) # Cap at 30 seconds 

164 log.warning( 

165 f'Retryable error (attempt {attempt + 1}/{self.max_retries}): {e}. Retrying in {wait_time}s...' 

166 ) 

167 time.sleep(wait_time) 

168 

169 def _upload_file(self) -> str: 

170 """ 

171 PERFORMANCE OPTIMIZATION: Enhanced file upload with streaming consideration. 

172 

173 Uploads the file to Mistral for OCR processing. 

174 Uses context manager for file handling to ensure proper resource cleanup. 

175 Although streaming is not enabled for this endpoint, the file is opened 

176 in a context manager to minimize memory usage duration. 

177 """ 

178 log.info('Uploading file to Mistral API') 

179 url = f'{self.base_url}/files' 

180 

181 def upload_request(): 

182 # MEMORY OPTIMIZATION: Use context manager to minimize file handle lifetime 

183 # This ensures the file is closed immediately after reading, reducing memory usage 

184 with open(self.file_path, 'rb') as f: 

185 files = {'file': (self.file_name, f, 'application/pdf')} 

186 data = {'purpose': 'ocr'} 

187 

188 # NOTE: stream=False is required for this endpoint 

189 # The Mistral API doesn't support chunked uploads for this endpoint 

190 response = requests.post( 

191 url, 

192 headers=self.headers, 

193 files=files, 

194 data=data, 

195 timeout=self.upload_timeout, # Use specialized upload timeout 

196 stream=False, # Keep as False for this endpoint 

197 ) 

198 

199 return self._handle_response(response) 

200 

201 try: 

202 response_data = self._retry_request_sync(upload_request) 

203 file_id = response_data.get('id') 

204 if not file_id: 

205 raise ValueError('File ID not found in upload response.') 

206 log.info('File uploaded successfully. File ID: %s', file_id) 

207 return file_id 

208 except Exception as e: 

209 log.error(f'Failed to upload file: {e}') 

210 raise 

211 

212 def _get_signed_url(self, file_id: str) -> str: 

213 """Retrieves a temporary signed URL for the uploaded file.""" 

214 log.info('Getting signed URL for file ID: %s', file_id) 

215 url = f'{self.base_url}/files/{file_id}/url' 

216 params = {'expiry': 1} 

217 signed_url_headers = {**self.headers, 'Accept': 'application/json'} 

218 

219 def url_request(): 

220 response = requests.get(url, headers=signed_url_headers, params=params, timeout=self.url_timeout) 

221 return self._handle_response(response) 

222 

223 try: 

224 response_data = self._retry_request_sync(url_request) 

225 signed_url = response_data.get('url') 

226 if not signed_url: 

227 raise ValueError('Signed URL not found in response.') 

228 log.info('Signed URL received.') 

229 return signed_url 

230 except Exception as e: 

231 log.error(f'Failed to get signed URL: {e}') 

232 raise 

233 

234 def _process_ocr(self, signed_url: str) -> Dict[str, Any]: 

235 """Sends the signed URL to the OCR endpoint for processing.""" 

236 log.info('Processing OCR via Mistral API') 

237 url = f'{self.base_url}/ocr' 

238 ocr_headers = { 

239 **self.headers, 

240 'Content-Type': 'application/json', 

241 'Accept': 'application/json', 

242 } 

243 payload = { 

244 'model': 'mistral-ocr-latest', 

245 'document': { 

246 'type': 'document_url', 

247 'document_url': signed_url, 

248 }, 

249 'include_image_base64': False, 

250 } 

251 

252 def ocr_request(): 

253 response = requests.post(url, headers=ocr_headers, json=payload, timeout=self.ocr_timeout) 

254 return self._handle_response(response) 

255 

256 try: 

257 ocr_response = self._retry_request_sync(ocr_request) 

258 log.info('OCR processing done.') 

259 self._debug_log('OCR response: %s', ocr_response) 

260 return ocr_response 

261 except Exception as e: 

262 log.error(f'Failed during OCR processing: {e}') 

263 raise 

264 

265 def _get_file_data_url(self) -> str: 

266 with open(self.file_path, 'rb') as f: 

267 encoded_file = base64.b64encode(f.read()).decode('utf-8') 

268 return f'data:application/pdf;base64,{encoded_file}' 

269 

270 def _delete_file(self, file_id: str) -> None: 

271 """Deletes the file from Mistral storage.""" 

272 log.info('Deleting uploaded file ID: %s', file_id) 

273 url = f'{self.base_url}/files/{file_id}' 

274 

275 try: 

276 response = requests.delete(url, headers=self.headers, timeout=self.cleanup_timeout) 

277 delete_response = self._handle_response(response) 

278 log.info('File deleted successfully: %s', delete_response) 

279 except Exception as e: 

280 # Log error but don't necessarily halt execution if deletion fails 

281 log.error(f'Failed to delete file ID {file_id}: {e}') 

282 

283 def _process_results(self, ocr_response: Dict[str, Any]) -> List[Document]: 

284 """Process OCR results into Document objects with enhanced metadata and memory efficiency.""" 

285 pages_data = ocr_response.get('pages') 

286 if not pages_data: 

287 log.warning('No pages found in OCR response.') 

288 return [ 

289 Document( 

290 page_content='No text content found', 

291 metadata={'error': 'no_pages', 'file_name': self.file_name}, 

292 ) 

293 ] 

294 

295 documents = [] 

296 total_pages = len(pages_data) 

297 skipped_pages = 0 

298 

299 # Process pages in a memory-efficient way 

300 for page_data in pages_data: 

301 page_content = page_data.get('markdown') 

302 page_index = page_data.get('index') # API uses 0-based index 

303 

304 if page_content is None or page_index is None: 

305 skipped_pages += 1 

306 self._debug_log( 

307 "Skipping page due to missing 'markdown' or 'index'. Data keys: %s", list(page_data.keys()) 

308 ) 

309 continue 

310 

311 # Clean up content efficiently with early exit for empty content 

312 if isinstance(page_content, str): 

313 cleaned_content = page_content.strip() 

314 else: 

315 cleaned_content = str(page_content).strip() 

316 

317 if not cleaned_content: 

318 skipped_pages += 1 

319 self._debug_log('Skipping empty page %s', page_index) 

320 continue 

321 

322 # Create document with optimized metadata 

323 documents.append( 

324 Document( 

325 page_content=cleaned_content, 

326 metadata={ 

327 'page': page_index, # 0-based index from API 

328 'page_label': page_index + 1, # 1-based label for convenience 

329 'total_pages': total_pages, 

330 'file_name': self.file_name, 

331 'file_size': self.file_size, 

332 'processing_engine': 'mistral-ocr', 

333 'content_length': len(cleaned_content), 

334 }, 

335 ) 

336 ) 

337 

338 if skipped_pages > 0: 

339 log.info('Processed %s pages, skipped %s empty/invalid pages', len(documents), skipped_pages) 

340 

341 if not documents: 

342 # Case where pages existed but none had valid markdown/index 

343 log.warning('OCR response contained pages, but none had valid content/index.') 

344 return [ 

345 Document( 

346 page_content='No valid text content found in document', 

347 metadata={ 

348 'error': 'no_valid_pages', 

349 'total_pages': total_pages, 

350 'file_name': self.file_name, 

351 }, 

352 ) 

353 ] 

354 

355 return documents 

356 

357 def load(self) -> List[Document]: 

358 """ 

359 Executes the full OCR workflow: upload, get URL, process OCR, delete file. 

360 

361 Returns: 

362 A list of Document objects, one for each page processed. 

363 """ 

364 file_id = None 

365 start_time = time.time() 

366 

367 try: 

368 if self.use_base64: 

369 documents = self._process_results(self._process_ocr(self._get_file_data_url())) 

370 total_time = time.time() - start_time 

371 log.info('Sync OCR workflow completed in %.2fs, produced %s documents', total_time, len(documents)) 

372 return documents 

373 

374 # 1. Upload file 

375 file_id = self._upload_file() 

376 

377 # 2. Get Signed URL 

378 signed_url = self._get_signed_url(file_id) 

379 

380 # 3. Process OCR 

381 ocr_response = self._process_ocr(signed_url) 

382 

383 # 4. Process results 

384 documents = self._process_results(ocr_response) 

385 

386 total_time = time.time() - start_time 

387 log.info('Sync OCR workflow completed in %.2fs, produced %s documents', total_time, len(documents)) 

388 

389 return documents 

390 

391 except Exception as e: 

392 total_time = time.time() - start_time 

393 log.error(f'An error occurred during the loading process after {total_time:.2f}s: {e}') 

394 # Return an error document on failure 

395 return [ 

396 Document( 

397 page_content=f'Error during processing: {e}', 

398 metadata={ 

399 'error': 'processing_failed', 

400 'file_name': self.file_name, 

401 }, 

402 ) 

403 ] 

404 finally: 

405 # 5. Delete file (attempt even if prior steps failed after upload) 

406 if file_id: 

407 try: 

408 self._delete_file(file_id) 

409 except Exception as del_e: 

410 # Log deletion error, but don't overwrite original error if one occurred 

411 log.error(f'Cleanup error: Could not delete file ID {file_id}. Reason: {del_e}')