Agentic-Service-Data-Eyond-Catalog

Sleeping

App Files Files Community

Agentic-Service-Data-Eyond-Catalog / src /knowledge_parsing /mistral_ocr.py

ishaq101

sofhiaazzhr

/feat knowledge management (#20)

b68816f about 2 months ago

Raw History Blame Contribute Delete

4.23 kB

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

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