litellm/tests/ocr_tests/test_ocr_matrix.py

"""Live provider x auth x input coverage for ``litellm.ocr`` / ``litellm.aocr``.

Each ``Case`` is one hand-picked cell, not the full cross product: every provider
exercises each of its credential kinds in both ``explicit`` (kwargs) and ``env``
(monkeypatched environment) mode at least once, and every input kind a provider
accepts is exercised at least once. Sync and async are spread across the cells.
Every cell also checks the success callback saw the same response and cost.
"""

from __future__ import annotations

import asyncio
import base64
import io
import os
import re
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from types import MappingProxyType
from typing import Final, Literal

import pytest

import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.base_llm.ocr.transformation import OCRResponse

Document = Mapping[str, object]
AuthMode = Literal["explicit", "env"]
CallStyle = Literal["sync", "async"]


@dataclass(frozen=True, slots=True)
class LoggedCall:
    payload: Mapping[str, object]
    response: object


class RecordingLogger(CustomLogger):
    def __init__(self) -> None:
        super().__init__()  # pyright: ignore[reportUnknownMemberType]  # CustomLogger.__init__ is untyped
        self.calls: Final[list[LoggedCall]] = []  # mutable-ok: append-only sink the callback hooks write into

    def _record(self, kwargs: Mapping[str, object], response_obj: object) -> None:
        payload: Final = _string_keyed(kwargs.get("standard_logging_object"))
        self.calls.append(LoggedCall(payload, response_obj))

    def log_success_event(
        self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
    ) -> None:
        self._record(kwargs, response_obj)

    async def async_log_success_event(
        self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
    ) -> None:
        self._record(kwargs, response_obj)

    async def wait_for_call(self, timeout: float = 10.0) -> LoggedCall:
        deadline: Final = asyncio.get_running_loop().time() + timeout
        while not self.calls:
            assert asyncio.get_running_loop().time() < deadline, "success callback never fired"
            await asyncio.sleep(0.05)
        assert len(self.calls) == 1, self.calls
        return self.calls[0]


@pytest.fixture
def logger(monkeypatch: pytest.MonkeyPatch) -> RecordingLogger:
    recorder: Final = RecordingLogger()
    monkeypatch.setattr(litellm, "callbacks", [recorder])
    for registry in ("success_callback", "_async_success_callback", "failure_callback", "_async_failure_callback"):
        monkeypatch.setattr(litellm, registry, [])
    return recorder


TESTS_DIR: Final = Path(__file__).resolve().parents[1]
PDF_PATH: Final = TESTS_DIR / "llm_translation" / "fixtures" / "dummy.pdf"
PNG_PATH: Final = TESTS_DIR / "image_gen_tests" / "test_image.png"
PINNED_CDN: Final = "https://cdn.jsdelivr.net/gh/BerriAI/litellm@d769e81c90d453240c61fc572cdb27fae06a89d0"
PDF_URL: Final = f"{PINNED_CDN}/tests/llm_translation/fixtures/dummy.pdf"
PNG_URL: Final = f"{PINNED_CDN}/tests/image_gen_tests/test_image.png"
PDF_TEXT: Final = "Test PDF File"
PNG_TEXT: Final = "LiteLLM"


class _NamedReader(io.BytesIO):
    def __init__(self, path: Path) -> None:
        super().__init__(path.read_bytes())
        self.name: Final = path.name


def _data_uri(path: Path, mime: str) -> str:
    return f"data:{mime};base64,{base64.b64encode(path.read_bytes()).decode()}"


@dataclass(frozen=True, slots=True)
class Input:
    id: str
    build: Callable[[], Document]
    expected_text: str


PDF_BY_URL: Final = Input("pdf_url", lambda: {"type": "document_url", "document_url": PDF_URL}, PDF_TEXT)
PNG_BY_URL: Final = Input("image_url", lambda: {"type": "image_url", "image_url": PNG_URL}, PNG_TEXT)
PDF_DATA_URI: Final = Input(
    "pdf_data_uri",
    lambda: {"type": "document_url", "document_url": _data_uri(PDF_PATH, "application/pdf")},
    PDF_TEXT,
)
PNG_DATA_URI: Final = Input(
    "image_data_uri", lambda: {"type": "image_url", "image_url": _data_uri(PNG_PATH, "image/png")}, PNG_TEXT
)
PDF_AS_PATH: Final = Input("pdf_path", lambda: {"type": "file", "file": PDF_PATH}, PDF_TEXT)
PDF_AS_BYTES: Final = Input(
    "pdf_bytes", lambda: {"type": "file", "file": PDF_PATH.read_bytes(), "mime_type": "application/pdf"}, PDF_TEXT
)
PNG_AS_BYTES: Final = Input(
    "image_bytes", lambda: {"type": "file", "file": PNG_PATH.read_bytes(), "mime_type": "image/png"}, PNG_TEXT
)
PNG_AS_FILE_OBJECT: Final = Input(
    "image_file_object", lambda: {"type": "file", "file": _NamedReader(PNG_PATH)}, PNG_TEXT
)


