rag/ocr.py

import os
import base64
from typing import List
from mistralai import Mistral


def _pdf_to_base64(pdf_path: str) -> str:
    with open(pdf_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def mistral_ocr_pdf(pdf_bytes: bytes) -> List[str]:
    """
    OCR PDF bytes directly (no filesystem).
    """
    print("🚀 OCR START | Mistral (in-memory)")

    client = Mistral(api_key=os.environ["MISTRAL_API_KEY"])
    b64 = base64.b64encode(pdf_bytes).decode("utf-8")

    resp = client.ocr.process(
        model="mistral-ocr-latest",
        document={
            "type": "document_url",
            "document_url": f"data:application/pdf;base64,{b64}",
        },
        include_image_base64=False,
    )

    pages = []
    for page in resp.pages:
        if getattr(page, "markdown", None):
            pages.append(page.markdown)
        elif getattr(page, "text", None):
            pages.append(page.text)
        else:
            pages.append("")

    print(f"✅ OCR DONE | pages={len(pages)}")
    return pages