urieljareth
239a7510c6
- 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
289 lines
10 KiB
Python
289 lines
10 KiB
Python
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:
|
|
"""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:
|
|
self._require_key()
|
|
data_url = self._image_data_url(image_path)
|
|
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,
|
|
}
|
|
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'}"
|
|
)
|
|
|
|
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 []
|
|
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 str(data["markdown"]).strip()
|
|
return ""
|
|
|
|
def _image_data_url(self, image_path: Path) -> str:
|
|
encoded = base64.b64encode(image_path.read_bytes()).decode("ascii")
|
|
return f"data:image/jpeg;base64,{encoded}"
|