OCR Mistral: modelo mistral-ocr-4-0, batch + fallback con reintentos y errores etiquetados
- Modelo por defecto mistral-ocr-4-0 (.env.example, config.py, docker-compose) - MistralClient: OCR en lote via API Batch (/v1/ocr) subiendo JSONL, polling y descarga de resultados - Reintentos con backoff exponencial en llamadas a Mistral (429/5xx/red); sin reintento en 401/400 - PdfProcessor: lote cuando >= umbral de paginas, fallback sincrono pagina a pagina, placeholder rastreable por pagina fallida - Errores etiquetados ([AUTH_MISTRAL], [BATCH], [NETWORK], [OCR_EMPTY], [PDF_PARSE]...) en el mensaje del job - Nuevas variables: MISTRAL_OCR_BATCH_THRESHOLD/POLL_SECONDS/TIMEOUT_SECONDS/RETRY_ATTEMPTS/RETRY_BACKOFF_BASE
This commit is contained in:
urieljareth
@@ -1,45 +1,287 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
|
||||
from backend.app.config import settings
|
||||
from backend.app.services.errors import (
|
||||
MistralAuthError,
|
||||
MistralBatchError,
|
||||
MistralOCREmptyError,
|
||||
MistralRateLimitError,
|
||||
MistralServerError,
|
||||
NetworkError,
|
||||
NON_RETRYABLE_CODES,
|
||||
ProcessingError,
|
||||
classify_requests_error,
|
||||
)
|
||||
|
||||
|
||||
class MistralClient:
|
||||
endpoint = "https://api.mistral.ai/v1/ocr"
|
||||
"""Cliente Mistral para OCR.
|
||||
|
||||
- ``ocr_image``: OCR sincrono de una pagina, con reintentos y backoff.
|
||||
- ``ocr_images_batch``: OCR en lote via la API de Batch de Mistral
|
||||
(endpoint ``/v1/ocr``), subiendo un fichero JSONL. Devuelve un dict
|
||||
``{page_index: markdown}``. Si el job falla o caduca se lanza
|
||||
``MistralBatchError`` para que el llamador caiga al modo sincrono.
|
||||
"""
|
||||
|
||||
base_url = "https://api.mistral.ai/v1"
|
||||
ocr_url = f"{base_url}/ocr"
|
||||
files_url = f"{base_url}/files"
|
||||
batch_url = f"{base_url}/batch/jobs"
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._api_key = settings.mistral_api_key
|
||||
self._model = settings.mistral_ocr_model
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# API publica
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def ocr_image(self, image_path: Path) -> str:
|
||||
if not settings.mistral_api_key:
|
||||
raise RuntimeError("MISTRAL_API_KEY no esta configurada")
|
||||
self._require_key()
|
||||
data_url = self._image_data_url(image_path)
|
||||
response = requests.post(
|
||||
self.endpoint,
|
||||
headers={
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {settings.mistral_api_key}",
|
||||
},
|
||||
json={
|
||||
"model": settings.mistral_ocr_model,
|
||||
return self._retry(lambda: self._ocr_sync(data_url))
|
||||
|
||||
def ocr_images_batch(self, image_paths: list[Path]) -> dict[int, str]:
|
||||
"""OCR en lote. Devuelve ``{posicion_en_lista: markdown}``.
|
||||
|
||||
``posicion_en_lista`` es el indice (0-based) dentro de ``image_paths``,
|
||||
NO el numero de pagina del PDF; el llamador hace esa traduccion.
|
||||
"""
|
||||
self._require_key()
|
||||
if not image_paths:
|
||||
return {}
|
||||
|
||||
payload = self._build_batch_payload(image_paths)
|
||||
upload = self._retry(lambda: self._upload_batch_file(payload))
|
||||
file_id = upload.get("id")
|
||||
if not file_id:
|
||||
raise MistralBatchError("la subida del fichero batch no devolvio id")
|
||||
|
||||
job = self._retry(lambda: self._create_batch_job(file_id))
|
||||
job_id = job.get("id")
|
||||
if not job_id:
|
||||
raise MistralBatchError("la creacion del job batch no devolvio id")
|
||||
|
||||
final = self._poll_batch_job(job_id)
|
||||
if final.get("status") != "SUCCESS":
|
||||
reason = final.get("errors") or final.get("status") or "sin detalle"
|
||||
raise MistralBatchError(f"job {job_id} finalizo como {final.get('status')}: {reason}")
|
||||
|
||||
output_file_id = final.get("output_file")
|
||||
if not output_file_id:
|
||||
raise MistralBatchError(f"job {job_id} sin output_file")
|
||||
|
||||
content = self._retry(lambda: self._download_file(output_file_id))
|
||||
return self._parse_batch_output(content)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# OCR sincrono
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _ocr_sync(self, data_url: str) -> str:
|
||||
try:
|
||||
response = requests.post(
|
||||
self.ocr_url,
|
||||
headers=self._json_headers(),
|
||||
json={
|
||||
"model": self._model,
|
||||
"document": {"type": "image_url", "image_url": data_url},
|
||||
"include_image_base64": False,
|
||||
},
|
||||
timeout=180,
|
||||
)
|
||||
except requests.RequestException as exc:
|
||||
raise classify_requests_error(exc) from exc
|
||||
self._raise_for_status(response)
|
||||
markdown = self._extract_markdown(response.json())
|
||||
if not markdown:
|
||||
raise MistralOCREmptyError()
|
||||
return markdown
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Batch: construccion, subida, job, polling, descarga, parseo
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _build_batch_payload(self, image_paths: list[Path]) -> bytes:
|
||||
lines: list[str] = []
|
||||
for idx, path in enumerate(image_paths):
|
||||
data_url = self._image_data_url(path)
|
||||
body = {
|
||||
"document": {"type": "image_url", "image_url": data_url},
|
||||
"include_image_base64": False,
|
||||
},
|
||||
timeout=180,
|
||||
}
|
||||
lines.append(
|
||||
json.dumps(
|
||||
{"custom_id": str(idx), "body": body},
|
||||
separators=(",", ":"),
|
||||
ensure_ascii=False,
|
||||
)
|
||||
)
|
||||
return ("\n".join(lines)).encode("utf-8")
|
||||
|
||||
def _upload_batch_file(self, payload: bytes) -> dict:
|
||||
try:
|
||||
response = requests.post(
|
||||
self.files_url,
|
||||
headers={"Authorization": f"Bearer {self._api_key}"},
|
||||
data={"purpose": "batch"},
|
||||
files={"file": ("ocr-batch.jsonl", payload, "application/jsonl")},
|
||||
timeout=300,
|
||||
)
|
||||
except requests.RequestException as exc:
|
||||
raise classify_requests_error(exc) from exc
|
||||
self._raise_for_status(response)
|
||||
return response.json()
|
||||
|
||||
def _create_batch_job(self, file_id: str) -> dict:
|
||||
try:
|
||||
response = requests.post(
|
||||
self.batch_url,
|
||||
headers=self._json_headers(),
|
||||
json={
|
||||
"input_files": [file_id],
|
||||
"model": self._model,
|
||||
"endpoint": "/v1/ocr",
|
||||
},
|
||||
timeout=120,
|
||||
)
|
||||
except requests.RequestException as exc:
|
||||
raise classify_requests_error(exc) from exc
|
||||
self._raise_for_status(response)
|
||||
return response.json()
|
||||
|
||||
def _poll_batch_job(self, job_id: str) -> dict:
|
||||
terminal = {"SUCCESS", "FAILED", "TIMEOUT_EXCEEDED", "CANCELLED"}
|
||||
deadline = time.time() + max(60, settings.mistral_ocr_batch_timeout_seconds)
|
||||
delay = max(2, settings.mistral_ocr_batch_poll_seconds)
|
||||
last: dict = {}
|
||||
while time.time() < deadline:
|
||||
try:
|
||||
response = requests.get(
|
||||
f"{self.batch_url}/{job_id}",
|
||||
headers=self._json_headers(),
|
||||
timeout=60,
|
||||
)
|
||||
except requests.RequestException:
|
||||
time.sleep(delay)
|
||||
continue
|
||||
if response.ok:
|
||||
last = response.json()
|
||||
if last.get("status") in terminal:
|
||||
return last
|
||||
time.sleep(delay)
|
||||
raise MistralBatchError(
|
||||
f"job {job_id} excedio el timeout de polling "
|
||||
f"({settings.mistral_ocr_batch_timeout_seconds}s); ultimo estado: "
|
||||
f"{last.get('status') if last else 'desconocido'}"
|
||||
)
|
||||
if response.status_code == 401:
|
||||
raise RuntimeError("API Key de Mistral invalida")
|
||||
if not response.ok:
|
||||
raise RuntimeError(f"Mistral error {response.status_code}: {response.text}")
|
||||
data = response.json()
|
||||
|
||||
def _download_file(self, file_id: str) -> bytes:
|
||||
try:
|
||||
response = requests.get(
|
||||
f"{self.files_url}/{file_id}/content",
|
||||
headers=self._json_headers(),
|
||||
timeout=300,
|
||||
)
|
||||
except requests.RequestException as exc:
|
||||
raise classify_requests_error(exc) from exc
|
||||
self._raise_for_status(response)
|
||||
return response.content
|
||||
|
||||
def _parse_batch_output(self, content: bytes) -> dict[int, str]:
|
||||
results: dict[int, str] = {}
|
||||
for raw in content.decode("utf-8", errors="replace").splitlines():
|
||||
line = raw.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
obj = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
custom_id = obj.get("custom_id")
|
||||
if custom_id is None:
|
||||
continue
|
||||
try:
|
||||
idx = int(custom_id)
|
||||
except (TypeError, ValueError):
|
||||
continue
|
||||
response = obj.get("response") or {}
|
||||
status_code = response.get("status_code", 200)
|
||||
body = response.get("body") or {}
|
||||
if status_code >= 400 or obj.get("error"):
|
||||
# Pagina fallida en lote: se omite y el llamador cae a sync.
|
||||
continue
|
||||
markdown = self._extract_markdown(body)
|
||||
if markdown:
|
||||
results[idx] = markdown
|
||||
return results
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Helpers compartidos
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _require_key(self) -> None:
|
||||
if not self._api_key:
|
||||
raise ProcessingError("CONFIG", "MISTRAL_API_KEY no esta configurada")
|
||||
|
||||
def _json_headers(self) -> dict[str, str]:
|
||||
return {
|
||||
"Authorization": f"Bearer {self._api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "application/json",
|
||||
}
|
||||
|
||||
def _raise_for_status(self, response: requests.Response) -> None:
|
||||
if response.ok:
|
||||
return
|
||||
try:
|
||||
body = response.text
|
||||
except Exception:
|
||||
body = ""
|
||||
status = response.status_code
|
||||
if status == 401:
|
||||
raise MistralAuthError()
|
||||
if status == 429:
|
||||
raise MistralRateLimitError()
|
||||
if status >= 500:
|
||||
raise MistralServerError(status, body)
|
||||
raise ProcessingError("MISTRAL_CLIENT", f"HTTP {status}: {body[:200]}")
|
||||
|
||||
def _retry(self, fn):
|
||||
attempts = max(1, settings.mistral_ocr_retry_attempts)
|
||||
base_delay = max(0.5, settings.mistral_ocr_retry_backoff_base)
|
||||
last_exc: Exception | None = None
|
||||
for attempt in range(attempts):
|
||||
try:
|
||||
return fn()
|
||||
except ProcessingError as exc:
|
||||
if exc.code in NON_RETRYABLE_CODES:
|
||||
raise
|
||||
last_exc = exc
|
||||
except Exception as exc:
|
||||
last_exc = exc
|
||||
if attempt < attempts - 1:
|
||||
time.sleep(base_delay * (2 ** attempt))
|
||||
if last_exc is not None:
|
||||
raise last_exc
|
||||
raise MistralBatchError("reintentos agotados sin excepcion conocida")
|
||||
|
||||
def _extract_markdown(self, data: dict) -> str:
|
||||
pages = data.get("pages") or []
|
||||
markdown_pages = [page.get("markdown", "").strip() for page in pages if page.get("markdown")]
|
||||
if markdown_pages:
|
||||
return "\n\n".join(markdown_pages)
|
||||
chunks = [page.get("markdown", "").strip() for page in pages if page.get("markdown")]
|
||||
if chunks:
|
||||
return "\n\n".join(chunks)
|
||||
if data.get("markdown"):
|
||||
return data["markdown"].strip()
|
||||
raise RuntimeError("Mistral OCR no devolvio Markdown")
|
||||
return str(data["markdown"]).strip()
|
||||
return ""
|
||||
|
||||
def _image_data_url(self, image_path: Path) -> str:
|
||||
encoded = base64.b64encode(image_path.read_bytes()).decode("ascii")
|
||||
|
||||
Reference in New Issue
Block a user