office_translator/services/mistral_ocr.py

"""
Mistral OCR client — text extraction for scanned PDFs.

Image-only PDFs have no extractable text layer: PyMuPDF sees empty pages
and the layout-preserving pipeline would output an empty document. This
client calls the Mistral OCR API to recover the text.

API reference (2026): POST https://api.mistral.ai/v1/ocr with Bearer auth,
body {"model", "document": {"type": "document_url", "document_url":
"data:application/pdf;base64,..."}, "pages": [0, 1, ...]}. The response
contains {"pages": [{"index", "markdown", "dimensions"}, ...]}.

The document is sent in chunks of PAGES_PER_REQUEST pages: it bounds the
request payload and lets us report progress page by page.
"""

import base64
import time
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional

import requests

from core.logging import get_logger

logger = get_logger(__name__)

MISTRAL_OCR_URL = "https://api.mistral.ai/v1/ocr"
DEFAULT_MODEL = "mistral-ocr-latest"
PAGES_PER_REQUEST = 8

MISTRAL_INVALID_KEY = "MISTRAL_INVALID_KEY"
MISTRAL_QUOTA_EXCEEDED = "MISTRAL_QUOTA_EXCEEDED"
MISTRAL_TIMEOUT = "MISTRAL_TIMEOUT"
MISTRAL_SERVICE_ERROR = "MISTRAL_SERVICE_ERROR"


class MistralOCRError(Exception):
    """Raised when the Mistral OCR API cannot extract the PDF text."""

    def __init__(
        self, code: str, message: str, details: Optional[Dict[str, Any]] = None
    ):
        self.code = code
        self.message = message
        self.details = details or {}
        super().__init__(self.message)


class MistralOCRClient:
    """Thin synchronous client around the Mistral OCR endpoint."""

    def __init__(
        self,
        api_key: str,
        model: str = DEFAULT_MODEL,
        timeout: int = 180,
        max_retries: int = 2,
        retry_delay: float = 2.0,
    ):
        self._api_key = (api_key or "").strip()
        self._model = model or DEFAULT_MODEL
        self._timeout = timeout
        self._max_retries = max_retries
        self._retry_delay = retry_delay

    def is_available(self) -> bool:
        """True when an API key is configured."""
        return bool(self._api_key)

    def _post_ocr(self, data_uri: str, pages: List[int]) -> List[Dict[str, Any]]:
        """POST one OCR request for the given 0-based page list, with retries."""
        payload = {
            "model": self._model,
            "document": {"type": "document_url", "document_url": data_uri},
            "pages": pages,
        }
        headers = {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
        }

        last_error: Optional[Exception] = None
        for attempt in range(self._max_retries + 1):
            try:
                response = requests.post(
                    MISTRAL_OCR_URL,
                    json=payload,
                    headers=headers,
                    timeout=self._timeout,
                )

                if response.status_code == 401:
                    raise MistralOCRError(
                        MISTRAL_INVALID_KEY,
                        "Clé API Mistral invalide (MISTRAL_API_KEY).",
                        {"status_code": 401},
                    )
                if response.status_code in (402, 429):
                    raise MistralOCRError(
                        MISTRAL_QUOTA_EXCEEDED,
                        "Quota Mistral OCR épuisé ou limite de débit atteinte.",
                        {"status_code": response.status_code},
                    )
                if response.status_code >= 500:
                    raise MistralOCRError(
                        MISTRAL_SERVICE_ERROR,
                        f"Service Mistral OCR indisponible (HTTP {response.status_code}).",
                        {"status_code": response.status_code},
                    )
                if response.status_code != 200:
                    raise MistralOCRError(
                        MISTRAL_SERVICE_ERROR,
                        f"Erreur Mistral OCR (HTTP {response.status_code}): {response.text[:200]}",
                        {"status_code": response.status_code},
                    )

                pages_out = response.json().get("pages", [])
                if not pages_out:
                    raise MistralOCRError(
                        MISTRAL_SERVICE_ERROR,
                        "Réponse Mistral OCR vide.",
                    )
                return pages_out

            except MistralOCRError as e:
                if e.code in (MISTRAL_INVALID_KEY, MISTRAL_QUOTA_EXCEEDED):
                    raise  # not transient
                last_error = e
            except requests.exceptions.Timeout as e:
                last_error = MistralOCRError(
                    MISTRAL_TIMEOUT,
                    f"Délai d'attente Mistral OCR dépassé ({self._timeout}s).",
                )
            except requests.exceptions.RequestException as e:
                last_error = MistralOCRError(
                    MISTRAL_SERVICE_ERROR,
                    f"Mistral OCR injoignable: {str(e)[:150]}",
                )

            if attempt < self._max_retries:
                delay = self._retry_delay * (2**attempt)
                logger.info(
                    "mistral_ocr_retry",
                    attempt=attempt + 1,
                    delay_s=round(delay, 2),
                    error=last_error.code if last_error else "unknown",
                )
                time.sleep(delay)

        raise last_error or MistralOCRError(
            MISTRAL_SERVICE_ERROR, "Erreur Mistral OCR inconnue."
        )

    def extract_pdf_text(
        self,
        pdf_path: Path,
        progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
    ) -> List[str]:
        """OCR a PDF file and return one text (markdown) string per page.

        Pages are processed in chunks of ``PAGES_PER_REQUEST``; the returned
        list is ordered by page index, empty strings for pages OCR returned
        nothing for.
        """
        import fitz

        pdf_path = Path(pdf_path)
        data_b64 = base64.b64encode(pdf_path.read_bytes()).decode("ascii")
        data_uri = f"data:application/pdf;base64,{data_b64}"

        with fitz.open(str(pdf_path)) as doc:
            total_pages = len(doc)
        if total_pages == 0:
            raise MistralOCRError(MISTRAL_SERVICE_ERROR, "PDF vide (0 page).")

        chunks = [
            list(range(start, min(start + PAGES_PER_REQUEST, total_pages)))
            for start in range(0, total_pages, PAGES_PER_REQUEST)
        ]

        page_texts: List[str] = [""] * total_pages
        done_chunks = 0
        for chunk in chunks:
            pages_out = self._post_ocr(data_uri, chunk)
            for page in pages_out:
                idx = int(page.get("index", -1))
                if 0 <= idx < total_pages:
                    page_texts[idx] = page.get("markdown", "") or ""

            done_chunks += 1
            logger.info(
                "mistral_ocr_chunk_done",
                pages_done=min(done_chunks * PAGES_PER_REQUEST, total_pages),
                total_pages=total_pages,
            )
            if progress_callback and chunks:
                pct = int(5 + 20 * done_chunks / len(chunks))
                progress_callback(
                    {
                        "current": done_chunks,
                        "total": len(chunks),
                        "phase": f"OCR (Mistral) {min(done_chunks * PAGES_PER_REQUEST, total_pages)}/{total_pages}",
                        "paragraph": done_chunks,
                        "total_paragraphs": len(chunks),
                        "progress_override": pct,
                    }
                )

        extracted = sum(1 for t in page_texts if t.strip())
        logger.info(
            "mistral_ocr_extracted",
            pages_with_text=extracted,
            total_pages=total_pages,
        )
        return page_texts