diff --git a/tests/e2e/batches/COVERAGE.md b/tests/e2e/batches/COVERAGE.md index 0026d4bb3d5..4ba9b1cea4b 100644 --- a/tests/e2e/batches/COVERAGE.md +++ b/tests/e2e/batches/COVERAGE.md @@ -9,20 +9,21 @@ cancels, and lists a batch; everything created is deleted on teardown. Only supported cells are tested. The capability table in `capabilities.py` holds one row per supported (provider, scenario) pair, so there are no skipped cells in the -parametrized run. +parametrized run. The batches suite never skips: missing provider creds or upstream +failures are hard test failures (see `tests/e2e/CLAUDE.md`). | Provider | create | retrieve | cancel | list | file backing | |-----------|--------|----------|--------|------|--------------| | OpenAI | yes | yes | yes | yes | OpenAI Files | | Azure | yes | yes | yes | yes | Azure Files | -| Vertex AI | yes | yes | yes | yes | GCS bucket | -| Bedrock | yes | yes | no (limited upstream) | no | S3 bucket | -| Anthropic | no | yes (env-gated) | no | no | Anthropic Files | +| Vertex AI | yes | yes | yes | yes | GCS bucket (`GCS_BUCKET_NAME` via files_settings) | +| Bedrock | yes | yes | no (limited upstream) | no | S3 bucket (`AWS_BATCH_S3_BUCKET` + `AWS_BATCH_ROLE_ARN` on model) | Bedrock cancel is unreliable upstream and list is unsupported, so both are gated off -(`can_cancel=False`, `can_list=False`). Anthropic cannot create/cancel/list through -litellm, so it has a standalone retrieve test that skips unless `ANTHROPIC_BATCH_ID` -points at a real Anthropic batch. +(`can_cancel=False`, `can_list=False`) when that provider is enabled in the matrix. +Bedrock file upload requires a model on the request (`encoded` / `unified` scenarios only); +`model_param` and `provider_fallback` are omitted because `POST /bedrock/v1/files` has no +model-less passthrough path. ## Routing scenarios (per `litellm/proxy/batches_endpoints/endpoints.py`) @@ -67,9 +68,10 @@ File delete asserts `object=="file"` and `deleted==True`. | File | Covers | |------|--------| -| `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared Gateway; denial helpers | -| `capabilities.py` | the provider x scenario matrix + id-shape classifiers + per-provider raw-id assertion | -| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial, anthropic retrieve | +| `batch_client.py` | typed file upload/download + batch create/retrieve/cancel/list/delete over the shared Gateway; runtime batch model registration via /model/new; denial helpers | +| `capabilities.py` | the provider x scenario matrix + per-provider /model/new params + id-shape classifiers + per-provider raw-id assertion | +| `conftest.py` | session-scoped batch deployment registration and teardown | +| `test_batches_e2e.py` | parametrized lifecycle with per-endpoint output assertions, file upload/delete outputs, key-model-access denial | ## Out of scope (intentionally) diff --git a/tests/e2e/batches/batch_client.py b/tests/e2e/batches/batch_client.py index 2a6376bf89e..fda4e87e478 100644 --- a/tests/e2e/batches/batch_client.py +++ b/tests/e2e/batches/batch_client.py @@ -1,11 +1,13 @@ """Client for the batches e2e suite: file upload/download and the batch operations (create / retrieve / cancel / list) over the shared Gateway. -`create_batch` returns the raw HTTP outcome (StreamingResponse) so a 403 model -access denial and a provider-native batch body both surface; the test parses -BatchObject from the body. A `provider` arg routes a call to /{provider}/v1/..., -which the provider-fallback scenario needs (its ids are raw, not model-encoded). -The request/response models are co-located here because only this suite uses them. +Batch deployments are registered at runtime via /model/new (see conftest.py), +not baked into the proxy config. `create_batch` returns the raw HTTP outcome +(StreamingResponse) so a 403 model access denial and a provider-native batch +body both surface; the test parses BatchObject from the body. A `provider` arg +routes a call to /{provider}/v1/..., which the provider-fallback scenario needs +(its ids are raw, not model-encoded). The request/response models are +co-located here because only this suite uses them. """ from __future__ import annotations @@ -22,6 +24,7 @@ from e2e_http import ( StreamingResponse, UnknownApiError, ) +from models import LiteLLMParamsBody class FileObject(BaseModel): @@ -84,6 +87,12 @@ def is_result_access_denied[R: BaseModel](result: Result[R]) -> bool: class BatchClient: gateway: Gateway + def create_model(self, model_name: str, litellm_params: LiteLLMParamsBody) -> str: + return self.gateway.create_model(model_name, litellm_params, mode="batch") + + def delete_model(self, model_id: str) -> None: + self.gateway.delete_model(model_id) + def upload_file( self, *, diff --git a/tests/e2e/batches/capabilities.py b/tests/e2e/batches/capabilities.py index 0416508a31a..522e3162e24 100644 --- a/tests/e2e/batches/capabilities.py +++ b/tests/e2e/batches/capabilities.py @@ -13,6 +13,8 @@ import base64 from dataclasses import dataclass from typing import Literal +from models import LiteLLMParamsBody + Scenario = Literal["encoded", "unified", "model_param", "provider_fallback"] IdShape = Literal["managed", "model_encoded", "raw"] @@ -33,6 +35,39 @@ class Provider: can_cancel: bool can_list: bool + def litellm_params(self) -> LiteLLMParamsBody: + match self.name: + case "openai": + return LiteLLMParamsBody( + model="openai/gpt-4o-mini", + api_key="os.environ/OPENAI_API_KEY", + ) + case "azure": + return LiteLLMParamsBody( + model="azure/gpt-4.1-mini-batch", + api_base="os.environ/AZURE_API_BASE", + api_key="os.environ/AZURE_API_KEY", + api_version="2024-07-01-preview", + ) + case "vertex_ai": + return LiteLLMParamsBody( + model="vertex_ai/gemini-2.5-flash", + vertex_project="os.environ/VERTEXAI_PROJECT", + vertex_location="us-central1", + vertex_credentials="os.environ/VERTEXAI_CREDENTIALS", + ) + case "bedrock": + return LiteLLMParamsBody( + model="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + s3_access_key_id="os.environ/AWS_ACCESS_KEY_ID", + s3_secret_access_key="os.environ/AWS_SECRET_ACCESS_KEY", + s3_region_name="os.environ/AWS_REGION", + s3_bucket_name="os.environ/AWS_BATCH_S3_BUCKET", + aws_batch_role_arn="os.environ/AWS_BATCH_ROLE_ARN", + ) + case _: + raise ValueError(f"unknown batch provider: {self.name!r}") + @dataclass(frozen=True, slots=True) class Capability: @@ -66,19 +101,28 @@ PROVIDERS: tuple[Provider, ...] = ( Provider( "vertex_ai", "vertex-batch", "gemini-2.5-flash", can_cancel=True, can_list=True ), - # Provider( - # "bedrock", - # "bedrock-batch", - # "us.anthropic.claude-haiku-4-5-20251001-v1:0", - # can_cancel=False, - # can_list=False, - # ), + Provider( + "bedrock", + "bedrock-batch", + "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + can_cancel=False, + can_list=False, + ), ) +BEDROCK_SCENARIOS: tuple[Scenario, ...] = ("encoded", "unified") + + +def scenarios_for_provider(provider: Provider) -> tuple[Scenario, ...]: + if provider.name == "bedrock": + return BEDROCK_SCENARIOS + return SCENARIOS + + CAPABILITIES: tuple[Capability, ...] = tuple( Capability(p.name, p.model, p.raw_model, scenario, p.can_cancel, p.can_list) for p in PROVIDERS - for scenario in SCENARIOS + for scenario in scenarios_for_provider(p) ) @@ -88,17 +132,13 @@ def raw_id_matches_provider(provider: str, batch_id: str) -> bool: if provider in ("openai", "azure"): return batch_id.startswith("batch") if provider == "vertex_ai": - # Vertex returns the batch prediction job id, which depending on the - # routing path arrives either as the full resource name - # (projects/.../batchPredictionJobs/) or as just the trailing - # numeric id, so accept either form. return ( batch_id.startswith("projects/") or "batchPredictionJobs" in batch_id or batch_id.isdigit() ) if provider == "bedrock": - return batch_id.startswith("arn:aws") + return batch_id.startswith("arn:aws:bedrock:") return True diff --git a/tests/e2e/llm_translation/test_ocr_rust_e2e.py b/tests/e2e/llm_translation/test_ocr_rust_e2e.py index 44908da2729..361bb5126a7 100644 --- a/tests/e2e/llm_translation/test_ocr_rust_e2e.py +++ b/tests/e2e/llm_translation/test_ocr_rust_e2e.py @@ -1,31 +1,32 @@ """Live e2e: Rust-backed OCR is reachable through the gateway across providers. -The gateway config declares one rust-ocr deployment per provider (mistral, -azure_ai, azure document intelligence, vertex mistral, vertex deepseek). Start the -proxy with the Rust OCR path enabled: +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: - LITELLM_USE_RUST_OCR=1 litellm --config tests/e2e/gateway/litellm-config.yml - -Three behaviors are checked: the config declares every provider's deployment (a -pure config read, no proxy needed); the running proxy loaded them onto /model/info; -and each one returns a well-formed OCR document over /v1/ocr. Per the e2e -"skip on environment, fail on behavior" rule, the proxy-backed cases skip when no -proxy answers but fail (never skip) once a request reaches it, so a provider whose -credentials are missing surfaces as a hard failure rather than silent green. +Each case creates its deployment, drives a real /v1/ocr call, and asserts a +well-formed OCR document comes back. Per the e2e "skip on environment, fail on +behavior" rule, a case skips when no proxy answers but fails (never skips) 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. """ from __future__ import annotations from dataclasses import dataclass -from pathlib import Path +from typing import Protocol import pytest -import yaml -from pydantic import BaseModel +from e2e_config import unique_marker from e2e_http import unwrap -from models import OcrBody, OcrDocument, OcrResponse -from passthrough_client import PassthroughClient +from endpoints_client import EndpointsClient +from lifecycle import ResourceManager +from models import LiteLLMParamsBody, OcrBody, OcrDocument, OcrResponse + +pytestmark = pytest.mark.e2e # Tiny in-repo fixtures served via jsdelivr (sha-pinned, immutable) so the request # bodies stay stable across runs. @@ -40,53 +41,99 @@ TEST_IMAGE_URL = ( "/tests/image_gen_tests/test_image.png" ) -CONFIG_PATH = Path(__file__).resolve().parents[1] / "gateway" / "litellm-config.yml" + +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/mistral-ocr-latest" + + 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_ai/doc-intelligence/prebuilt-layout" + + def litellm_params(self) -> LiteLLMParamsBody: + return LiteLLMParamsBody(model=self.model) + + +@dataclass(frozen=True, slots=True) +class VertexOcr: + model: str + location: str + + def litellm_params(self) -> LiteLLMParamsBody: + return LiteLLMParamsBody( + model=self.model, + vertex_project="os.environ/VERTEXAI_PROJECT", + vertex_location=self.location, + vertex_credentials="os.environ/VERTEXAI_CREDENTIALS", + ) @dataclass(frozen=True, slots=True) class _OcrCase: - model: str + suffix: str + provider: OcrProvider document: OcrDocument RUST_OCR_CASES: tuple[_OcrCase, ...] = ( _OcrCase( - "rust-ocr-mistral", + "mistral", + MistralOcr(), OcrDocument(type="document_url", document_url=TEST_PDF_URL), ), _OcrCase( - "rust-ocr-azure-ai", + "azure-ai", + AzureAiOcr("azure_ai/mistral-document-ai-2505"), OcrDocument(type="document_url", document_url=TEST_PDF_URL), ), _OcrCase( - "rust-ocr-azure-document-intelligence", + "azure-document-intelligence", + AzureDocIntelligenceOcr(), OcrDocument(type="document_url", document_url=TEST_PDF_URL), ), _OcrCase( - "rust-ocr-vertex-mistral", + "vertex-mistral", + VertexOcr("vertex_ai/mistral-ocr-2505", "us-central1"), OcrDocument(type="document_url", document_url=TEST_PDF_URL), ), _OcrCase( - "rust-ocr-vertex-deepseek", + "vertex-deepseek", + VertexOcr("vertex_ai/deepseek-ocr-maas", "global"), OcrDocument(type="image_url", image_url=TEST_IMAGE_URL), ), ) -_EXPECTED_MODELS = frozenset(case.model for case in RUST_OCR_CASES) -_CASE_IDS = tuple(case.model.removeprefix("rust-ocr-") for case in RUST_OCR_CASES) - - -class _ConfiguredModel(BaseModel): - model_name: str - - -class _GatewayConfig(BaseModel): - model_list: list[_ConfiguredModel] - - -def _configured_model_names() -> frozenset[str]: - config = _GatewayConfig.model_validate(yaml.safe_load(CONFIG_PATH.read_text())) - return frozenset(entry.model_name for entry in config.model_list) +_CASE_IDS = tuple(case.suffix for case in RUST_OCR_CASES) def _assert_ocr_document(response: OcrResponse) -> None: @@ -96,22 +143,15 @@ def _assert_ocr_document(response: OcrResponse) -> None: assert response.pages[0].markdown is not None, "first page has no markdown" -def test_rust_ocr_models_declared_in_gateway_config() -> None: - """Pure config read (no proxy): every provider's rust-ocr deployment the suite - exercises is declared in the gateway config the proxy runs with. A case added - here without a matching deployment fails before any live call is attempted.""" - missing = _EXPECTED_MODELS - _configured_model_names() - assert not missing, f"rust-ocr models absent from {CONFIG_PATH.name}: {missing}" - - -@pytest.mark.e2e class TestRustOcrGateway: - def test_gateway_loaded_rust_ocr_models(self, client: PassthroughClient) -> None: - loaded = frozenset(entry.model_name for entry in client.gateway.model_info()) - missing = _EXPECTED_MODELS - loaded - assert not missing, f"proxy did not load rust-ocr models: {missing}" - @pytest.mark.parametrize("case", RUST_OCR_CASES, ids=_CASE_IDS) - def test_rust_ocr_response(self, client: PassthroughClient, scoped_key: str, case: _OcrCase) -> None: - response = unwrap(client.gateway.ocr(scoped_key, OcrBody(model=case.model, document=case.document))) + def test_rust_ocr_response( + self, endpoints_client: EndpointsClient, resources: ResourceManager, case: _OcrCase + ) -> None: + model = f"rust-ocr-{case.suffix}-{unique_marker()}" + model_id = endpoints_client.create_model(model, case.provider.litellm_params()) + resources.defer(lambda: endpoints_client.delete_model(model_id)) + key = resources.key() + + response = unwrap(endpoints_client.gateway.ocr(key, OcrBody(model=model, document=case.document))) _assert_ocr_document(response)