@dataclass(frozen=True, slots=True)
class Secret:
    """One credential value: the ``litellm.ocr`` kwarg it travels in, the env var litellm reads
    when the kwarg is omitted, and the env var that holds the value in the test process."""

    kwarg: str
    env: str
    source: str | None = None

    @property
    def source_env(self) -> str:
        return self.source or self.env


@dataclass(frozen=True, slots=True)
class Credential:
    id: str
    secrets: tuple[Secret, ...]


@dataclass(frozen=True, slots=True)
class Provider:
    id: str
    model: str
    credentials: tuple[Credential, ...]
    params: Mapping[str, str] = MappingProxyType({})

    @property
    def env_vars(self) -> frozenset[str]:
        return frozenset(secret.env for credential in self.credentials for secret in credential.secrets)


MISTRAL_KEY: Final = Credential("api_key", (Secret("api_key", "MISTRAL_API_KEY"),))
COHERE_KEY: Final = Credential("api_key", (Secret("api_key", "COHERE_API_KEY"),))
REDUCTO_KEY: Final = Credential("api_key", (Secret("api_key", "REDUCTO_API_KEY"),))

AZURE_ENTRA_SECRETS: Final = (
    Secret("tenant_id", "AZURE_TENANT_ID", "AZURE_FOUNDRY_TENANT_ID"),
    Secret("client_id", "AZURE_CLIENT_ID", "AZURE_FOUNDRY_ADMIN_CLIENT_ID"),
    Secret("client_secret", "AZURE_CLIENT_SECRET", "AZURE_FOUNDRY_ADMIN_CLIENT_SECRET"),
)
AZURE_AI_BASE: Final = Secret("api_base", "AZURE_AI_API_BASE")
AZURE_AI_KEY: Final = Credential("api_key", (AZURE_AI_BASE, Secret("api_key", "AZURE_AI_API_KEY")))
AZURE_AI_ENTRA: Final = Credential("entra", (AZURE_AI_BASE, *AZURE_ENTRA_SECRETS))

AZURE_DI_BASE: Final = Secret("api_base", "AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT")
AZURE_DI_KEY: Final = Credential("api_key", (AZURE_DI_BASE, Secret("api_key", "AZURE_DOCUMENT_INTELLIGENCE_API_KEY")))
AZURE_DI_ENTRA: Final = Credential("entra", (AZURE_DI_BASE, *AZURE_ENTRA_SECRETS))

VERTEX_SERVICE_ACCOUNT: Final = Credential(
    "service_account",
    (Secret("vertex_credentials", "VERTEXAI_CREDENTIALS"), Secret("vertex_project", "VERTEXAI_PROJECT")),
)

MISTRAL: Final = Provider("mistral", "mistral/mistral-ocr-latest", (MISTRAL_KEY,))
AZURE_AI_MISTRAL: Final = Provider(
    "azure_ai_mistral", "azure_ai/mistral-document-ai-2512", (AZURE_AI_KEY, AZURE_AI_ENTRA)
)
AZURE_DOC_INTELLIGENCE: Final = Provider(
    "azure_doc_intelligence", "azure_ai/doc-intelligence/prebuilt-layout", (AZURE_DI_KEY, AZURE_DI_ENTRA)
)
COHERE: Final = Provider("cohere", "cohere/parse-v5.0", (COHERE_KEY,))
REDUCTO_V3: Final = Provider("reducto_v3", "reducto/parse-v3", (REDUCTO_KEY,))
REDUCTO_LEGACY: Final = Provider("reducto_legacy", "reducto/parse-legacy", (REDUCTO_KEY,))
VERTEX_MISTRAL: Final = Provider(
    "vertex_mistral",
    "vertex_ai/mistral-ocr-2505",
    (VERTEX_SERVICE_ACCOUNT,),
    MappingProxyType({"vertex_location": "us-central1"}),
)


@dataclass(frozen=True, slots=True)
class Case:
    provider: Provider
    credential: Credential
    auth: AuthMode
    document: Input
    call: CallStyle

    @property
    def id(self) -> str:
        return f"{self.provider.id}-{self.credential.id}-{self.auth}-{self.document.id}-{self.call}"

    def bind_credentials(self, monkeypatch: pytest.MonkeyPatch) -> Mapping[str, str]:
        """Clear every env var the provider could fall back to, then supply this case's values via kwargs or env."""
        values: Final = {secret: os.environ.get(secret.source_env) for secret in self.credential.secrets}
        missing: Final = tuple(secret.source_env for secret, value in values.items() if not value)
        if missing:
            pytest.skip(f"{', '.join(missing)} not set")
        for env_var in self.provider.env_vars:
            monkeypatch.delenv(env_var, raising=False)
        if self.auth == "explicit":
            return {secret.kwarg: value for secret, value in values.items() if value}
        for secret, value in values.items():
            monkeypatch.setenv(secret.env, value or "")
        return {}

    async def run(self, credentials: Mapping[str, str]) -> OCRResponse:
        kwargs: Final = {**self.provider.params, **credentials}
        document: Final = self.document.build()
        response: Final = (
            await litellm.aocr(model=self.provider.model, document=document, **kwargs)  # pyright: ignore[reportUnknownMemberType]  # @client erases the signature
            if self.call == "async"
            else litellm.ocr(model=self.provider.model, document=document, **kwargs)
        )
        assert isinstance(response, OCRResponse)
        return response


