app.py

import os
import base64
import gradio as gr
from mistralai import Mistral, DocumentURLChunk, ImageURLChunk, TextChunk
from mistralai.models import OCRResponse
from pathlib import Path
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

# 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:
            models = self.client.models.list()  # Validate API key
            if not models:
                raise ValueError("No models available")
        except Exception as e:
            raise ValueError(f"Invalid API key: {str(e)}")

    @staticmethod
    def _encode_image(image_path: str) -> str:
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')

    @staticmethod
    @contextmanager
    def _temp_file(content: bytes, suffix: str, keep_alive: bool = False) -> str:
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix, dir=tempfile.gettempdir())
        try:
            logger.info(f"Creating temp file: {temp_file.name}")
            temp_file.write(content)
            temp_file.close()
            yield temp_file.name
        finally:
            if not keep_alive and os.path.exists(temp_file.name):
                logger.info(f"Cleaning up temp file: {temp_file.name}")
                os.unlink(temp_file.name)
            else:
                logger.info(f"Keeping temp file alive: {temp_file.name}")

    @retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
    def _call_ocr_api(self, document: Union[DocumentURLChunk, ImageURLChunk]) -> OCRResponse:
        try:
            return self.client.ocr.process(model="mistral-ocr-latest", document=document, include_image_base64=True)
        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, bytes]) -> bytes:
        if isinstance(file_input, str):
            if file_input.startswith("http"):
                response = requests.get(file_input)
                response.raise_for_status()
                return response.content
            else:
                with open(file_input, "rb") as f:
                    return f.read()
        return file_input.read() if hasattr(file_input, 'read') else file_input

    def ocr_pdf_url(self, pdf_url: str) -> str:
        logger.info(f"Processing PDF URL: {pdf_url}")
        try:
            response = self._call_ocr_api(DocumentURLChunk(document_url=pdf_url))
            return self._get_combined_markdown(response)
        except Exception as e:
            return self._handle_error("PDF URL processing", e)

    def ocr_uploaded_pdf(self, pdf_file: Union[str, bytes]) -> tuple[str, 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", keep_alive=True) 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(DocumentURLChunk(document_url=signed_url.url))
                return self._get_combined_markdown(response), temp_path
        except Exception as e:
            return self._handle_error("uploaded PDF processing", e), None

    def ocr_image_url(self, image_url: str) -> str:
        logger.info(f"Processing image URL: {image_url}")
        try:
            response = self._call_ocr_api(ImageURLChunk(image_url=image_url))
            return self._get_combined_markdown(response)
        except Exception as e:
            return self._handle_error("image URL processing", e)

    def ocr_uploaded_image(self, image_file: Union[str, bytes]) -> tuple[str, 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", keep_alive=True) as temp_path:
                encoded_image = self._encode_image(temp_path)
                base64_url = f"data:image/jpeg;base64,{encoded_image}"
                response = self._call_ocr_api(ImageURLChunk(image_url=base64_url))
                return self._get_combined_markdown(response), temp_path
        except Exception as e:
            return self._handle_error("uploaded image processing", e), None

    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": [
                TextChunk(text=question),
                DocumentURLChunk(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, bytes]) -> tuple[str, 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", keep_alive=True) 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(ImageURLChunk(image_url=base64_url))
                markdown = self._get_combined_markdown(ocr_response)

                chat_response = self._call_chat_complete(
                    model="pixtral-12b-latest",
                    messages=[{
                        "role": "user",
                        "content": [
                            ImageURLChunk(image_url=base64_url),
                            TextChunk(text=(
                                f"This is image's OCR in markdown:\n<BEGIN_IMAGE_OCR>\n{markdown}\n<END_IMAGE_OCR>.\n"
                                "Convert this into a sensible structured json response with file_name, topics, languages, and ocr_contents fields"
                            ))
                        ]
                    }],
                    response_format={"type": "json_object"},
                    temperature=0
                )

                response_content = chat_response.choices[0].message.content
                content = json.loads(response_content)
                return self._format_structured_response(temp_path, content), temp_path
        except Exception as e:
            return self._handle_error("structured OCR", e), None

    def _get_combined_markdown(self, response: OCRResponse) -> str:
        markdowns = []
        for page in response.pages:
            image_data = {}
            for img in page.images:
                image_data[img.id] = img.image_base64
            markdown = page.markdown
            for img_name, base64_str in image_data.items():
                markdown = markdown.replace(f"![{img_name}]({img_name})", f"![{img_name}]({base64_str})")
            markdowns.append(markdown)
        return "\n\n".join(markdowns)

    @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')}
        content_languages = content["languages"] if "languages" in content else [DEFAULT_LANGUAGE]
        valid_langs = [l for l in content_languages if l in languages.values()]

        response = {
            "file_name": Path(file_path).name,
            "topics": content["topics"] if "topics" in content else [],
            "languages": valid_langs or [DEFAULT_LANGUAGE],
            "ocr_contents": content["ocr_contents"] if "ocr_contents" in content else {}
        }
        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.")
        gr.Markdown("**Note:** After entering your API key, click 'Set API Key' to validate and use it.")

        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()
        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, None),
            ("OCR with Uploaded PDF", gr.File, "ocr_uploaded_pdf", "Upload PDF", SUPPORTED_PDF_TYPES, gr.File),
            ("OCR with Image URL", gr.Textbox, "ocr_image_url", "Image URL", None, None),
            ("OCR with Uploaded Image", gr.File, "ocr_uploaded_image", "Upload Image", SUPPORTED_IMAGE_TYPES, gr.Image),
            ("Structured OCR", gr.File, "structured_ocr", "Upload Image", SUPPORTED_IMAGE_TYPES, gr.Image),
        ]

        for name, input_type, fn_name, label, file_types, preview_type 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)
                
                with gr.Row():
                    output = gr.Markdown(label="Result")
                    preview = preview_type(label="Preview") if preview_type else None
                
                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.", None
                    fn = getattr(processor, fn_name)
                    return fn(input_data)  # Returns tuple (result, preview_path)

                gr.Button(f"Process {button_label}").click(
                    fn=process_with_api,
                    inputs=[processor_state, inputs],
                    outputs=[output, preview] if preview else [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__":
    print(f"===== Application Startup at {os.environ.get('START_TIME', 'Unknown')} =====")
    create_interface().launch(share=True, debug=True)