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}"