File size: 4,227 Bytes
b68816f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | """Mistral OCR API call — the Mistral counterpart of MinerU's `do_parse`.
Kept out of parse.py so the API client (upload → ocr.process) stays separate
from the caching/cache-key logic. Returns the raw response dict; parse.py owns
persisting and caching it. normalize_mistral.py turns it into chunks.
No GPU: this is a hosted API. The only precondition is `MISTRAL_API_KEY`. It is
read through `src.config.settings` first — the same path every other credential
in this repo takes, and the one that works in a deployed HuggingFace Space —
then falls back to the environment / repo `.env` so the standalone parser still
runs on a box without the full app config. Either way the GPU check in
parse.require_gpu never runs for this backend.
"""
from __future__ import annotations
import os
from pathlib import Path
from typing import Any
from .config import PipelineConfig
def _key_from_settings() -> str | None:
"""The key from the app Settings singleton, if it can be constructed.
Guarded on purpose: the full `Settings` has required fields (Postgres, Redis,
Langfuse) with no defaults, and the parser runs standalone on boxes that may
not have them. A failure to build Settings must fall through to the
environment, not abort a parse — so any construction error yields None.
"""
try:
from src.config.settings import settings
except Exception: # noqa: BLE001 - Settings needs the full app env; fall back
return None
return getattr(settings, "mistral_api_key", "") or None
def _api_key() -> str:
key = _key_from_settings()
if not key:
key = os.environ.get("MISTRAL_API_KEY")
if not key:
try:
from dotenv import load_dotenv
load_dotenv() # picks up the repo .env, same file the app reads
key = os.environ.get("MISTRAL_API_KEY")
except ImportError:
pass
if not key:
raise RuntimeError(
"MISTRAL_API_KEY is not set.\n"
" -> set it in the app settings / HuggingFace Space secret, or put it "
"in the repo .env (MISTRAL_API_KEY=...) or the environment."
)
return key
def run_ocr(source: Path, cfg: PipelineConfig) -> dict[str, Any]:
"""Upload one PDF and run Mistral OCR. Returns the response as a dict."""
try:
from mistralai import Mistral # type: ignore[attr-defined]
except ImportError: # SDK 2.x moved the class to a submodule
from mistralai.client import Mistral
client = Mistral(api_key=_api_key())
with source.open("rb") as f:
uploaded = client.files.upload(
file={"file_name": source.name, "content": f}, purpose="ocr"
)
kwargs: dict[str, Any] = {
"model": cfg.mistral_model,
"table_format": cfg.mistral_table_format,
"include_blocks": cfg.mistral_include_blocks,
"include_image_base64": cfg.mistral_describe_images,
}
if cfg.mistral_describe_images:
# Per-image description (generated_description). Guarded: an SDK without
# the annotation helpers must degrade to OCR-without-descriptions, not
# abort — exactly as MinerU's `pipeline` backend produces none.
try:
from mistralai.extra import response_format_from_pydantic_model
from pydantic import BaseModel
class _ImageDescription(BaseModel):
description: str
chart_type: str | None = None
kwargs["bbox_annotation_format"] = response_format_from_pydantic_model(
_ImageDescription
)
except Exception: # noqa: BLE001, S110 - optional capability, never fatal
pass
# file_id is the newer path; if this SDK/model rejects it, fall back to a
# signed URL, which works on every version.
try:
resp = client.ocr.process(document={"file_id": uploaded.id}, **kwargs)
except Exception: # noqa: BLE001
signed = client.files.get_signed_url(file_id=uploaded.id)
resp = client.ocr.process(
document={"type": "document_url", "document_url": signed.url}, **kwargs
)
dumped: dict[str, Any] = resp.model_dump()
return dumped
|