litellm/tests/e2e/llm_translation/test_ocr_rust_e2e.py

"""Live e2e: Rust-backed OCR is reachable through the gateway across providers.

Each provider's OCR deployment is registered at runtime via /model/new and deleted
on teardown, so nothing is hardcoded into the gateway config. Every provider is its
own typed OcrProvider below: it owns the model id and the os.environ/* credential
references the proxy resolves at call time, so adding a provider is a new type
rather than another inline body. Start the proxy with the Rust OCR path enabled:

Each case creates its deployment, drives a real /v1/ocr call, and asserts a
well-formed OCR document comes back. Per the e2e hard-fail contract, a case
fails when no proxy answers and also fails once a request reaches it: the proxy
fetches each provider's referenced secrets, so a
missing credential surfaces as a live provider error rather than silent green.

The provider keys are shared with other pipelines, so a provider's rate limit can
hold across the bounded retries. The case then accepts the gateway's faithful relay
of that 429 (throttling_error, code 429) as its second expected outcome; any other
non-success still fails at once.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Final, Protocol

import pytest
from e2e_config import unique_marker
from e2e_metadata import Domain, Mode, Provider, Route, Subject, meta
from e2e_http import (
    PROVIDER_RATE_LIMIT_ATTEMPTS,
    RateLimitedError,
    Success,
    assert_client_error,
    relayed_provider_rate_limit,
    tolerate_provider_rate_limit,
)
from lifecycle import ResourceManager
from models import LiteLLMParamsBody, OcrBody, OcrDocument, OcrResponse
from proxy_client import ProxyClient
from pydantic import BaseModel

pytestmark = pytest.mark.e2e


class _OptionalOcrBody(BaseModel):
    model: str | None = None
    document: dict[str, object] | None = None


# Tiny in-repo fixtures served via jsdelivr (sha-pinned, immutable) so the request
# bodies stay stable across runs.
TEST_PDF_URL = (
    "https://cdn.jsdelivr.net/gh/BerriAI/litellm"
    "@d769e81c90d453240c61fc572cdb27fae06a89d0"
    "/tests/llm_translation/fixtures/dummy.pdf"
)
TEST_IMAGE_URL = (
    "https://cdn.jsdelivr.net/gh/BerriAI/litellm"
    "@d769e81c90d453240c61fc572cdb27fae06a89d0"
    "/tests/image_gen_tests/test_image.png"
)


MISTRAL_OCR_MODEL: Final = "mistral/mistral-ocr-latest"
AZURE_AI_OCR_MODEL: Final = "azure_ai/mistral-document-ai-2512"
AZURE_DOC_INTELLIGENCE_MODEL: Final = "azure_ai/doc-intelligence/prebuilt-layout"
VERTEX_OCR_MODEL: Final = "vertex_ai/mistral-ocr-2505"
COHERE_OCR_MODEL: Final = "cohere/parse-v5.0"


class OcrProvider(Protocol):
    """One OCR provider's deployment config: its model id plus the os.environ/*
    credential references the proxy resolves at call time. Each provider owns which
    env vars it reads, so a new provider is a new type, not another inline body."""

    def litellm_params(self) -> LiteLLMParamsBody: ...


@dataclass(frozen=True, slots=True)
class MistralOcr:
    model: str = MISTRAL_OCR_MODEL

    def litellm_params(self) -> LiteLLMParamsBody:
        return LiteLLMParamsBody(model=self.model, api_key="os.environ/MISTRAL_API_KEY")


@dataclass(frozen=True, slots=True)
class AzureAiOcr:
    """azure_ai (mistral) OCR. The rust OCR path resolves credentials itself from
    AZURE_AI_API_BASE / AZURE_AI_API_KEY when the deployment leaves them unset; it
    does NOT unwrap an `os.environ/*` reference passed as api_base (it would be sent
    to Azure verbatim), so we omit them and let litellm read the env vars by name."""

    model: str

    def litellm_params(self) -> LiteLLMParamsBody:
        return LiteLLMParamsBody(model=self.model)


@dataclass(frozen=True, slots=True)
class AzureDocIntelligenceOcr:
    """azure_ai Document Intelligence OCR. A separate Azure resource from the
    mistral one, so it has its own env vars: AZURE_DOCUMENT_INTELLIGENCE_ENDPOINT /
    AZURE_DOCUMENT_INTELLIGENCE_API_KEY, which the OCR config resolves from the
    doc-intelligence model name when api_base/api_key are left unset."""

    model: str = AZURE_DOC_INTELLIGENCE_MODEL

    def litellm_params(self) -> LiteLLMParamsBody:
        return LiteLLMParamsBody(model=self.model)


@dataclass(frozen=True, slots=True)
class VertexOcr:
    """Vertex AI OCR (Mistral publisher). Only the location (not a secret) is set;
    the project and credentials are left unset so the gateway resolves VERTEXAI_PROJECT
    and VERTEXAI_CREDENTIALS from its own environment by name, keeping every secret on
    the gateway like the azure_ai cases above. This is deliberate: the OCR path reads
    vertex_project verbatim from litellm_params and never unwraps an `os.environ/*`
    ref, so passing one would put the literal string in the request URL."""

    model: str
    location: str

    def litellm_params(self) -> LiteLLMParamsBody:
        return LiteLLMParamsBody(model=self.model, vertex_location=self.location)


@dataclass(frozen=True, slots=True)
class CohereOcr:
    model: str = COHERE_OCR_MODEL

    def litellm_params(self) -> LiteLLMParamsBody:
        return LiteLLMParamsBody(model=self.model, api_key="os.environ/COHERE_API_KEY")


@dataclass(frozen=True, slots=True)
class _OcrCase:
    suffix: str
    provider: OcrProvider
    document: OcrDocument


RUST_OCR_CASES: tuple[_OcrCase, ...] = (
    _OcrCase(
        "mistral",
        MistralOcr(),
        OcrDocument(type="document_url", document_url=TEST_PDF_URL),
    ),
    _OcrCase(
        "azure-ai",
        AzureAiOcr(AZURE_AI_OCR_MODEL),
        OcrDocument(type="document_url", document_url=TEST_PDF_URL),
    ),
    _OcrCase(
        "azure-document-intelligence",
        AzureDocIntelligenceOcr(),
        OcrDocument(type="document_url", document_url=TEST_PDF_URL),
    ),
    _OcrCase(
        "vertex-mistral",
        VertexOcr(VERTEX_OCR_MODEL, "us-central1"),
        OcrDocument(type="document_url", document_url=TEST_PDF_URL),
    ),
)

PDF_TEXT: Final = "test pdf file"
IMAGE_TEXT: Final = "litellm"
PDF_DOCUMENT: Final = OcrDocument(type="document_url", document_url=TEST_PDF_URL)
IMAGE_DOCUMENT: Final = OcrDocument(type="image_url", image_url=TEST_IMAGE_URL)


@dataclass(frozen=True, slots=True)
class _OcrContentCase:
    suffix: str
    provider: OcrProvider
    document: OcrDocument
    expected_text: str


OCR_CONTENT_CASES: Final = (
    _OcrContentCase("mistral-pdf", MistralOcr(), PDF_DOCUMENT, PDF_TEXT),
    _OcrContentCase("mistral-image", MistralOcr(), IMAGE_DOCUMENT, IMAGE_TEXT),
    _OcrContentCase("azure-ai-image", AzureAiOcr(AZURE_AI_OCR_MODEL), IMAGE_DOCUMENT, IMAGE_TEXT),
    _OcrContentCase(
        "vertex-mistral-image", VertexOcr(VERTEX_OCR_MODEL, "us-central1"), IMAGE_DOCUMENT, IMAGE_TEXT
    ),
    _OcrContentCase("cohere-image", CohereOcr(), IMAGE_DOCUMENT, IMAGE_TEXT),
)


def _ocr_subject(provider: OcrProvider) -> Subject:
    match provider:
        case MistralOcr():
            vendor, model = Provider.MISTRAL, MISTRAL_OCR_MODEL
        case AzureAiOcr():
            vendor, model = Provider.AZURE_AI, AZURE_AI_OCR_MODEL
        case AzureDocIntelligenceOcr():
            vendor, model = Provider.AZURE_AI, AZURE_DOC_INTELLIGENCE_MODEL
        case VertexOcr():
            vendor, model = Provider.VERTEX_AI, VERTEX_OCR_MODEL
        case CohereOcr():
            vendor, model = Provider.COHERE, COHERE_OCR_MODEL
        case _:
            raise TypeError(f"no OCR subject for {provider!r}")
    return Subject(
        domain=Domain.LLM_TRANSLATION,
        route=Route.OCR,
        providers=(vendor,),
        models=(model,),
        mode=Mode.NONSTREAM,
    )


def _assert_ocr_document(response: OcrResponse) -> None:
    assert response.object == "ocr", f"expected object='ocr', got {response.object!r}"
    assert response.model, "response missing the resolved model name"
    assert response.pages, "OCR returned no pages"
    assert response.pages[0].markdown is not None, "first page has no markdown"


def _assert_provider_rate_limit_relayed(model: str, outcome: RateLimitedError) -> None:
    detail: Final = relayed_provider_rate_limit(outcome)
    assert detail is not None, f"{model}: the 429 is the gateway's own, not the provider's: {outcome.body}"
    assert (detail.type, detail.code) == ("throttling_error", "429"), f"{model}: provider 429 relayed as {detail!r}"
    print(
        f"{model}: the provider's rate limit held across {PROVIDER_RATE_LIMIT_ATTEMPTS} attempts; "
        f"the gateway relayed it as {detail.type} {detail.code}",
        flush=True,
    )


class TestRustOcrGateway:
    @pytest.mark.parametrize(
        "case",
        [pytest.param(case, marks=meta(_ocr_subject(case.provider)), id=case.suffix) for case in RUST_OCR_CASES],
    )
    def test_rust_ocr_response(self, proxy: ProxyClient, resources: ResourceManager, case: _OcrCase) -> None:
        model = f"rust-ocr-{case.suffix}-{unique_marker()}"
        model_id = proxy.create_model(model, case.provider.litellm_params())
        resources.defer(lambda: proxy.delete_model(model_id))
        key = resources.key()

        match tolerate_provider_rate_limit(lambda: proxy.ocr(key, OcrBody(model=model, document=case.document))):
            case Success(data=response):
                _assert_ocr_document(response)
            case RateLimitedError() as outcome:
                _assert_provider_rate_limit_relayed(model, outcome)
            case outcome:
                pytest.fail(f"{model}: {outcome!r}")

    @pytest.mark.skip(reason="stage red: product gap, /v1/ocr 500s (aocr TypeError) on missing document instead of 400")
    @pytest.mark.covers("llm.ocr.openai.input_validation.nonstream.works")
    @meta(
        Subject(
            domain=Domain.LLM_TRANSLATION,
            route=Route.OCR,
            providers=(),
            models=(),
        )
    )
    def test_missing_document_returns_error(self, proxy: ProxyClient, resources: ResourceManager) -> None:
        model = f"rust-ocr-val-{unique_marker()}"
        model_id = proxy.create_model(model, MistralOcr().litellm_params())
        resources.defer(lambda: proxy.delete_model(model_id))
        key = resources.key()
        result = proxy.transport.send(
            "/v1/ocr",
            headers=proxy.transport.bearer(key),
            json=_OptionalOcrBody(model=model),
        )
        assert_client_error(result, "ocr missing document")


class TestOcrDocumentContent:
    @pytest.mark.parametrize(
        "case",
        [
            pytest.param(case, marks=meta(_ocr_subject(case.provider)), id=case.suffix)
            for case in OCR_CONTENT_CASES
        ],
    )
    def test_ocr_reads_the_document_and_bills_its_pages(
        self, proxy: ProxyClient, resources: ResourceManager, case: _OcrContentCase
    ) -> None:
        model = f"ocr-content-{case.suffix}-{unique_marker()}"
        model_id = proxy.create_model(model, case.provider.litellm_params())
        resources.defer(lambda: proxy.delete_model(model_id))

        result = proxy.transport.send(
            "/v1/ocr",
            headers=proxy.transport.bearer(resources.key()),
            json=OcrBody(model=model, document=case.document),
        )
        assert result.status_code == 200, f"{model}: /v1/ocr failed with {result.status_code}: {result.body[:300]}"
        response = OcrResponse.model_validate_json(result.body)
        assert response.object == "ocr", f"expected object='ocr', got {response.object!r}"
        assert [page.index for page in response.pages] == list(range(len(response.pages))), (
            f"page indexes are not contiguous from 0: {[page.index for page in response.pages]}"
        )
        text = " ".join(" ".join(page.markdown for page in response.pages).split()).lower()
        assert case.expected_text in text, f"{model}: OCR text lost the document content: {text[:300]!r}"
        assert response.usage_info is not None and response.usage_info.pages_processed == len(response.pages), (
            f"usage_info.pages_processed disagrees with the returned pages: {response.usage_info!r}"
        )
        assert result.response_cost is not None and result.response_cost > 0, (
            f"{model}: OCR call was not costed: x-litellm-response-cost={result.response_cost!r}"
        )