import os import json import time import requests from typing import Optional, Dict, Any, Tuple class MistralClient: """ Client for interacting with the Mistral AI OCR API. This client handles authentication, file uploads, and OCR processing requests to the Mistral AI API. Attributes: BASE_URL (str): Base URL for the Mistral AI API MAX_FILE_SIZE (int): Maximum allowed file size in bytes (52 MB) api_key (str): Mistral AI API key for authentication session (requests.Session): Session object for making HTTP requests """ BASE_URL = "https://api.mistral.ai/v1" MAX_FILE_SIZE = 52 * 1024 * 1024 # 52 MB def __init__(self, api_key: Optional[str] = None): """ Initialize the Mistral AI client. Args: api_key (Optional[str]): Mistral AI API key. If not provided, will look for MISTRAL_API_KEY environment variable. Raises: ValueError: If no API key is provided or found in environment variables. """ self.api_key = api_key or os.environ.get("MISTRAL_API_KEY") if not self.api_key: raise ValueError("API key must be provided or set as MISTRAL_API_KEY environment variable") self.session = requests.Session() self.session.headers.update({ "Authorization": f"Bearer {self.api_key}", "Accept": "application/json" }) def upload_file(self, file_path: str) -> str: """ Upload a file to Mistral API for OCR processing. Args: file_path (str): Path to the local file to upload Returns: str: File ID returned by the API Raises: ValueError: If the file is too large or if the upload fails requests.RequestException: If there's an error communicating with the API """ # Check file size file_size = os.path.getsize(file_path) if file_size > self.MAX_FILE_SIZE: raise ValueError(f"File is too large ({file_size/1024/1024:.2f} MB). Maximum allowed size is {self.MAX_FILE_SIZE/1024/1024:.2f} MB") # Retry logic max_retries = 3 retry_delay = 3 last_error = None for attempt in range(1, max_retries + 1): try: with open(file_path, 'rb') as f: files = {'file': f} data = {'purpose': 'ocr'} response = self.session.post( f"{self.BASE_URL}/files", files=files, data=data ) response.raise_for_status() if not response.content: last_error = ValueError("Received empty response from API") time.sleep(retry_delay) continue file_response = response.json() file_id = file_response.get('id') if not file_id: last_error = ValueError("Received response without file ID") time.sleep(retry_delay) continue return file_id except requests.RequestException as e: last_error = e # Retry on server errors or rate limiting if hasattr(e, 'response') and e.response is not None: status_code = e.response.status_code if status_code >= 500 or status_code == 429: time.sleep(retry_delay) continue raise raise last_error or ValueError(f"Failed to upload file after {max_retries} attempts") def get_file_url(self, file_id: str) -> str: """ Get a signed URL for an uploaded file. Args: file_id (str): ID of the file previously uploaded to the API Returns: str: Signed URL that can be used for OCR processing Raises: ValueError: If the API response does not contain a URL requests.RequestException: If there's an error communicating with the API """ response = self.session.get(f"{self.BASE_URL}/files/{file_id}/url?expiry=24") response.raise_for_status() url_response = response.json() url = url_response.get('url') if not url: raise ValueError("API response did not contain a URL") return url def process_ocr(self, doc_type: str, doc_source: str, include_image_base64: bool = False) -> bytes: """ Process a document with OCR. Args: doc_type (str): Type of document, either "document_url" or "image_url" doc_source (str): URL of the document to process include_image_base64 (bool, optional): Whether to include base64-encoded images in the response. Defaults to False. Returns: bytes: JSON response from the API containing OCR results Raises: ValueError: If the document type is unsupported or if processing fails requests.RequestException: If there's an error communicating with the API """ if doc_type not in ["document_url", "image_url"]: raise ValueError(f"Unsupported document type: {doc_type}") document_map = {"type": doc_type} if doc_type == "document_url": document_map["document_url"] = doc_source elif doc_type == "image_url": document_map["image_url"] = doc_source request_body = { "model": "mistral-ocr-latest", "document": document_map, "include_image_base64": include_image_base64 } # Retry logic max_retries = 5 retry_delay = 10 last_error = None for attempt in range(1, max_retries + 1): try: response = self.session.post( f"{self.BASE_URL}/ocr", json=request_body, headers={"Content-Type": "application/json"} ) # Check for non-200 status codes if response.status_code != 200: error_msg = response.text or response.reason # Retry on server errors or rate limiting if response.status_code >= 500 or response.status_code == 429: last_error = ValueError(f"API returned error status: {response.status_code} - {error_msg}") time.sleep(retry_delay) continue # For other errors, don't retry raise ValueError(f"API returned error status: {response.status_code} - {error_msg}") # Check for empty response if not response.content: last_error = ValueError("Received empty response from API") adjusted_delay = retry_delay * attempt time.sleep(adjusted_delay) continue # Check if response is valid JSON try: json.loads(response.content) except json.JSONDecodeError: last_error = ValueError("Received invalid JSON response from API") time.sleep(retry_delay) continue # If we got here, we have a valid response return response.content except requests.RequestException as e: last_error = e time.sleep(retry_delay) continue raise last_error or ValueError(f"Failed after {max_retries} attempts")