import os import base64 import gradio as gr from mistralai import Mistral from mistralai.models import OCRResponse from pathlib import Path from pydantic import BaseModel import pycountry import json import logging from tenacity import retry, stop_after_attempt, wait_fixed import tempfile from typing import Union, Dict, List from contextlib import contextmanager import requests from enum import Enum # Constants DEFAULT_LANGUAGE = "English" SUPPORTED_IMAGE_TYPES = [".jpg", ".png"] SUPPORTED_PDF_TYPES = [".pdf"] TEMP_FILE_EXPIRY = 7200 # 2 hours in seconds # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) class OCRProcessor: def __init__(self, api_key: str): if not api_key: raise ValueError("API key must be provided") self.api_key = api_key self.client = Mistral(api_key=self.api_key) try: self.client.models.list() # Validate API key except Exception as e: raise ValueError(f"Invalid API key: {str(e)}") @staticmethod def _encode_image(image_path: str) -> str: try: with open(image_path, "rb") as image_file: return base64.b64encode(image_file.read()).decode('utf-8') except Exception as e: logger.error(f"Error encoding image {image_path}: {str(e)}") raise @staticmethod @contextmanager def _temp_file(content: bytes, suffix: str) -> str: temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix) try: temp_file.write(content) temp_file.close() yield temp_file.name finally: if os.path.exists(temp_file.name): os.unlink(temp_file.name) @retry(stop=stop_after_attempt(3), wait=wait_fixed(2)) def _call_ocr_api(self, document: Dict) -> OCRResponse: try: return self.client.ocr.process(model="mistral-ocr-latest", document=document) except Exception as e: logger.error(f"OCR API call failed: {str(e)}") raise @retry(stop=stop_after_attempt(3), wait=wait_fixed(2)) def _call_chat_complete(self, model: str, messages: List[Dict], **kwargs) -> Dict: try: return self.client.chat.complete(model=model, messages=messages, **kwargs) except Exception as e: logger.error(f"Chat complete API call failed: {str(e)}") raise def _get_file_content(self, file_input: Union[str, object]) -> bytes: try: if isinstance(file_input, str) and file_input.startswith(("http://", "https://")): # Handle URLs response = requests.get(file_input, timeout=10) response.raise_for_status() return response.content elif isinstance(file_input, str): # File path with open(file_input, "rb") as f: return f.read() elif hasattr(file_input, 'read'): # File-like object return file_input.read() else: raise ValueError("Invalid file input: must be a URL, path, or file-like object") except Exception as e: logger.error(f"Error getting file content: {str(e)}") raise def ocr_pdf_url(self, pdf_url: str) -> str: logger.info(f"Processing PDF URL: {pdf_url}") try: response = self._call_ocr_api({"type": "document_url", "document_url": pdf_url}) return self._extract_markdown(response) except Exception as e: return self._handle_error("PDF URL processing", e) def ocr_uploaded_pdf(self, pdf_file: Union[str, object]) -> str: file_name = getattr(pdf_file, 'name', 'unknown') logger.info(f"Processing uploaded PDF: {file_name}") try: content = self._get_file_content(pdf_file) with self._temp_file(content, ".pdf") as temp_path: uploaded_file = self.client.files.upload( file={"file_name": temp_path, "content": open(temp_path, "rb")}, purpose="ocr" ) signed_url = self.client.files.get_signed_url(file_id=uploaded_file.id, expiry=TEMP_FILE_EXPIRY) response = self._call_ocr_api({"type": "document_url", "document_url": signed_url.url}) return self._extract_markdown(response) except Exception as e: return self._handle_error("uploaded PDF processing", e) def ocr_image_url(self, image_url: str) -> str: logger.info(f"Processing image URL: {image_url}") try: response = self._call_ocr_api({"type": "image_url", "image_url": image_url}) return self._extract_markdown(response) except Exception as e: return self._handle_error("image URL processing", e) def ocr_uploaded_image(self, image_file: Union[str, object]) -> str: file_name = getattr(image_file, 'name', 'unknown') logger.info(f"Processing uploaded image: {file_name}") try: content = self._get_file_content(image_file) with self._temp_file(content, ".jpg") as temp_path: encoded_image = self._encode_image(temp_path) base64_url = f"data:image/jpeg;base64,{encoded_image}" response = self._call_ocr_api({"type": "image_url", "image_url": base64_url}) return self._extract_markdown(response) except Exception as e: return self._handle_error("uploaded image processing", e) def document_understanding(self, doc_url: str, question: str) -> str: logger.info(f"Document understanding - URL: {doc_url}, Question: {question}") try: messages = [{"role": "user", "content": [ {"type": "text", "text": question}, {"type": "document_url", "document_url": doc_url} ]}] response = self._call_chat_complete(model="mistral-small-latest", messages=messages) return response.choices[0].message.content if response.choices else "No response received" except Exception as e: return self._handle_error("document understanding", e) def structured_ocr(self, image_file: Union[str, object]) -> str: file_name = getattr(image_file, 'name', 'unknown') logger.info(f"Processing structured OCR for: {file_name}") try: content = self._get_file_content(image_file) with self._temp_file(content, ".jpg") as temp_path: encoded_image = self._encode_image(temp_path) base64_url = f"data:image/jpeg;base64,{encoded_image}" ocr_response = self._call_ocr_api({"type": "image_url", "image_url": base64_url}) markdown = self._extract_markdown(ocr_response) chat_response = self._call_chat_complete( model="pixtral-12b-latest", messages=[{ "role": "user", "content": [ {"type": "image_url", "image_url": base64_url}, {"type": "text", "text": ( f"OCR result:\n\n{markdown}\n\n" "Convert to structured JSON with file_name, topics, languages, and ocr_contents" )} ] }], response_format={"type": "json_object"}, temperature=0 ) content = chat_response.choices[0].message.content if chat_response.choices else "{}" try: response_dict = json.loads(content) if isinstance(response_dict, list): # Handle unexpected list response response_dict = response_dict[0] if response_dict else {} except json.JSONDecodeError: logger.error("Invalid JSON response from chat API") response_dict = {} return self._format_structured_response(temp_path, response_dict) except Exception as e: return self._handle_error("structured OCR", e) @staticmethod def _extract_markdown(response: OCRResponse) -> str: try: return response.pages[0].markdown if response.pages else "No text extracted" except AttributeError: return "Invalid OCR response format" @staticmethod def _handle_error(context: str, error: Exception) -> str: logger.error(f"Error in {context}: {str(error)}") return f"**Error:** {str(error)}" @staticmethod def _format_structured_response(file_path: str, content: Dict) -> str: languages = {lang.alpha_2: lang.name for lang in pycountry.languages if hasattr(lang, 'alpha_2')} valid_langs = [l for l in (content.get("languages") or [DEFAULT_LANGUAGE]) if l in languages.values()] response = { "file_name": Path(file_path).name, "topics": content.get("topics", []), "languages": valid_langs or [DEFAULT_LANGUAGE], "ocr_contents": content.get("ocr_contents", {}) } return f"```json\n{json.dumps(response, indent=4)}\n```" def create_interface(): with gr.Blocks(title="Mistral OCR & Structured Output App") as demo: gr.Markdown("# Mistral OCR & Structured Output App") gr.Markdown("Enter your Mistral API key below to use the app. Extract text from PDFs and images or get structured JSON output.") api_key_input = gr.Textbox( label="Mistral API Key", placeholder="Enter your Mistral API key here", type="password" ) def initialize_processor(api_key): try: processor = OCRProcessor(api_key) return processor, "**Success:** API key set and validated!" except ValueError as e: return None, f"**Error:** {str(e)}" except Exception as e: return None, f"**Error:** Unexpected error: {str(e)}" processor_state = gr.State(value=None) api_status = gr.Markdown("API key not set. Please enter and set your key.") set_api_button = gr.Button("Set API Key") set_api_button.click( fn=initialize_processor, inputs=api_key_input, outputs=[processor_state, api_status] ) tabs = [ ("OCR with PDF URL", gr.Textbox, "ocr_pdf_url", "PDF URL", None), ("OCR with Uploaded PDF", gr.File, "ocr_uploaded_pdf", "Upload PDF", SUPPORTED_PDF_TYPES), ("OCR with Image URL", gr.Textbox, "ocr_image_url", "Image URL", None), ("OCR with Uploaded Image", gr.File, "ocr_uploaded_image", "Upload Image", SUPPORTED_IMAGE_TYPES), ("Structured OCR", gr.File, "structured_ocr", "Upload Image", SUPPORTED_IMAGE_TYPES), ] for name, input_type, fn_name, label, file_types in tabs: with gr.Tab(name): if input_type == gr.Textbox: inputs = input_type(label=label, placeholder=f"e.g., https://example.com/{label.lower().replace(' ', '')}") else: inputs = input_type(label=label, file_types=file_types) output = gr.Markdown(label="Result") button_label = name.replace("OCR with ", "").replace("Structured ", "Get Structured ") def process_with_api(processor, input_data): if not processor: return "**Error:** Please set a valid API key first." fn = getattr(processor, fn_name) return fn(input_data) gr.Button(f"Process {button_label}").click( fn=process_with_api, inputs=[processor_state, inputs], outputs=output ) with gr.Tab("Document Understanding"): doc_url = gr.Textbox(label="Document URL", placeholder="e.g., https://arxiv.org/pdf/1805.04770") question = gr.Textbox(label="Question", placeholder="e.g., What is the last sentence?") output = gr.Markdown(label="Answer") def doc_understanding_with_api(processor, url, q): if not processor: return "**Error:** Please set a valid API key first." return processor.document_understanding(url, q) gr.Button("Ask Question").click( fn=doc_understanding_with_api, inputs=[processor_state, doc_url, question], outputs=output ) return demo if __name__ == "__main__": create_interface().launch(share=True, debug=True)