CASES: Final = (
    Case(MISTRAL, MISTRAL_KEY, "explicit", PDF_BY_URL, "sync"),
    Case(MISTRAL, MISTRAL_KEY, "env", PNG_BY_URL, "async"),
    Case(MISTRAL, MISTRAL_KEY, "explicit", PDF_AS_PATH, "sync"),
    Case(MISTRAL, MISTRAL_KEY, "explicit", PNG_AS_BYTES, "async"),
    Case(MISTRAL, MISTRAL_KEY, "explicit", PNG_AS_FILE_OBJECT, "sync"),
    Case(AZURE_AI_MISTRAL, AZURE_AI_KEY, "explicit", PDF_BY_URL, "sync"),
    Case(AZURE_AI_MISTRAL, AZURE_AI_KEY, "env", PNG_BY_URL, "async"),
    Case(AZURE_AI_MISTRAL, AZURE_AI_ENTRA, "explicit", PDF_AS_PATH, "sync"),
    Case(AZURE_AI_MISTRAL, AZURE_AI_ENTRA, "env", PDF_DATA_URI, "async"),
    Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_KEY, "explicit", PDF_BY_URL, "sync"),
    Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_KEY, "env", PNG_AS_BYTES, "async"),
    Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_ENTRA, "explicit", PNG_BY_URL, "async"),
    Case(AZURE_DOC_INTELLIGENCE, AZURE_DI_ENTRA, "env", PDF_AS_PATH, "sync"),
    Case(COHERE, COHERE_KEY, "explicit", PNG_BY_URL, "sync"),
    Case(COHERE, COHERE_KEY, "env", PNG_DATA_URI, "async"),
    Case(REDUCTO_V3, REDUCTO_KEY, "explicit", PDF_AS_PATH, "sync"),
    Case(REDUCTO_V3, REDUCTO_KEY, "env", PNG_AS_BYTES, "async"),
    Case(REDUCTO_V3, REDUCTO_KEY, "explicit", PDF_DATA_URI, "async"),
    Case(REDUCTO_LEGACY, REDUCTO_KEY, "explicit", PDF_AS_BYTES, "sync"),
    Case(VERTEX_MISTRAL, VERTEX_SERVICE_ACCOUNT, "explicit", PDF_BY_URL, "sync"),
    Case(VERTEX_MISTRAL, VERTEX_SERVICE_ACCOUNT, "env", PNG_BY_URL, "async"),
)


def _response_cost(response: OCRResponse) -> float:
    response_cost: Final[object] = response._hidden_params.get("response_cost")  # pyright: ignore[reportPrivateUsage, reportUnknownMemberType, reportUnknownVariableType]  # response_cost is only surfaced on _hidden_params
    assert isinstance(response_cost, float) and response_cost > 0
    return response_cost


def _assert_ocr_response(response: OCRResponse, model: str, expected_text: str) -> None:
    assert response.object == "ocr"
    assert response.model == model.split("/", 1)[1]
    assert [page.index for page in response.pages] == list(range(len(response.pages)))
    text: Final = re.sub(r"\s+", " ", " ".join(page.markdown for page in response.pages))
    assert expected_text.lower() in text.lower(), text
    assert response.usage_info is not None
    assert response.usage_info.pages_processed == len(response.pages)
    _response_cost(response)


def _string_keyed(value: object) -> Mapping[str, object]:
    assert isinstance(value, Mapping), type(value)
    items: Final = tuple(value.items())  # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType, reportUnknownArgumentType]  # narrowed from object
    return MappingProxyType({str(key): value for key, value in items})  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]  # narrowed from object


def _assert_logged(logged: LoggedCall, response: OCRResponse, model: str, logged_model: str, call: CallStyle) -> None:
    assert isinstance(logged.response, OCRResponse)
    assert logged.response.pages == response.pages
    assert logged.payload["status"] == "success"
    assert logged.payload["call_type"] == ("aocr" if call == "async" else "ocr")
    assert logged.payload["custom_llm_provider"] == model.split("/", 1)[0]
    assert logged.payload["model"] == logged_model
    assert logged.payload["response_cost"] == _response_cost(response)


@pytest.mark.parametrize("case", CASES, ids=[case.id for case in CASES])
async def test_ocr(case: Case, monkeypatch: pytest.MonkeyPatch, logger: RecordingLogger) -> None:
    credentials: Final = case.bind_credentials(monkeypatch)
    response: Final = await case.run(credentials)
    _assert_ocr_response(response, case.provider.model, case.document.expected_text)
    _assert_logged(await logger.wait_for_call(), response, case.provider.model, response.model, case.call)