litellm/tests/unit/batches/test_main.py
mika-lt d1d6d55063
fix(hosted_vllm): accept the vLLM production-stack router's file and batch objects (#49068)
* fix(files): default a missing file object type to "file" so OpenAI-compatible servers that omit it stop 500ing

* fix(batches): accept the vLLM production-stack router's batch object

The router leaves out the batch object's "object" field and reports "pending" and "running" statuses, so batch create and retrieve through a hosted_vllm deployment failed validation with a 500. A missing or null object now reads as "batch", pending maps to validating and running to in_progress, and any other unknown value is still rejected

* fix(batches): return OpenAI batch objects from batch list behind the vLLM router

* fix(batches): type the batch list conversion without a TypeVar and refuse an async client on the sync path

* fix(batches): keep batch list items as OpenAI's Batch so an odd item never fails a list

* refactor(types): reuse the shared JSON value alias on the batch type

* fix: return the provider page for an async client in sync batch list and pass through non-string list statuses

* test: inject the SDK's own transport in the sync batch list test so it passes on OpenAI SDK 3.x

---------

Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
2026-10-30 23:11:34 -07:00

1388 lines
51 KiB
Python

"""
Provider-dispatch contract tests for litellm/batches/main.py
main.py is the SDK layer beneath the proxy batch endpoints: each of
create/retrieve/list/cancel_batch is a switch on `custom_llm_provider` (and, for
create/retrieve, on whether a provider-config + model is present) that hands off
to exactly one provider handler. These tests lock that dispatch:
1. DISPATCH - exactly which provider seam fired (openai_batches_instance vs
azure vs vertex vs anthropic vs base_llm_http_handler vs the
Bedrock ARN handlers), with every sibling seam asserted NOT
called. A reordered/negated branch flips this.
2. PAYLOAD - the request object (CreateBatchRequest/RetrieveBatchRequest/...)
and the _is_async flag forwarded to the handler.
3. RESULT - the handler's return value is what the function returns.
4. DELEGATION - the async wrappers (a*_batch) forward to the sync function in an
executor with the right "_is_async" flag, and pass the result
back untouched.
Only the provider handler instances are mocked (true network boundaries). The
real public functions run (including the @client decorator) so dispatch reflects
production. Provider env vars are not required: missing creds resolve to None and
flow through harmlessly because the handler is mocked.
"""
from contextlib import ExitStack
from dataclasses import dataclass
from typing import Any, Dict
from types import MappingProxyType
from unittest.mock import MagicMock, patch
import pytest
import litellm
import litellm.batches.main as bm
import asyncio
import datetime
import json
from collections.abc import Awaitable, Callable, Mapping, Sequence
from typing import Final
import httpx
import respx
from pydantic import TypeAdapter, ValidationError
from typing_extensions import ReadOnly, TypedDict
from litellm.integrations.custom_logger import CustomLogger
from openai.types.batch import Batch
from litellm.types.utils import LiteLLMBatch
# --------------------------------------------------------------------------- #
# Seam harness - one mock per provider handler instance + the Bedrock ARN
# handler. Each handler method auto-returns a unique sentinel (its
# return_value), so "result is seam.<method>.return_value" verifies dispatch.
# --------------------------------------------------------------------------- #
@dataclass
class Seams:
openai: MagicMock
azure: MagicMock
vertex: MagicMock
anthropic: MagicMock
base_http: MagicMock
bedrock_arn: MagicMock
@pytest.fixture
def seams():
openai_i = MagicMock(name="openai_batches_instance")
azure_i = MagicMock(name="azure_batches_instance")
vertex_i = MagicMock(name="vertex_ai_batches_instance")
anthropic_i = MagicMock(name="anthropic_batches_instance")
base_http = MagicMock(name="base_llm_http_handler")
bedrock_arn = MagicMock(name="BedrockBatchesHandler")
with ExitStack() as stack:
stack.enter_context(patch.object(bm, "openai_batches_instance", openai_i))
stack.enter_context(patch.object(bm, "azure_batches_instance", azure_i))
stack.enter_context(patch.object(bm, "vertex_ai_batches_instance", vertex_i))
stack.enter_context(patch.object(bm, "anthropic_batches_instance", anthropic_i))
stack.enter_context(patch.object(bm, "base_llm_http_handler", base_http))
stack.enter_context(patch.object(bm, "BedrockBatchesHandler", bedrock_arn))
yield Seams(
openai=openai_i,
azure=azure_i,
vertex=vertex_i,
anthropic=anthropic_i,
base_http=base_http,
bedrock_arn=bedrock_arn,
)
# Every <op> handler method across all provider instances - used to assert
# "no sibling seam fired" exhaustively.
def _all_seam_methods(seams: Seams, op: str):
return [
getattr(seams.openai, op),
getattr(seams.azure, op),
getattr(seams.vertex, op),
getattr(seams.anthropic, op),
getattr(seams.base_http, op),
]
def _assert_only(fired, seams: Seams, op: str):
"""Assert `fired` was called exactly once and every other op seam was not."""
assert fired.call_count == 1
for m in _all_seam_methods(seams, op):
if m is not fired:
m.assert_not_called()
CREATE_KW: Dict[str, Any] = dict(
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-abc",
)
# =========================================================================== #
# create_batch
# =========================================================================== #
def test_create__openai_dispatch_and_payload(seams):
result = bm.create_batch(**CREATE_KW, custom_llm_provider="openai")
# DISPATCH + RESULT
assert result is seams.openai.create_batch.return_value
_assert_only(seams.openai.create_batch, seams, "create_batch")
seams.bedrock_arn.handle_async_invoke_status.assert_not_called()
# PAYLOAD - request object built from the call, sync flag off.
kw = seams.openai.create_batch.call_args.kwargs
assert kw["create_batch_data"] == {
"completion_window": "24h",
"endpoint": "/v1/chat/completions",
"input_file_id": "file-abc",
"metadata": None,
"extra_headers": None,
"extra_body": None,
}
assert kw["_is_async"] is False
assert kw["timeout"] == 600.0
def test_create__hosted_vllm_routes_to_openai_instance(seams):
"""hosted_vllm is in OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS, so it shares
the openai handler. Locks that set membership."""
result = bm.create_batch(**CREATE_KW, custom_llm_provider="hosted_vllm")
assert result is seams.openai.create_batch.return_value
_assert_only(seams.openai.create_batch, seams, "create_batch")
def test_create__azure_dispatch(seams):
result = bm.create_batch(**CREATE_KW, custom_llm_provider="azure")
assert result is seams.azure.create_batch.return_value
_assert_only(seams.azure.create_batch, seams, "create_batch")
def test_create__vertex_ai_dispatch(seams):
result = bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai")
assert result is seams.vertex.create_batch.return_value
_assert_only(seams.vertex.create_batch, seams, "create_batch")
def test_create__vertex_ai_forwards_custom_endpoint(seams):
"""The vertex handler owns the custom_endpoint batch rejection (LIT-6899), so the dispatcher
must forward the flag for the handler to act on."""
bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai", custom_endpoint=True)
assert seams.vertex.create_batch.call_args.kwargs["custom_endpoint"] is True
def test_create__provider_config_routes_to_base_http_handler(seams):
"""model + a provider batches config (bedrock-style) routes to the generic
base_llm_http_handler, NOT the per-provider instance."""
with patch.object(
bm.ProviderConfigManager,
"get_provider_batches_config",
return_value=MagicMock(name="provider_config"),
):
result = bm.create_batch(**CREATE_KW, custom_llm_provider="bedrock", model="bedrock/my-batch-model")
assert result is seams.base_http.create_batch.return_value
_assert_only(seams.base_http.create_batch, seams, "create_batch")
def test_create__unsupported_provider_raises_badrequest(seams):
with pytest.raises(litellm.exceptions.BadRequestError):
bm.create_batch(**CREATE_KW, custom_llm_provider="cohere") # type: ignore[arg-type]
for m in _all_seam_methods(seams, "create_batch"):
m.assert_not_called()
@pytest.mark.asyncio
async def test_create__async_path_propagates_is_async(seams):
"""Through the real async wrapper, the handler is invoked with _is_async=True.
(Calling the @client sync create_batch with acreate_batch=True directly is not
a real code path - logging-obj setup only happens on the async wrapper path.)"""
await bm.acreate_batch(**CREATE_KW, custom_llm_provider="openai")
assert seams.openai.create_batch.call_args.kwargs["_is_async"] is True
# =========================================================================== #
# retrieve_batch
# =========================================================================== #
def test_retrieve__openai_dispatch_and_payload(seams):
result = bm.retrieve_batch(batch_id="batch-1", custom_llm_provider="openai")
assert result is seams.openai.retrieve_batch.return_value
_assert_only(seams.openai.retrieve_batch, seams, "retrieve_batch")
kw = seams.openai.retrieve_batch.call_args.kwargs
assert kw["retrieve_batch_data"] == {
"batch_id": "batch-1",
"extra_headers": None,
"extra_body": None,
}
assert kw["_is_async"] is False
def test_retrieve__hosted_vllm_routes_to_openai_instance(seams):
result = bm.retrieve_batch(batch_id="batch-1", custom_llm_provider="hosted_vllm")
assert result is seams.openai.retrieve_batch.return_value
_assert_only(seams.openai.retrieve_batch, seams, "retrieve_batch")
def test_retrieve__azure_dispatch(seams):
result = bm.retrieve_batch(batch_id="batch-1", custom_llm_provider="azure")
assert result is seams.azure.retrieve_batch.return_value
_assert_only(seams.azure.retrieve_batch, seams, "retrieve_batch")
def test_retrieve__vertex_ai_dispatch(seams):
result = bm.retrieve_batch(batch_id="batch-1", custom_llm_provider="vertex_ai")
assert result is seams.vertex.retrieve_batch.return_value
_assert_only(seams.vertex.retrieve_batch, seams, "retrieve_batch")
def test_retrieve__anthropic_dispatch(seams):
"""anthropic is retrieve-capable (not in create's provider set)."""
result = bm.retrieve_batch(batch_id="batch-1", custom_llm_provider="anthropic")
assert result is seams.anthropic.retrieve_batch.return_value
_assert_only(seams.anthropic.retrieve_batch, seams, "retrieve_batch")
def test_retrieve__provider_config_routes_to_base_http_handler(seams):
with patch.object(
bm.ProviderConfigManager,
"get_provider_batches_config",
return_value=MagicMock(name="provider_config"),
):
result = bm.retrieve_batch(
batch_id="batch-1",
custom_llm_provider="bedrock",
model="bedrock/my-batch-model",
)
assert result is seams.base_http.retrieve_batch.return_value
_assert_only(seams.base_http.retrieve_batch, seams, "retrieve_batch")
def test_retrieve__bedrock_async_invoke_arn(seams):
arn = "arn:aws:bedrock:us-east-1:123456789012:async-invoke/abc123"
result = bm.retrieve_batch(batch_id=arn, custom_llm_provider="bedrock")
seams.bedrock_arn.handle_async_invoke_status.assert_called_once()
assert result is seams.bedrock_arn.handle_async_invoke_status.return_value
# provider instances untouched.
for m in _all_seam_methods(seams, "retrieve_batch"):
m.assert_not_called()
def test_retrieve__bedrock_model_invocation_job_arn(seams):
arn = "arn:aws:bedrock:us-east-1:123456789012:model-invocation-job/xyz789"
result = bm.retrieve_batch(batch_id=arn, custom_llm_provider="bedrock")
seams.bedrock_arn.handle_model_invocation_job_status.assert_called_once()
assert result is seams.bedrock_arn.handle_model_invocation_job_status.return_value
seams.bedrock_arn.handle_async_invoke_status.assert_not_called()
def test_retrieve__unsupported_provider_raises_badrequest(seams):
with pytest.raises(litellm.exceptions.BadRequestError):
bm.retrieve_batch(batch_id="batch-1", custom_llm_provider="cohere") # type: ignore[arg-type]
for m in _all_seam_methods(seams, "retrieve_batch"):
m.assert_not_called()
# =========================================================================== #
# list_batches (supported: openai, hosted_vllm, azure, vertex_ai)
# =========================================================================== #
def test_list__openai_dispatch_and_payload(seams):
result = bm.list_batches(custom_llm_provider="openai", after="cur", limit=5)
assert result is seams.openai.list_batches.return_value
_assert_only(seams.openai.list_batches, seams, "list_batches")
kw = seams.openai.list_batches.call_args.kwargs
assert kw["after"] == "cur"
assert kw["limit"] == 5
assert kw["_is_async"] is False
def test_list__hosted_vllm_routes_to_openai_instance(seams):
result = bm.list_batches(custom_llm_provider="hosted_vllm")
assert result is seams.openai.list_batches.return_value
_assert_only(seams.openai.list_batches, seams, "list_batches")
def test_list__azure_dispatch(seams):
result = bm.list_batches(custom_llm_provider="azure")
assert result is seams.azure.list_batches.return_value
_assert_only(seams.azure.list_batches, seams, "list_batches")
def test_list__vertex_ai_dispatch(seams):
result = bm.list_batches(custom_llm_provider="vertex_ai")
assert result is seams.vertex.list_batches.return_value
_assert_only(seams.vertex.list_batches, seams, "list_batches")
def test_list__unsupported_provider_raises_badrequest(seams):
# anthropic supports retrieve but NOT list - good negative case.
with pytest.raises(litellm.exceptions.BadRequestError):
bm.list_batches(custom_llm_provider="anthropic") # type: ignore[arg-type]
for m in _all_seam_methods(seams, "list_batches"):
m.assert_not_called()
# =========================================================================== #
# cancel_batch (supported: openai, hosted_vllm, azure, vertex_ai; no @client)
# =========================================================================== #
def test_cancel__openai_dispatch_and_payload(seams):
result = bm.cancel_batch(batch_id="batch-1", custom_llm_provider="openai")
assert result is seams.openai.cancel_batch.return_value
_assert_only(seams.openai.cancel_batch, seams, "cancel_batch")
kw = seams.openai.cancel_batch.call_args.kwargs
assert kw["cancel_batch_data"] == {
"batch_id": "batch-1",
"extra_headers": None,
"extra_body": None,
}
assert kw["_is_async"] is False
def test_cancel__azure_dispatch(seams):
result = bm.cancel_batch(batch_id="batch-1", custom_llm_provider="azure")
assert result is seams.azure.cancel_batch.return_value
_assert_only(seams.azure.cancel_batch, seams, "cancel_batch")
def test_cancel__vertex_ai_dispatch(seams):
result = bm.cancel_batch(batch_id="batch-1", custom_llm_provider="vertex_ai")
assert result is seams.vertex.cancel_batch.return_value
_assert_only(seams.vertex.cancel_batch, seams, "cancel_batch")
def test_cancel__unsupported_provider_raises_badrequest(seams):
with pytest.raises(litellm.exceptions.BadRequestError):
bm.cancel_batch(batch_id="batch-1", custom_llm_provider="cohere")
for m in _all_seam_methods(seams, "cancel_batch"):
m.assert_not_called()
def test_cancel__async_flag_propagates_is_async(seams):
bm.cancel_batch(batch_id="batch-1", custom_llm_provider="openai", acancel_batch=True)
assert seams.openai.cancel_batch.call_args.kwargs["_is_async"] is True
# =========================================================================== #
# Async wrappers - delegate to the sync function in an executor, set the right
# "_is_async" flag, and return the result untouched.
# =========================================================================== #
@pytest.mark.asyncio
async def test_acreate_batch_delegates_to_create_batch():
with patch.object(bm, "create_batch", MagicMock(return_value="SENTINEL")) as m:
result = await bm.acreate_batch(**CREATE_KW, custom_llm_provider="openai")
assert result == "SENTINEL"
assert m.call_count == 1
assert m.call_args.kwargs.get("acreate_batch") is True
# positional handoff: (completion_window, endpoint, input_file_id, provider, ...)
assert m.call_args.args[0] == "24h"
assert m.call_args.args[2] == "file-abc"
assert m.call_args.args[3] == "openai"
@pytest.mark.asyncio
async def test_aretrieve_batch_delegates_to_retrieve_batch():
with patch.object(bm, "retrieve_batch", MagicMock(return_value="SENTINEL")) as m:
result = await bm.aretrieve_batch(batch_id="batch-1", custom_llm_provider="azure")
assert result == "SENTINEL"
assert m.call_count == 1
assert m.call_args.kwargs.get("aretrieve_batch") is True
assert m.call_args.args[0] == "batch-1"
assert m.call_args.args[1] == "azure"
@pytest.mark.asyncio
async def test_alist_batches_delegates_to_list_batches():
with patch.object(bm, "list_batches", MagicMock(return_value="SENTINEL")) as m:
result = await bm.alist_batches(after="cur", limit=3, custom_llm_provider="vertex_ai")
assert result == "SENTINEL"
assert m.call_count == 1
assert m.call_args.kwargs.get("alist_batches") is True
assert m.call_args.args[0] == "cur"
assert m.call_args.args[1] == 3
assert m.call_args.args[2] == "vertex_ai"
@pytest.mark.asyncio
async def test_acancel_batch_delegates_to_cancel_batch():
with patch.object(bm, "cancel_batch", MagicMock(return_value="SENTINEL")) as m:
result = await bm.acancel_batch(batch_id="batch-1", custom_llm_provider="openai")
assert result == "SENTINEL"
assert m.call_count == 1
assert m.call_args.kwargs.get("acancel_batch") is True
assert m.call_args.args[0] == "batch-1"
# =========================================================================== #
# Credential passthrough - when the caller supplies credentials in kwargs, they
# must reach the provider handler. Explicit kwargs win over litellm.* globals and
# env vars (they are first in each `optional_params.x or litellm.x or env` chain),
# so these assertions are deterministic regardless of the test environment.
#
# The credential-resolution blocks are copy-pasted per provider in EACH of
# create/retrieve/list/cancel, so a regression can land in any one independently;
# every function is checked.
# =========================================================================== #
# Distinct values so a cross-wired field (e.g. api_key forwarded as api_base) is
# impossible to miss.
OPENAI_CREDS: Dict[str, Any] = dict(
api_key="sk-user-openai",
api_base="https://openai.user.test",
organization="org-user-123",
max_retries=7,
)
AZURE_CREDS: Dict[str, Any] = dict(
api_key="sk-user-azure",
api_base="https://azure.user.test",
api_version="2024-12-99",
)
VERTEX_CREDS: Dict[str, Any] = dict(
vertex_project="proj-user",
vertex_location="loc-user",
vertex_credentials="cred-user",
api_base="https://vertex.user.test",
)
def _sent(mock_method, *keys):
"""Subset of the call kwargs limited to `keys`, for exact comparison."""
kw = mock_method.call_args.kwargs
return {k: kw.get(k) for k in keys}
# ---- create_batch ---------------------------------------------------------- #
def test_create__openai_credentials_passthrough(seams):
bm.create_batch(**CREATE_KW, custom_llm_provider="openai", **OPENAI_CREDS)
assert _sent(seams.openai.create_batch, "api_key", "api_base", "organization", "max_retries") == {
"api_key": "sk-user-openai",
"api_base": "https://openai.user.test",
"organization": "org-user-123",
"max_retries": 7,
}
def test_create__azure_credentials_passthrough(seams):
bm.create_batch(**CREATE_KW, custom_llm_provider="azure", **AZURE_CREDS)
assert _sent(seams.azure.create_batch, "api_key", "api_base", "api_version") == {
"api_key": "sk-user-azure",
"api_base": "https://azure.user.test",
"api_version": "2024-12-99",
}
def test_create__vertex_credentials_passthrough(seams):
bm.create_batch(**CREATE_KW, custom_llm_provider="vertex_ai", **VERTEX_CREDS)
assert _sent(
seams.vertex.create_batch,
"vertex_project",
"vertex_location",
"vertex_credentials",
"api_base",
) == {
"vertex_project": "proj-user",
"vertex_location": "loc-user",
"vertex_credentials": "cred-user",
"api_base": "https://vertex.user.test",
}
def test_create__provider_config_credentials_passthrough(seams):
with patch.object(
bm.ProviderConfigManager,
"get_provider_batches_config",
return_value=MagicMock(name="provider_config"),
):
bm.create_batch(
**CREATE_KW,
custom_llm_provider="bedrock",
model="bedrock/my-batch-model",
api_key="sk-user-bedrock",
api_base="https://bedrock.user.test",
)
assert _sent(seams.base_http.create_batch, "api_key", "api_base") == {
"api_key": "sk-user-bedrock",
"api_base": "https://bedrock.user.test",
}
# ---- retrieve_batch -------------------------------------------------------- #
def test_retrieve__openai_credentials_passthrough(seams):
bm.retrieve_batch(batch_id="b1", custom_llm_provider="openai", **OPENAI_CREDS)
assert _sent(seams.openai.retrieve_batch, "api_key", "api_base", "organization") == {
"api_key": "sk-user-openai",
"api_base": "https://openai.user.test",
"organization": "org-user-123",
}
def test_retrieve__azure_credentials_passthrough(seams):
bm.retrieve_batch(batch_id="b1", custom_llm_provider="azure", **AZURE_CREDS)
assert _sent(seams.azure.retrieve_batch, "api_key", "api_base", "api_version") == {
"api_key": "sk-user-azure",
"api_base": "https://azure.user.test",
"api_version": "2024-12-99",
}
def test_retrieve__vertex_credentials_passthrough(seams):
bm.retrieve_batch(batch_id="b1", custom_llm_provider="vertex_ai", **VERTEX_CREDS)
assert _sent(
seams.vertex.retrieve_batch,
"vertex_project",
"vertex_location",
"vertex_credentials",
) == {
"vertex_project": "proj-user",
"vertex_location": "loc-user",
"vertex_credentials": "cred-user",
}
def test_retrieve__anthropic_credentials_passthrough(seams):
bm.retrieve_batch(
batch_id="b1",
custom_llm_provider="anthropic",
api_key="sk-user-anthropic",
api_base="https://anthropic.user.test",
)
assert _sent(seams.anthropic.retrieve_batch, "api_key", "api_base") == {
"api_key": "sk-user-anthropic",
"api_base": "https://anthropic.user.test",
}
def test_retrieve__provider_config_credentials_passthrough(seams):
with patch.object(
bm.ProviderConfigManager,
"get_provider_batches_config",
return_value=MagicMock(name="provider_config"),
):
bm.retrieve_batch(
batch_id="b1",
custom_llm_provider="bedrock",
model="bedrock/my-batch-model",
api_key="sk-user-bedrock",
api_base="https://bedrock.user.test",
)
assert _sent(seams.base_http.retrieve_batch, "api_key", "api_base") == {
"api_key": "sk-user-bedrock",
"api_base": "https://bedrock.user.test",
}
# ---- list_batches ---------------------------------------------------------- #
def test_list__openai_credentials_passthrough(seams):
bm.list_batches(custom_llm_provider="openai", **OPENAI_CREDS)
assert _sent(seams.openai.list_batches, "api_key", "api_base", "organization") == {
"api_key": "sk-user-openai",
"api_base": "https://openai.user.test",
"organization": "org-user-123",
}
def test_list__azure_credentials_passthrough(seams):
bm.list_batches(custom_llm_provider="azure", **AZURE_CREDS)
assert _sent(seams.azure.list_batches, "api_key", "api_base", "api_version") == {
"api_key": "sk-user-azure",
"api_base": "https://azure.user.test",
"api_version": "2024-12-99",
}
def test_list__vertex_credentials_passthrough(seams):
bm.list_batches(custom_llm_provider="vertex_ai", **VERTEX_CREDS)
assert _sent(
seams.vertex.list_batches,
"vertex_project",
"vertex_location",
"vertex_credentials",
) == {
"vertex_project": "proj-user",
"vertex_location": "loc-user",
"vertex_credentials": "cred-user",
}
# ---- cancel_batch ---------------------------------------------------------- #
def test_cancel__openai_credentials_passthrough(seams):
bm.cancel_batch(batch_id="b1", custom_llm_provider="openai", **OPENAI_CREDS)
assert _sent(seams.openai.cancel_batch, "api_key", "api_base", "organization") == {
"api_key": "sk-user-openai",
"api_base": "https://openai.user.test",
"organization": "org-user-123",
}
def test_cancel__azure_credentials_passthrough(seams):
bm.cancel_batch(batch_id="b1", custom_llm_provider="azure", **AZURE_CREDS)
assert _sent(seams.azure.cancel_batch, "api_key", "api_base", "api_version") == {
"api_key": "sk-user-azure",
"api_base": "https://azure.user.test",
"api_version": "2024-12-99",
}
def test_cancel__vertex_credentials_passthrough(seams):
bm.cancel_batch(batch_id="b1", custom_llm_provider="vertex_ai", **VERTEX_CREDS)
assert _sent(
seams.vertex.cancel_batch,
"vertex_project",
"vertex_location",
"vertex_credentials",
) == {
"vertex_project": "proj-user",
"vertex_location": "loc-user",
"vertex_credentials": "cred-user",
}
# =========================================================================== #
# _resolve_timeout - pure helper (used by create_batch).
# =========================================================================== #
def _params(**kw):
from litellm.types.router import GenericLiteLLMParams
return GenericLiteLLMParams(**kw)
def test_resolve_timeout__explicit_numeric():
assert bm._resolve_timeout(_params(timeout=30), {}, "openai") == 30.0
def test_resolve_timeout__default_when_unset():
assert bm._resolve_timeout(_params(), {}, "openai") == 600.0
def test_resolve_timeout__request_timeout_kwarg_fallback():
assert bm._resolve_timeout(_params(), {"request_timeout": 45}, "openai") == 45.0
def test_resolve_timeout__httpx_timeout_returns_float_read():
import httpx
t = httpx.Timeout(99.0, connect=5.0)
resolved = bm._resolve_timeout(_params(timeout=t), {}, "openai")
assert isinstance(resolved, float)
assert resolved == 99.0
def test_retrieve__forwards_trusted_model_credentials_into_litellm_params(seams):
"""The batch's cost is computed by reading its output file after the retrieve, and
Bedrock resolves that bucket only from this immutable snapshot. get_litellm_params has
a fixed signature that drops it, so without re-adding it here the snapshot never
reaches the logging object and cost accounting fails on a bucket that is configured."""
snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket"})
logging_obj = MagicMock()
bm.retrieve_batch(
batch_id="batch-1",
custom_llm_provider="openai",
litellm_logging_obj=logging_obj,
_litellm_internal_model_credentials=snapshot,
)
litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"]
assert litellm_params["_litellm_internal_model_credentials"] is snapshot
def test_retrieve__omits_trusted_model_credentials_when_not_supplied(seams):
"""A retrieve with no snapshot must not invent an empty one, which would read as a
configured bucket of nothing."""
logging_obj = MagicMock()
bm.retrieve_batch(batch_id="batch-1", custom_llm_provider="openai", litellm_logging_obj=logging_obj)
litellm_params = logging_obj.update_from_kwargs.call_args.kwargs["litellm_params"]
assert "_litellm_internal_model_credentials" not in litellm_params
# =========================================================================== #
# mistral - a provider-config provider, like bedrock, so it requires `model`
# =========================================================================== #
def test_create__mistral_ocr_routes_to_base_http_handler_with_mistral_config(seams):
result = bm.create_batch(
completion_window="24h",
endpoint="/v1/ocr",
input_file_id="file-abc",
custom_llm_provider="mistral",
model="mistral/mistral-ocr-latest",
)
assert result is seams.base_http.create_batch.return_value
_assert_only(seams.base_http.create_batch, seams, "create_batch")
forwarded = seams.base_http.create_batch.call_args.kwargs
assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig"
assert forwarded["model"] == "mistral-ocr-latest"
assert forwarded["create_batch_data"]["endpoint"] == "/v1/ocr"
def test_create__mistral_without_model_raises_badrequest(seams):
with pytest.raises(litellm.exceptions.BadRequestError):
bm.create_batch(**CREATE_KW, custom_llm_provider="mistral")
for m in _all_seam_methods(seams, "create_batch"):
m.assert_not_called()
def test_retrieve__mistral_routes_to_base_http_handler_with_mistral_config(seams):
result = bm.retrieve_batch(batch_id="job-1", custom_llm_provider="mistral", model="mistral/mistral-ocr-latest")
assert result is seams.base_http.retrieve_batch.return_value
_assert_only(seams.base_http.retrieve_batch, seams, "retrieve_batch")
forwarded = seams.base_http.retrieve_batch.call_args.kwargs
assert type(forwarded["provider_config"]).__name__ == "MistralBatchesConfig"
assert forwarded["batch_id"] == "job-1"
@pytest.mark.asyncio()
async def test_batch_logging_azure_credentials_regression():
"""
Regression test: LoggingWorker Missing Azure Credentials When Fetching Batch Output
This test ensures that Azure credentials are properly passed when fetching batch
output files during logging, preventing "Missing credentials" errors.
Bug: The LoggingWorker failed when processing completed Azure batches because
it attempted to fetch batch output file content without Azure credentials.
Fix: Pass litellm_params (containing credentials) from the logging object
through to the file content retrieval functions.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from litellm.batches.batch_utils import (
extract_file_access_credentials,
_fetch_batch_output_file_content,
handle_completed_batch,
)
from litellm.types.llms.openai import Batch, HttpxBinaryResponseContent
import httpx
print("\n=== Regression Test: Azure Batch Logging Credentials ===")
mock_batch = Batch(
id="batch-azure-test",
object="batch",
endpoint="/v1/chat/completions",
errors=None,
input_file_id="file-input-azure",
completion_window="24h",
status="completed",
output_file_id="file-output-azure",
error_file_id=None,
created_at=1234567890,
in_progress_at=1234567900,
expires_at=1234654290,
finalizing_at=1234568000,
completed_at=1234568100,
failed_at=None,
expired_at=None,
cancelling_at=None,
cancelled_at=None,
request_counts=None,
metadata=None,
)
azure_credentials = {
"api_key": "test-azure-key-regression",
"api_base": "https://test-regression.openai.azure.com",
"api_version": "2024-02-15-preview",
"organization": "test-org",
"timeout": 600,
}
batch_output = b'{"id": "batch_req_1", "custom_id": "request-1", "response": {"status_code": 200, "body": {"id": "chatcmpl-azure", "object": "chat.completion", "model": "gpt-4", "usage": {"prompt_tokens": 15, "completion_tokens": 25, "total_tokens": 40}}}}\n'
print("\n1. Testing credential extraction...")
extracted_creds = extract_file_access_credentials(azure_credentials)
assert "api_key" in extracted_creds, "api_key should be extracted"
assert (
extracted_creds["api_key"] == "test-azure-key-regression"
), "Incorrect api_key"
assert "api_base" in extracted_creds, "api_base should be extracted"
assert "api_version" in extracted_creds, "api_version should be extracted"
assert "timeout" in extracted_creds, "timeout should be extracted"
print(" ✓ Credentials extracted correctly")
print(f" ✓ Extracted keys: {list(extracted_creds.keys())}")
print("\n2. Testing credentials passed to afile_content...")
credentials_received = {"value": False, "params": None}
async def mock_afile_content_tracker(**kwargs):
if "api_key" in kwargs and "api_base" in kwargs and "api_version" in kwargs:
credentials_received["value"] = True
credentials_received["params"] = {
"api_key": kwargs.get("api_key"),
"api_base": kwargs.get("api_base"),
"api_version": kwargs.get("api_version"),
}
mock_response = httpx.Response(
status_code=200,
content=batch_output,
headers={"content-type": "application/octet-stream"},
)
return HttpxBinaryResponseContent(response=mock_response)
with patch(
"litellm.files.main.afile_content", side_effect=mock_afile_content_tracker
):
result = await _fetch_batch_output_file_content(
batch=mock_batch,
custom_llm_provider="azure",
litellm_params=azure_credentials,
)
assert credentials_received[
"value"
], "REGRESSION: Azure credentials not passed to afile_content! This causes 'Missing credentials' error."
assert (
credentials_received["params"]["api_key"] == "test-azure-key-regression"
), "REGRESSION: Incorrect api_key"
assert (
credentials_received["params"]["api_base"]
== "https://test-regression.openai.azure.com"
), "REGRESSION: Incorrect api_base"
print(" ✓ Credentials passed to afile_content")
print(f" ✓ api_key: {credentials_received['params']['api_key']}")
print(f" ✓ api_base: {credentials_received['params']['api_base']}")
print("\n3. Testing full logging flow...")
credentials_received["value"] = False
credentials_received["params"] = None
with patch(
"litellm.files.main.afile_content", side_effect=mock_afile_content_tracker
):
result = await handle_completed_batch(
batch=mock_batch,
custom_llm_provider="azure",
litellm_params=azure_credentials,
)
assert credentials_received[
"value"
], "REGRESSION: Credentials not passed through _handle_completed_batch"
assert result.cost > 0, "Cost should be calculated"
assert result.usage.total_tokens == 40, "Usage should be calculated correctly"
print(" ✓ Credentials passed through full flow")
print(f" ✓ Cost: {result.cost}")
print(f" ✓ Usage: {result.usage.total_tokens} tokens")
print(f" ✓ Models: {result.models}")
print("\n4. Testing 'Missing credentials' error prevention...")
with patch("litellm.files.main.afile_content") as mock_afile_content_fail:
mock_afile_content_fail.side_effect = Exception(
"Missing credentials. Please pass one of `api_key`, `azure_ad_token`, "
"`azure_ad_token_provider`, or the `AZURE_OPENAI_API_KEY` or "
"`AZURE_OPENAI_AD_TOKEN` environment variables."
)
with patch(
"litellm.files.main.afile_content", side_effect=mock_afile_content_tracker
):
try:
result = await handle_completed_batch(
batch=mock_batch,
custom_llm_provider="azure",
litellm_params=azure_credentials,
)
print(" ✓ No 'Missing credentials' error with fix")
except Exception as e:
if "Missing credentials" in str(e):
pytest.fail(
f"REGRESSION: 'Missing credentials' error occurred! "
f"Credentials not being passed. Error: {str(e)}"
)
raise
print("\n5. Testing backwards compatibility...")
with patch("litellm.files.main.afile_content") as mock_afile_content:
mock_response = httpx.Response(
status_code=200,
content=batch_output,
headers={"content-type": "application/octet-stream"},
)
mock_afile_content.return_value = HttpxBinaryResponseContent(
response=mock_response
)
result = await _fetch_batch_output_file_content(
batch=mock_batch,
custom_llm_provider="openai",
litellm_params=None,
)
assert len(result) > 0, "Should return file content"
print(" ✓ Backwards compatibility maintained")
print(" ✓ Works without litellm_params for OpenAI")
print("\n=== Regression Test Passed ===")
print("✓ Azure credentials properly passed from logging to file retrieval")
print("✓ 'Missing credentials' error prevented")
print("✓ Batch output files can be fetched with Azure credentials")
print("✓ Cost and usage tracking works for Azure batches")
print("✓ Backwards compatibility maintained\n")
_OPENAI_FILE_JSON: Final = MappingProxyType(
{
"id": "file-abc123",
"object": "file",
"purpose": "batch",
"filename": "batch.jsonl",
"bytes": 416,
"created_at": 1739598666,
"status": "processed",
}
)
_OPENAI_BATCH_JSON: Final = MappingProxyType(
{
"id": "batch_abc123",
"object": "batch",
"endpoint": "/v1/chat/completions",
"input_file_id": "file-abc123",
"status": "validating",
"completion_window": "24h",
"created_at": 1739598666,
}
)
class _KeyAliasMetadata(TypedDict):
user_api_key_alias: ReadOnly[str | None]
user_api_key_team_alias: ReadOnly[str | None]
class _LoggedCall(TypedDict):
call_type: ReadOnly[str]
metadata: ReadOnly[_KeyAliasMetadata]
_LOGGED_CALL: Final = TypeAdapter(_LoggedCall)
class _SuccessPayloadRecorder(CustomLogger):
def __init__(self, call_type: str) -> None:
super().__init__()
self._call_type: Final = call_type
self.logged: Final = asyncio.Event()
self.payload: _LoggedCall | None = None
async def async_log_success_event(
self,
kwargs: Mapping[str, object],
response_obj: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
) -> None:
payload: Final = _LOGGED_CALL.validate_python(kwargs["standard_logging_object"])
if payload["call_type"] != self._call_type:
return
self.payload = payload
self.logged.set()
@pytest.mark.asyncio
async def test_acreate_batch_full_crud_and_logging_metadata(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
litellm.logging_callback_manager._reset_all_callbacks()
recorder: Final = _SuccessPayloadRecorder("acreate_batch")
monkeypatch.setattr(litellm, "callbacks", [recorder])
upload_route: Final = respx_mock.post("https://api.openai.com/v1/files").mock(
return_value=httpx.Response(200, json=dict(_OPENAI_FILE_JSON))
)
create_route: Final = respx_mock.post("https://api.openai.com/v1/batches").mock(
return_value=httpx.Response(200, json=dict(_OPENAI_BATCH_JSON))
)
retrieve_route: Final = respx_mock.get("https://api.openai.com/v1/batches/batch_abc123").mock(
return_value=httpx.Response(200, json=dict(_OPENAI_BATCH_JSON))
)
list_batches_route: Final = respx_mock.get("https://api.openai.com/v1/batches").mock(
return_value=httpx.Response(200, json={"object": "list", "data": [dict(_OPENAI_BATCH_JSON)]})
)
respx_mock.get("https://api.openai.com/v1/files/file-abc123/content").mock(
return_value=httpx.Response(200, content=b'{"custom_id": "request-1"}\n')
)
respx_mock.get("https://api.openai.com/v1/files/file-abc123").mock(
return_value=httpx.Response(200, json=dict(_OPENAI_FILE_JSON))
)
respx_mock.delete("https://api.openai.com/v1/files/file-abc123").mock(
return_value=httpx.Response(200, json={"id": "file-abc123", "object": "file", "deleted": True})
)
list_files_route: Final = respx_mock.get("https://api.openai.com/v1/files").mock(
return_value=httpx.Response(200, json={"object": "list", "data": [dict(_OPENAI_FILE_JSON)]})
)
cancel_route: Final = respx_mock.post("https://api.openai.com/v1/batches/batch_abc123/cancel").mock(
return_value=httpx.Response(200, json={**_OPENAI_BATCH_JSON, "status": "cancelling"})
)
batch_file: Final = (
"batch.jsonl",
b'{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", '
b'"body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}}\n',
"application/jsonl",
)
file_obj: Final = await litellm.acreate_file(
file=batch_file, purpose="batch", custom_llm_provider="openai", api_key="fake-key"
)
assert file_obj.id == "file-abc123"
upload_body: Final = upload_route.calls.last.request.content
assert b'name="purpose"\r\n\r\nbatch' in upload_body
assert batch_file[1] in upload_body
extra_metadata_field: Final = {
"user_api_key_alias": "special_api_key_alias",
"user_api_key_team_alias": "special_team_alias",
}
create_batch_response: Final = await litellm.acreate_batch(
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id=file_obj.id,
custom_llm_provider="openai",
api_key="fake-key",
metadata={"key1": "value1", "key2": "value2"},
litellm_metadata=extra_metadata_field,
)
assert json.loads(create_route.calls.last.request.content) == {
"completion_window": "24h",
"endpoint": "/v1/chat/completions",
"input_file_id": "file-abc123",
"metadata": {"key1": "value1", "key2": "value2"},
}
assert create_batch_response.id == "batch_abc123"
assert create_batch_response.endpoint == "/v1/chat/completions"
assert create_batch_response.input_file_id == file_obj.id
await asyncio.wait_for(recorder.logged.wait(), timeout=10)
assert recorder.payload is not None
standard_logging_object: Final = recorder.payload
assert standard_logging_object["metadata"]["user_api_key_alias"] == extra_metadata_field["user_api_key_alias"]
assert (
standard_logging_object["metadata"]["user_api_key_team_alias"]
== extra_metadata_field["user_api_key_team_alias"]
)
retrieved_batch: Final = await litellm.aretrieve_batch(
batch_id=create_batch_response.id, custom_llm_provider="openai", api_key="fake-key"
)
assert retrieve_route.called
assert retrieved_batch.id == create_batch_response.id
list_batches: Final = await litellm.alist_batches(custom_llm_provider="openai", limit=2, api_key="fake-key")
assert list_batches_route.calls.last.request.url.params["limit"] == "2"
assert [batch.id for batch in list_batches.data] == ["batch_abc123"]
file_content: Final = await litellm.afile_content(
file_id=file_obj.id, custom_llm_provider="openai", api_key="fake-key"
)
assert file_content.content == b'{"custom_id": "request-1"}\n'
retrieved_file: Final = await litellm.afile_retrieve(
file_id=file_obj.id, custom_llm_provider="openai", api_key="fake-key"
)
assert retrieved_file.id == file_obj.id
delete_file_response: Final = await litellm.afile_delete(
file_id=file_obj.id, custom_llm_provider="openai", api_key="fake-key"
)
assert delete_file_response.id == file_obj.id
all_files_list: Final = await litellm.afile_list(custom_llm_provider="openai", api_key="fake-key")
assert list_files_route.called
assert [file.id for file in all_files_list.data] == ["file-abc123"]
cancel_batch_response: Final = await litellm.acancel_batch(
batch_id=create_batch_response.id, custom_llm_provider="openai", api_key="fake-key"
)
assert cancel_route.called
assert cancel_batch_response.id == create_batch_response.id
@pytest.mark.asyncio
async def test_delete_batch_output_file(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch):
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
batch_with_output: Final = {
**_OPENAI_BATCH_JSON,
"status": "completed",
"output_file_id": "file-output123",
}
respx_mock.get("https://api.openai.com/v1/batches/batch_abc123").mock(
return_value=httpx.Response(200, json=batch_with_output)
)
delete_route: Final = respx_mock.delete("https://api.openai.com/v1/files/file-output123").mock(
return_value=httpx.Response(200, json={"id": "file-output123", "object": "file", "deleted": True})
)
batch: Final = await litellm.aretrieve_batch(
batch_id="batch_abc123", custom_llm_provider="openai", api_key="fake-key"
)
assert batch.output_file_id == "file-output123"
delete_response: Final = await litellm.afile_delete(
file_id=batch.output_file_id, custom_llm_provider="openai", api_key="fake-key"
)
assert delete_route.call_count == 1
assert delete_response.id == "file-output123"
assert delete_response.deleted is True
VLLM_ROUTER_API_BASE: Final = "http://vllm-router.test/v1"
VLLM_ROUTER_BATCH: Final = MappingProxyType(
{
"id": "batch_5f8dac",
"status": "pending",
"input_file_id": "file-b73f4c",
"created_at": 1793239207,
"endpoint": "/v1/chat/completions",
"completion_window": "24h",
"output_file_id": None,
"total_requests": None,
"completed_requests": 0,
"failed_requests": 0,
"metadata": {},
}
)
async def _create_on_vllm_router() -> LiteLLMBatch:
return await litellm.acreate_batch(
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-b73f4c",
custom_llm_provider="hosted_vllm",
api_key="sk-test",
api_base=VLLM_ROUTER_API_BASE,
)
async def _retrieve_from_vllm_router() -> LiteLLMBatch:
return await litellm.aretrieve_batch(
batch_id="batch_5f8dac", custom_llm_provider="hosted_vllm", api_key="sk-test", api_base=VLLM_ROUTER_API_BASE
)
VLLM_ROUTER_BATCH_CALLS: Final = (
pytest.param("POST", "/batches", _create_on_vllm_router, id="acreate_batch"),
pytest.param("GET", "/batches/batch_5f8dac", _retrieve_from_vllm_router, id="aretrieve_batch"),
)
@pytest.mark.parametrize("method, path, call_router", VLLM_ROUTER_BATCH_CALLS)
@pytest.mark.parametrize(
"router_status, openai_status",
(
("pending", "validating"),
("running", "in_progress"),
("completed", "completed"),
("failed", "failed"),
("cancelled", "cancelled"),
),
)
async def test_hosted_vllm_router_batch_comes_back_as_an_openai_batch(
respx_mock: respx.MockRouter,
method: str,
path: str,
call_router: Callable[[], Awaitable[LiteLLMBatch]],
router_status: str,
openai_status: str,
):
respx_mock.route(method=method, url=f"{VLLM_ROUTER_API_BASE}{path}").respond(
200, json={**VLLM_ROUTER_BATCH, "status": router_status}
)
batch: Final = await call_router()
assert (batch.id, batch.object, batch.status) == ("batch_5f8dac", "batch", openai_status)
@pytest.mark.parametrize("method, path, call_router", VLLM_ROUTER_BATCH_CALLS)
@pytest.mark.parametrize("field, value", (("status", "queued"), ("object", "file")))
async def test_hosted_vllm_router_batch_with_an_unknown_value_is_rejected(
respx_mock: respx.MockRouter,
method: str,
path: str,
call_router: Callable[[], Awaitable[LiteLLMBatch]],
field: str,
value: str,
):
respx_mock.route(method=method, url=f"{VLLM_ROUTER_API_BASE}{path}").respond(
200, json={**VLLM_ROUTER_BATCH, field: value}
)
with pytest.raises(ValidationError) as exc_info:
await call_router()
assert [error["loc"] for error in exc_info.value.errors()] == [(field,)]
async def _list_on_vllm_router() -> Sequence[Batch]:
page: Final = await litellm.alist_batches(
custom_llm_provider="hosted_vllm", api_key="sk-test", api_base=VLLM_ROUTER_API_BASE
)
return page.data
async def _list_on_vllm_router_sync() -> Sequence[Batch]:
return litellm.list_batches(custom_llm_provider="hosted_vllm", api_key="sk-test", api_base=VLLM_ROUTER_API_BASE).data
@pytest.mark.parametrize("list_router", (_list_on_vllm_router, _list_on_vllm_router_sync))
async def test_hosted_vllm_router_batch_list_comes_back_as_openai_batches(
respx_mock: respx.MockRouter, list_router: Callable[[], Awaitable[Sequence[Batch]]]
):
respx_mock.get(f"{VLLM_ROUTER_API_BASE}/batches").respond(
200,
json={
"object": "list",
"data": [
{**VLLM_ROUTER_BATCH, "id": "batch_5f8dac", "status": "pending"},
{**VLLM_ROUTER_BATCH, "id": "batch_c74802", "status": "running"},
],
"first_id": "batch_5f8dac",
"last_id": "batch_c74802",
"has_more": False,
},
)
batches: Final = await list_router()
assert [(batch.id, batch.object, batch.status) for batch in batches] == [
("batch_5f8dac", "batch", "validating"),
("batch_c74802", "batch", "in_progress"),
]
async def test_hosted_vllm_router_batch_list_passes_through_a_status_that_is_not_a_string(
respx_mock: respx.MockRouter,
):
respx_mock.get(f"{VLLM_ROUTER_API_BASE}/batches").respond(
200,
json={
"object": "list",
"data": [{**VLLM_ROUTER_BATCH, "status": ["pending"]}],
"first_id": "batch_5f8dac",
"last_id": "batch_5f8dac",
"has_more": False,
},
)
batches: Final = await _list_on_vllm_router()
assert [(batch.id, batch.object, batch.status) for batch in batches] == [("batch_5f8dac", "batch", ["pending"])]
async def test_openai_batch_list_items_come_back_as_openai_sent_them(respx_mock: respx.MockRouter):
respx_mock.get("http://openai.test/v1/batches").respond(
200,
json={
"object": "list",
"data": [
{
**_OPENAI_BATCH_JSON,
"status": "completed",
"usage": {
"input_tokens": 10,
"input_tokens_details": {"cached_tokens": 4},
"output_tokens": 5,
"output_tokens_details": {"reasoning_tokens": 2},
"total_tokens": 15,
},
},
{**_OPENAI_BATCH_JSON, "id": "batch_def456", "status": "queued"},
],
"first_id": "batch_abc123",
"last_id": "batch_def456",
"has_more": False,
},
)
page: Final = await litellm.alist_batches(
custom_llm_provider="openai", api_key="sk-test", api_base="http://openai.test/v1"
)
first, second = page.data
assert first.usage is not None
assert (first.usage.input_tokens_details.cached_tokens, first.usage.output_tokens_details.reasoning_tokens) == (4, 2)
assert (second.id, second.object, second.status) == ("batch_def456", "batch", "queued")