app.py

"""
Author : Janarddan Sarkar
file_name : mistral_ocr_st.py 
date : 10-03-2025
description : Fixed JSON parsing by using manual JSON extraction with response_format=json_object.
"""
import os
import json
import base64
import time
import streamlit as st
from mistralai import Mistral
from dotenv import find_dotenv, load_dotenv
from mistralai import DocumentURLChunk, ImageURLChunk, TextChunk
from mistralai.models import OCRResponse
from mistralai.models.sdkerror import SDKError
from enum import Enum
from pydantic import BaseModel
import pycountry

# Load environment variables
load_dotenv(find_dotenv())
api_key = os.environ.get("MISTRAL_API_KEY")
client = Mistral(api_key=api_key)

# Define Language Enum (保留仅供其他用途，但不在模型中使用)
languages = {lang.alpha_2: lang.name for lang in pycountry.languages if hasattr(lang, 'alpha_2')}

class LanguageMeta(Enum.__class__):
    def __new__(metacls, cls, bases, classdict):
        for code, name in languages.items():
            classdict[name.upper().replace(' ', '_')] = name
        return super().__new__(metacls, cls, bases, classdict)

class Language(Enum, metaclass=LanguageMeta):
    pass

# 不再使用 Pydantic 模型作为 response_format，仅定义结构供参考
class StructuredOCR(BaseModel):
    file_name: str
    topics: list[str]
    languages: list[str]
    ocr_contents: dict

def replace_images_in_markdown(markdown_str: str, images_dict: dict) -> str:
    for img_name, base64_str in images_dict.items():
        markdown_str = markdown_str.replace(f"![{img_name}]({img_name})", f"![{img_name}]({base64_str})")
    return markdown_str

def get_combined_markdown(ocr_response: OCRResponse) -> str:
    markdowns: list[str] = []
    for page in ocr_response.pages:
        image_data = {img.id: img.image_base64 for img in page.images}
        markdowns.append(replace_images_in_markdown(page.markdown, image_data))
    return "\n\n".join(markdowns)

# ---------- 重试装饰器 ----------
def retry_on_rate_limit(max_retries=5, initial_delay=1):
    def decorator(func):
        def wrapper(*args, **kwargs):
            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except SDKError as e:
                    if "429" in str(e) and attempt < max_retries - 1:
                        delay = initial_delay * (2 ** attempt)
                        st.warning(f"Rate limit hit. Retrying in {delay}s... (attempt {attempt+1}/{max_retries})")
                        time.sleep(delay)
                    else:
                        raise
            return None
        return wrapper
    return decorator

@retry_on_rate_limit(max_retries=5, initial_delay=1)
def ocr_process_with_retry(**kwargs):
    return client.ocr.process(**kwargs)

@retry_on_rate_limit(max_retries=3, initial_delay=1)
def chat_complete_with_retry(**kwargs):
    return client.chat.complete(**kwargs)

# ---------- 处理函数 ----------
def process_pdf(pdf_bytes, file_name):
    uploaded_file = client.files.upload(
        file={"file_name": file_name, "content": pdf_bytes},
        purpose="ocr",
    )
    signed_url = client.files.get_signed_url(file_id=uploaded_file.id, expiry=1)
    pdf_response = ocr_process_with_retry(
        document=DocumentURLChunk(document_url=signed_url.url),
        model="mistral-ocr-latest",
        include_image_base64=True,
    )
    if isinstance(pdf_response, dict):
        pdf_response = OCRResponse(**pdf_response)
    return pdf_response

def process_image(image_bytes, file_name):
    """Process an image: returns (markdown_with_images, structured_json)"""
    encoded_image = base64.b64encode(image_bytes).decode()
    base64_data_url = f"data:image/jpeg;base64,{encoded_image}"
    
    # OCR 调用
    image_response = ocr_process_with_retry(
        document=ImageURLChunk(image_url=base64_data_url),
        model="mistral-ocr-latest",
        include_image_base64=True
    )
    page = image_response.pages[0]
    images_dict = {img.id: img.image_base64 for img in page.images if img.image_base64}
    markdown_with_images = replace_images_in_markdown(page.markdown, images_dict)
    
    # 默认 fallback 结构（包含完整的 OCR 文本）
    default_structured = {
        "file_name": file_name,
        "topics": [],
        "languages": [],
        "ocr_contents": {"full_text": markdown_with_images}
    }
    
    try:
        response = chat_complete_with_retry(
            model="pixtral-12b-latest",
            messages=[
                {
                    "role": "user",
                    "content": [
                        ImageURLChunk(image_url=base64_data_url),
                        TextChunk(
                            text=(
                                "Based on the image and its OCR markdown below, output a JSON object with the following fields:\n"
                                "- file_name: string, the original file name.\n"
                                "- topics: list of strings, main subjects.\n"
                                "- languages: list of strings, languages detected (e.g., 'English', 'Chinese').\n"
                                "- ocr_contents: dictionary, where keys are section titles or 'full_text', and values are the corresponding text. Include the full OCR text under key 'full_text'.\n\n"
                                f"OCR Markdown:\n{markdown_with_images}\n\n"
                                "Output ONLY valid JSON, no other text."
                            )
                        ),
                    ],
                },
            ],
            response_format={"type": "json_object"},
            temperature=0,
        )
        content = response.choices[0].message.content
        parsed_json = json.loads(content)
        structured_json = {
            "file_name": parsed_json.get("file_name", file_name),
            "topics": parsed_json.get("topics", []),
            "languages": parsed_json.get("languages", []),
            "ocr_contents": parsed_json.get("ocr_contents", {"full_text": markdown_with_images})
        }
        # 确保类型正确
        if not isinstance(structured_json["topics"], list):
            structured_json["topics"] = []
        if not isinstance(structured_json["languages"], list):
            structured_json["languages"] = []
        if not isinstance(structured_json["ocr_contents"], dict):
            structured_json["ocr_contents"] = {"full_text": markdown_with_images}
            
    except Exception as e:
        st.warning(f"Could not parse structured JSON from model: {e}. Using fallback OCR text as ocr_contents.")
        structured_json = default_structured
    
    return markdown_with_images, structured_json

# ---------- Streamlit UI ----------
st.title("OLMOCR")

uploaded_file = st.file_uploader("Upload a PDF or Image", type=["pdf", "png", "jpg", "jpeg"])

if uploaded_file:
    file_type = uploaded_file.type
    file_bytes = uploaded_file.read()
    file_name = uploaded_file.name

    if st.button("Submit"):
        st.write(f"**Processing file:** {file_name}")

        if "pdf" in file_type:
            pdf_response = process_pdf(file_bytes, file_name)
            st.markdown(get_combined_markdown(pdf_response))
        else:
            markdown_content, structured_data = process_image(file_bytes, file_name)
            st.markdown(markdown_content)
            with st.expander("📋 Structured Data (JSON)"):
                st.json(structured_data)