pdf/gateway/app/services/convert/ocr/mistral_hook.py

"""Mistral OCR product-feature client for convert rebuild.

Mistral OCR remains an in-product DocQube feature. Never sends files to ConvertAPI/Aspose.
Secrets (API keys) are never logged.
"""

from __future__ import annotations

import base64
import json
import os
import urllib.error
import urllib.request
from typing import Any

from app.services import logs, retry
from app.services.convert.idm.model import BBox, Block, BlockType, TextSpan

logger = logs.get_logger(__name__)


def _retry_after_hint(exc: BaseException) -> float | None:
    """Honour a 429/503's ``Retry-After`` rather than guessing at the delay."""
    if isinstance(exc, urllib.error.HTTPError):
        try:
            return retry.retry_after_seconds(exc.headers.get("Retry-After"))
        except Exception:
            return None
    return None


def mistral_configured() -> bool:
    return bool(os.environ.get("MISTRAL_API_KEY") or os.environ.get("DOCQUBE_MISTRAL_OCR_URL"))


def _endpoint() -> str:
    return (os.environ.get("DOCQUBE_MISTRAL_OCR_URL") or "https://api.mistral.ai/v1/ocr").rstrip(
        "/"
    )


def _map_payload_to_blocks(payload: dict[str, Any], start_order: int) -> list[Block]:
    blocks: list[Block] = []
    order = start_order
    pages = payload.get("pages") or payload.get("data") or [payload]
    if isinstance(pages, dict):
        pages = [pages]
    for page in pages:
        # Tables first if present
        for table in page.get("tables") or []:
            cells = table.get("cells") or table.get("rows") or []
            if cells and isinstance(cells[0], dict):
                # sparse cell list → skip complex; prefer rows
                continue
            if cells:
                blocks.append(
                    Block(
                        type=BlockType.table,
                        cells=cells,
                        table_confidence=float(table.get("confidence") or 0.8),
                        reading_order=order,
                    )
                )
                order += 1
        for line in page.get("lines") or page.get("markdown_lines") or []:
            if isinstance(line, str):
                text = line.strip()
                bbox = BBox()
            else:
                text = str(line.get("text") or line.get("content") or "").strip()
                b = line.get("bbox") or line.get("box") or {}
                bbox = BBox(
                    x=float(b.get("x", 0)),
                    y=float(b.get("y", 0)),
                    w=float(b.get("w", b.get("width", 0))),
                    h=float(b.get("h", b.get("height", 0))),
                )
            if not text:
                continue
            blocks.append(
                Block(
                    type=BlockType.paragraph,
                    text=text,
                    spans=[TextSpan(text=text)],
                    bbox=bbox,
                    reading_order=order,
                )
            )
            order += 1
        # Fallback: single markdown / text blob
        if not blocks:
            md = (page.get("markdown") or page.get("text") or "").strip()
            if md:
                for i, para in enumerate(md.splitlines()):
                    t = para.strip()
                    if t:
                        blocks.append(
                            Block(
                                type=BlockType.paragraph,
                                text=t,
                                spans=[TextSpan(text=t)],
                                reading_order=start_order + i,
                            )
                        )
    return blocks


def ocr_image_to_blocks_mistral(
    image_bytes: bytes,
    start_order: int = 0,
    *,
    timeout: float = 30.0,
    _opener=None,
) -> list[Block]:
    """
    Call Mistral / DocQube product OCR HTTPS endpoint and map to IDM blocks.
    Retries once on transient failure. Returns [] if not configured or on error.
    """
    if not mistral_configured():
        return []

    # Size guard
    if len(image_bytes) > 12_000_000:
        return []

    url = _endpoint()
    b64 = base64.b64encode(image_bytes).decode("ascii")
    body = {
        "model": os.environ.get("MISTRAL_OCR_MODEL", "mistral-ocr-latest"),
        "document": {
            "type": "image_url",
            "image_url": f"data:image/png;base64,{b64}",
        },
    }
    data = json.dumps(body).encode("utf-8")
    headers = {"Content-Type": "application/json", "Accept": "application/json"}
    key = os.environ.get("MISTRAL_API_KEY")
    if key:
        headers["Authorization"] = f"Bearer {key}"

    def _attempt() -> list[Block]:
        req = urllib.request.Request(url, data=data, headers=headers, method="POST")
        opener = _opener or urllib.request.urlopen
        with opener(req, timeout=timeout) as resp:
            raw = resp.read()
        payload = json.loads(raw.decode("utf-8"))
        return _map_payload_to_blocks(payload, start_order)

    # The previous shape here retried every failure twice, back to back, with
    # no delay: a 401 was retried (it will never succeed), and a 429 was
    # retried *instantly* with its ``Retry-After`` ignored, which adds load to
    # a service that has just said it has too much. The accumulated error was
    # then discarded, so a total provider outage reached the caller as a page
    # with no text on it and nothing anywhere saying why.
    try:
        return retry.call_with_retry(
            _attempt,
            label="mistral-ocr",
            retry_after=_retry_after_hint,
        )
    except retry.RetryExhausted as exc:
        logs.error(
            logger,
            "mistral_ocr_unavailable",
            attempts=exc.attempts,
            error=type(exc.last).__name__ if exc.last else "unknown",
            detail=str(exc.last)[:200] if exc.last else "",
        )
        return []
    except Exception as exc:  # classified as fatal by the policy above
        logs.error(
            logger,
            "mistral_ocr_failed",
            error=type(exc).__name__,
            detail=str(exc)[:200],
            retryable=False,
        )
        return []