mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
Already shaped ProxyException and HTTPException errors passing through the moderations, audio speech, Anthropic Messages, and handle_exception_on_proxy paths now answer with the x-litellm-call-id header the route logged under, without overwriting a header the exception was raised with. The GET /v1/batches failure hook receives the resolved request data so the spend log request_id matches the response header and the error log Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2805 lines
107 KiB
Python
2805 lines
107 KiB
Python
"""
|
|
Routing-contract tests for litellm/proxy/batches_endpoints/endpoints.py
|
|
|
|
These are not happy-path smoke tests. Each row of the matrix locks the full
|
|
contract of a single routing branch so that *any* behavior change in this layer
|
|
fails loudly:
|
|
|
|
1. DISPATCH - exactly which downstream seam fired (litellm.acreate_batch
|
|
vs llm_router.acreate_batch), and every sibling seam is
|
|
asserted NOT called. A reordered/negated branch flips this.
|
|
2. CREDENTIALS - the credential resolver receives the model derived from the
|
|
request, not a hardcoded value. The router's
|
|
get_deployment_credentials_with_provider is input-locked.
|
|
3. SEAM PAYLOAD - the *entire* kwargs dict forwarded to the provider call is
|
|
exact-matched. Because LiteLLMBatchCreateRequest is a
|
|
TypedDict (zero runtime filtering), nothing else stops a
|
|
newly-added param from silently reaching every provider.
|
|
This exact-match is that missing guard: a new key fails the
|
|
test and forces a reviewer to ask "does this work for all
|
|
providers, or just openai".
|
|
4. OUTPUT SHAPE - the id encode/decode round-trip clients depend on.
|
|
|
|
Only true I/O boundaries are mocked (provider call, router, proxy logging,
|
|
request parsing, pre-call enrichment). The pure encode/decode/credential-merge
|
|
helpers run for real so the payload assertions reflect production exactly.
|
|
|
|
The object mocks are spec'd to their real classes, so a brand-new method call
|
|
added to this layer raises instead of silently passing - the inventory of seams
|
|
cannot drift without a test failure.
|
|
"""
|
|
|
|
import base64
|
|
import json
|
|
import logging
|
|
from contextlib import ExitStack
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, Optional
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
|
|
|
|
import litellm
|
|
import litellm.proxy.batches_endpoints.endpoints as endpoints
|
|
import litellm.proxy.proxy_server as proxy_server
|
|
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
|
|
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
|
|
from litellm.proxy.openai_files_endpoints.common_utils import (
|
|
encode_file_id_with_model,
|
|
)
|
|
from litellm.proxy.utils import ProxyLogging
|
|
from litellm.router import Router
|
|
from litellm.types.llms.openai import BatchJobStatus
|
|
from litellm.types.utils import CredentialItem, LiteLLMBatch
|
|
|
|
from fastapi import Request, Response
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Fixtures: distinguishable credentials per model so a wrong/hardcoded model_id
|
|
# produces wrong creds (or KeyError) and is impossible to hide.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
CREDS: Dict[str, Dict[str, str]] = {
|
|
"azure/gpt-4o": {
|
|
"custom_llm_provider": "azure",
|
|
"api_key": "sk-azure",
|
|
"api_base": "https://azure.test",
|
|
"model": "azure/gpt-4o-deployment",
|
|
},
|
|
"vertex-model": {
|
|
"custom_llm_provider": "vertex_ai",
|
|
"api_key": "sk-vertex",
|
|
"api_base": "https://vertex.test",
|
|
"model": "vertex_ai/gemini-2.0",
|
|
},
|
|
}
|
|
|
|
# A real model-encoded file id: decodes to "azure/gpt-4o", strips to "file-original123".
|
|
AZURE_FILE_ID = encode_file_id_with_model("file-original123", "azure/gpt-4o", id_type="file")
|
|
|
|
|
|
def make_batch(
|
|
*,
|
|
id: str = "batch-provider-id",
|
|
output_file_id: Optional[str] = None,
|
|
error_file_id: Optional[str] = None,
|
|
input_file_id: Optional[str] = None,
|
|
status: BatchJobStatus = "validating",
|
|
) -> LiteLLMBatch:
|
|
batch = LiteLLMBatch(
|
|
id=id,
|
|
completion_window="24h",
|
|
created_at=1234567890,
|
|
endpoint="/v1/chat/completions",
|
|
input_file_id=input_file_id or "file-provider-input",
|
|
object="batch",
|
|
status=status,
|
|
)
|
|
if output_file_id is not None:
|
|
batch.output_file_id = output_file_id
|
|
if error_file_id is not None:
|
|
batch.error_file_id = error_file_id
|
|
batch._hidden_params = {}
|
|
return batch
|
|
|
|
|
|
class FakeRequest:
|
|
"""Minimal stand-in. The request is only read via .headers/.query_params on
|
|
the model-param fallback path; everything else that touches it is mocked."""
|
|
|
|
def __init__(
|
|
self,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
query: Optional[Dict[str, str]] = None,
|
|
):
|
|
self.headers = headers or {}
|
|
self.query_params = query or {}
|
|
|
|
|
|
@pytest.fixture
|
|
def openai_env_creds(monkeypatch):
|
|
"""Deterministic env creds so the implicit-openai fallback forwards instead
|
|
of tripping the no-creds 404 gate, regardless of the host environment."""
|
|
monkeypatch.setenv("OPENAI_API_KEY", "sk-env-openai")
|
|
|
|
|
|
@pytest.fixture
|
|
def no_openai_creds(monkeypatch):
|
|
"""Neutralize every credential source the 404 gate checks."""
|
|
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
|
|
monkeypatch.setattr(litellm, "api_key", None)
|
|
monkeypatch.setattr(litellm, "openai_key", None)
|
|
|
|
|
|
@dataclass
|
|
class Harness:
|
|
"""Holds every mocked seam so a test can configure inputs and assert calls."""
|
|
|
|
body: Dict[str, Any]
|
|
read_body: AsyncMock
|
|
pre_call: AsyncMock
|
|
get_headers: MagicMock
|
|
provider_from_headers: MagicMock
|
|
is_known_model: MagicMock
|
|
litellm_acreate: AsyncMock
|
|
router: MagicMock
|
|
logging: MagicMock
|
|
creds_resolver: MagicMock
|
|
|
|
@property
|
|
def router_acreate(self) -> AsyncMock:
|
|
return self.router.acreate_batch
|
|
|
|
def acreate_kwargs(self) -> Dict[str, Any]:
|
|
"""Exact kwargs forwarded to litellm.acreate_batch."""
|
|
assert self.litellm_acreate.call_count == 1
|
|
return dict(self.litellm_acreate.call_args.kwargs)
|
|
|
|
def router_kwargs(self) -> Dict[str, Any]:
|
|
assert self.router_acreate.call_count == 1
|
|
return dict(self.router_acreate.call_args.kwargs)
|
|
|
|
|
|
def _creds_lookup(*, model_id: str) -> Dict[str, str]:
|
|
# KeyError on an unknown/hardcoded model_id - the bug cannot hide.
|
|
return dict(CREDS[model_id])
|
|
|
|
|
|
@pytest.fixture
|
|
def harness():
|
|
"""Seam harness. Patches only true I/O boundaries; pure encode/decode/merge
|
|
helpers run for real. Object mocks are spec'd so unknown method calls raise."""
|
|
body_holder: Dict[str, Any] = {}
|
|
logging = MagicMock(spec=ProxyLogging)
|
|
logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"])
|
|
logging.post_call_failure_hook = AsyncMock()
|
|
logging.update_request_status = AsyncMock()
|
|
logging.get_proxy_hook = MagicMock(return_value=None)
|
|
|
|
router = MagicMock(spec=Router)
|
|
router.acreate_batch = AsyncMock(return_value=make_batch())
|
|
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
|
|
|
read_body = AsyncMock(side_effect=lambda request: body_holder["body"])
|
|
pre_call = AsyncMock(side_effect=lambda **kw: (body_holder["body"], MagicMock()))
|
|
get_headers = MagicMock(return_value={})
|
|
provider_from_headers = MagicMock(return_value=None)
|
|
is_known_model = MagicMock(return_value=False)
|
|
litellm_acreate = AsyncMock(return_value=make_batch())
|
|
|
|
with ExitStack() as stack:
|
|
stack.enter_context(patch.object(endpoints, "_read_request_body", read_body))
|
|
stack.enter_context(
|
|
patch.object(
|
|
ProxyBaseLLMRequestProcessing,
|
|
"common_processing_pre_call_logic",
|
|
pre_call,
|
|
)
|
|
)
|
|
stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
|
|
stack.enter_context(
|
|
patch.object(
|
|
endpoints,
|
|
"get_custom_llm_provider_from_request_headers",
|
|
provider_from_headers,
|
|
)
|
|
)
|
|
stack.enter_context(patch.object(endpoints, "is_known_model", is_known_model))
|
|
stack.enter_context(patch.object(litellm, "acreate_batch", litellm_acreate))
|
|
stack.enter_context(patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False))
|
|
stack.enter_context(patch.object(proxy_server, "llm_router", router))
|
|
stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging))
|
|
stack.enter_context(patch.object(proxy_server, "general_settings", {}))
|
|
stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock()))
|
|
stack.enter_context(patch.object(proxy_server, "version", "test-version"))
|
|
stack.enter_context(patch.object(proxy_server, "prisma_client", None))
|
|
|
|
h = Harness(
|
|
body=body_holder,
|
|
read_body=read_body,
|
|
pre_call=pre_call,
|
|
get_headers=get_headers,
|
|
provider_from_headers=provider_from_headers,
|
|
is_known_model=is_known_model,
|
|
litellm_acreate=litellm_acreate,
|
|
router=router,
|
|
logging=logging,
|
|
creds_resolver=router.get_deployment_credentials_with_provider,
|
|
)
|
|
yield h
|
|
|
|
|
|
def set_body(harness: Harness, body: Dict[str, Any]) -> None:
|
|
harness.body["body"] = body
|
|
|
|
|
|
async def call_create(
|
|
harness: Harness,
|
|
*,
|
|
provider: Optional[str] = None,
|
|
user: Optional[UserAPIKeyAuth] = None,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
query: Optional[Dict[str, str]] = None,
|
|
):
|
|
return await endpoints.create_batch(
|
|
request=FakeRequest(headers=headers, query=query),
|
|
fastapi_response=Response(),
|
|
provider=provider,
|
|
user_api_key_dict=user or UserAPIKeyAuth(api_key="sk-test"),
|
|
)
|
|
|
|
|
|
# =========================================================================== #
|
|
# SCENARIO 1 - input_file_id encoded with model. The full showcase: every
|
|
# assertion type from the design lives here.
|
|
# =========================================================================== #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__model_encoded_file_id(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": AZURE_FILE_ID,
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
resp = await call_create(harness)
|
|
|
|
# 1. DISPATCH - model-credential path fired via litellm, router did not.
|
|
assert harness.litellm_acreate.call_count == 1
|
|
harness.router_acreate.assert_not_called()
|
|
|
|
# 2. CREDENTIALS - resolved for the model decoded FROM the file id.
|
|
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
|
|
|
# 3. SEAM PAYLOAD - exact, whole dict. A new forwarded key breaks this.
|
|
assert harness.acreate_kwargs() == {
|
|
"custom_llm_provider": "azure",
|
|
"input_file_id": "file-original123", # encoding stripped by this layer
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"metadata": None, # sanitize_openai_provider_metadata(None)
|
|
"api_key": "sk-azure",
|
|
"api_base": "https://azure.test",
|
|
"model": "azure/gpt-4o-deployment",
|
|
}
|
|
|
|
# 4. OUTPUT SHAPE - ids re-encoded with the model; input_file_id restored.
|
|
assert resp.id == encode_file_id_with_model("batch-provider-id", "azure/gpt-4o", id_type="batch")
|
|
assert resp.input_file_id == AZURE_FILE_ID
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__model_encoded_file_id__encodes_output_and_error_ids(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": AZURE_FILE_ID,
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
harness.litellm_acreate.return_value = make_batch(
|
|
id="batch-xyz",
|
|
output_file_id="file-out-raw",
|
|
error_file_id="file-err-raw",
|
|
)
|
|
|
|
resp = await call_create(harness)
|
|
|
|
assert resp.output_file_id == encode_file_id_with_model("file-out-raw", "azure/gpt-4o")
|
|
assert resp.error_file_id == encode_file_id_with_model("file-err-raw", "azure/gpt-4o")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__model_encoded_file_id__resolver_gets_decoded_model(harness):
|
|
"""Regression guard: model_id for credential resolution must be derived from
|
|
the file id. A hardcode would call the resolver with the wrong model."""
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": AZURE_FILE_ID,
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
await call_create(harness)
|
|
|
|
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
|
|
|
|
|
# =========================================================================== #
|
|
# SCENARIO 2 - model from body / header / query. Locks source precedence.
|
|
# =========================================================================== #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__model_from_body(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"model": "vertex-model",
|
|
},
|
|
)
|
|
|
|
resp = await call_create(harness)
|
|
|
|
assert harness.litellm_acreate.call_count == 1
|
|
harness.router_acreate.assert_not_called()
|
|
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
|
|
payload = harness.acreate_kwargs()
|
|
assert payload["custom_llm_provider"] == "vertex_ai"
|
|
assert payload["input_file_id"] == "file-plain"
|
|
assert resp.id == encode_file_id_with_model("batch-provider-id", "vertex-model", id_type="batch")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__model_from_header(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
await call_create(harness, headers={"x-litellm-model": "vertex-model"})
|
|
|
|
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
|
|
harness.router_acreate.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__model_from_query(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
await call_create(harness, query={"model": "vertex-model"})
|
|
|
|
harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
|
|
harness.router_acreate.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__body_model_beats_header_and_query(harness):
|
|
"""Precedence row: body > header > query (data.get('model') first)."""
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"model": "azure/gpt-4o",
|
|
},
|
|
)
|
|
|
|
await call_create(
|
|
harness,
|
|
headers={"x-litellm-model": "vertex-model"},
|
|
query={"model": "vertex-model"},
|
|
)
|
|
|
|
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
|
|
|
|
|
# =========================================================================== #
|
|
# SCENARIO 3 - fallback to custom_llm_provider (env-var creds). MUST NOT touch
|
|
# the credential resolver.
|
|
# =========================================================================== #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__fallback_default_openai(harness, openai_env_creds):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
await call_create(harness)
|
|
|
|
assert harness.litellm_acreate.call_count == 1
|
|
harness.router_acreate.assert_not_called()
|
|
harness.creds_resolver.assert_not_called() # inverse-bug guard
|
|
assert harness.acreate_kwargs()["custom_llm_provider"] == "openai"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__fallback_no_creds_404(harness, no_openai_creds):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_create(harness)
|
|
|
|
assert exc.value.code == "404"
|
|
assert exc.value.type == "invalid_request_error"
|
|
assert exc.value.param is None
|
|
assert exc.value.message == "No such File object: file-plain"
|
|
harness.litellm_acreate.assert_not_called()
|
|
harness.router_acreate.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__fallback_explicit_provider_bypasses_not_found_gate(harness, no_openai_creds):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
await call_create(harness, provider="anthropic")
|
|
|
|
assert harness.acreate_kwargs()["custom_llm_provider"] == "anthropic"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__fallback_env_key_alone_forwards(harness, monkeypatch):
|
|
monkeypatch.setattr(litellm, "api_key", None)
|
|
monkeypatch.setattr(litellm, "openai_key", None)
|
|
monkeypatch.setenv("OPENAI_API_KEY", "sk-env-openai")
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
await call_create(harness)
|
|
|
|
assert harness.acreate_kwargs()["custom_llm_provider"] == "openai"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__fallback_provider_path_param(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
await call_create(harness, provider="anthropic")
|
|
|
|
harness.creds_resolver.assert_not_called()
|
|
assert harness.acreate_kwargs()["custom_llm_provider"] == "anthropic"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__fallback_body_custom_llm_provider(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"custom_llm_provider": "bedrock",
|
|
},
|
|
)
|
|
|
|
await call_create(harness)
|
|
|
|
payload = harness.acreate_kwargs()
|
|
assert payload["custom_llm_provider"] == "bedrock"
|
|
|
|
|
|
# =========================================================================== #
|
|
# Unified file id routing (-> llm_router). Helpers mocked only here because a
|
|
# real unified id is opaque base64; the routing contract is what we lock.
|
|
# =========================================================================== #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__unified_file_id_single_model_disables_cross_model_fallbacks(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "litellm_proxy_unified_id",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"disable_fallbacks": False,
|
|
},
|
|
)
|
|
with (
|
|
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
|
|
patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gpt-4o-mini"]),
|
|
):
|
|
resp = await call_create(harness)
|
|
|
|
# DISPATCH - router fired, direct litellm did not.
|
|
assert harness.router_acreate.call_count == 1
|
|
harness.litellm_acreate.assert_not_called()
|
|
# model injected from the unified id, input_file_id restored, hidden param set
|
|
assert harness.router_kwargs()["model"] == "gpt-4o-mini"
|
|
assert harness.router_kwargs()["disable_fallbacks"] is True
|
|
assert resp.input_file_id == "litellm_proxy_unified_id"
|
|
assert resp._hidden_params["unified_file_id"] == "unified-xyz"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("models", [[], ["m1", "m2"]])
|
|
async def test_create__unified_file_id_not_exactly_one_model_400(harness, models):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "litellm_proxy_unified_id",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
with (
|
|
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
|
|
patch.object(endpoints, "get_models_from_unified_file_id", return_value=models),
|
|
):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_create(harness)
|
|
|
|
assert exc.value.code == "400"
|
|
harness.router_acreate.assert_not_called()
|
|
harness.litellm_acreate.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__unified_file_id_resolves_real_storage_url(harness):
|
|
"""A base64 unified_file_id is a LiteLLM-internal token, not a real
|
|
provider-side file reference (e.g. Vertex AI's batch transformation parses
|
|
a `publishers/` segment out of the file URI and crashes on the opaque
|
|
base64 string). The real backend location (`storage_url`) must be looked
|
|
up from LiteLLM_ManagedFileTable and substituted before dispatch.
|
|
|
|
Regression lock on the lookup key: LiteLLM_ManagedFileTable.unified_file_id
|
|
stores the raw base64 file id (see schema.prisma and the enterprise
|
|
managed-files hook, which queries with the raw id), NOT the decoded
|
|
litellm_proxy:... string. Querying with the decoded string never matches
|
|
and silently falls back."""
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "litellm_proxy_unified_id",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
fake_db_file = MagicMock(
|
|
storage_url="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.0/abc",
|
|
)
|
|
find_first = AsyncMock(return_value=fake_db_file)
|
|
fake_repo_instance = MagicMock()
|
|
fake_repo_instance.table.find_first = find_first
|
|
fake_repo_cls = MagicMock(return_value=fake_repo_instance)
|
|
|
|
with (
|
|
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
|
|
patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]),
|
|
patch.object(proxy_server, "prisma_client", MagicMock()),
|
|
patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
|
|
):
|
|
resp = await call_create(harness)
|
|
|
|
assert harness.router_kwargs()["input_file_id"] == fake_db_file.storage_url
|
|
find_first.assert_awaited_once_with(where={"unified_file_id": "litellm_proxy_unified_id"})
|
|
assert resp.input_file_id == "litellm_proxy_unified_id"
|
|
assert resp._hidden_params["unified_file_id"] == "unified-xyz"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__unified_file_id_db_error_falls_back_to_raw_id(harness):
|
|
"""Resolution is additive and best-effort: a lookup error leaves the id
|
|
unresolved and dispatch falls back to the original id, exactly as before
|
|
this change (the managed-files deployment hook still maps it). No new
|
|
failure mode is introduced."""
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "litellm_proxy_unified_id",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
find_first = AsyncMock(side_effect=Exception("db unavailable"))
|
|
fake_repo_instance = MagicMock()
|
|
fake_repo_instance.table.find_first = find_first
|
|
fake_repo_cls = MagicMock(return_value=fake_repo_instance)
|
|
|
|
with (
|
|
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
|
|
patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]),
|
|
patch.object(proxy_server, "prisma_client", MagicMock()),
|
|
patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
|
|
):
|
|
await call_create(harness)
|
|
|
|
assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__multi_model_unified_file_with_loadbalancing_keeps_router_branch(harness):
|
|
"""Regression guard: a multi-model managed file dispatched with an explicit
|
|
router model under load balancing must keep taking the load-balanced router
|
|
branch, exactly as on the base revision, where the managed-files deployment
|
|
hook remaps the unified id per model. Routing it into the unified branch
|
|
instead would trip that branch's "exactly one model" 400 and break a path
|
|
that works today, so the unified-file resolution must not steal the
|
|
load-balanced branch."""
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "litellm_proxy_unified_id",
|
|
"model": "vertex-model",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
harness.is_known_model.return_value = True
|
|
|
|
with (
|
|
patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True),
|
|
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
|
|
patch.object(endpoints, "get_models_from_unified_file_id", return_value=["model-a", "model-b"]),
|
|
):
|
|
await call_create(harness)
|
|
|
|
assert harness.router_acreate.call_count == 1
|
|
assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"
|
|
harness.litellm_acreate.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__unified_file_id_missing_row_falls_back_to_raw_id(harness):
|
|
"""Resolution is additive: when no managed-file row exists there is nothing
|
|
to substitute, so dispatch falls back to the original id exactly as before
|
|
this change (the managed-files deployment hook still maps it). No new
|
|
failure mode is introduced for this case."""
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "litellm_proxy_unified_id",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
find_first = AsyncMock(return_value=None)
|
|
fake_repo_instance = MagicMock()
|
|
fake_repo_instance.table.find_first = find_first
|
|
fake_repo_cls = MagicMock(return_value=fake_repo_instance)
|
|
|
|
with (
|
|
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
|
|
patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]),
|
|
patch.object(proxy_server, "prisma_client", MagicMock()),
|
|
patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
|
|
):
|
|
await call_create(harness)
|
|
|
|
assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__unified_file_id_legacy_row_without_storage_url_dispatches_raw(
|
|
harness,
|
|
):
|
|
"""A managed file whose row predates the storage_url column still dispatches
|
|
the original id (the managed-files deployment hook maps it); the row exists,
|
|
so this is not the missing-row fail-closed case."""
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "litellm_proxy_unified_id",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
fake_db_file = MagicMock(storage_url=None)
|
|
find_first = AsyncMock(return_value=fake_db_file)
|
|
fake_repo_instance = MagicMock()
|
|
fake_repo_instance.table.find_first = find_first
|
|
fake_repo_cls = MagicMock(return_value=fake_repo_instance)
|
|
|
|
with (
|
|
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
|
|
patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]),
|
|
patch.object(proxy_server, "prisma_client", MagicMock()),
|
|
patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
|
|
):
|
|
await call_create(harness)
|
|
|
|
assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__model_encoded_beats_unified(harness):
|
|
"""Precedence row: a file id that is BOTH model-encoded and (pretend) unified
|
|
must take the model-encoded branch (checked first)."""
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": AZURE_FILE_ID,
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
with (
|
|
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
|
|
patch.object(endpoints, "get_models_from_unified_file_id", return_value=["something-else"]),
|
|
):
|
|
await call_create(harness)
|
|
|
|
assert harness.litellm_acreate.call_count == 1
|
|
harness.router_acreate.assert_not_called()
|
|
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
|
|
|
|
|
# =========================================================================== #
|
|
# Loadbalancing branch (-> llm_router) and its precedence vs model-encoded.
|
|
# =========================================================================== #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__loadbalancing_routes_to_router(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"model": "lb-model",
|
|
},
|
|
)
|
|
harness.is_known_model.return_value = True
|
|
with patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True):
|
|
await call_create(harness)
|
|
|
|
harness.is_known_model.assert_called_once_with(model="lb-model", llm_router=harness.router)
|
|
assert harness.router_acreate.call_count == 1
|
|
harness.litellm_acreate.assert_not_called()
|
|
harness.creds_resolver.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__model_encoded_beats_loadbalancing(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": AZURE_FILE_ID,
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"model": "lb-model",
|
|
},
|
|
)
|
|
harness.is_known_model.return_value = True
|
|
with patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True):
|
|
await call_create(harness)
|
|
|
|
assert harness.litellm_acreate.call_count == 1
|
|
harness.router_acreate.assert_not_called()
|
|
harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"body, missing_param",
|
|
[
|
|
({"endpoint": "/v1/chat/completions", "completion_window": "24h"}, "input_file_id"),
|
|
({"input_file_id": "file-abc", "completion_window": "24h"}, "endpoint"),
|
|
({"input_file_id": "file-abc", "endpoint": "/v1/chat/completions"}, "completion_window"),
|
|
({}, "input_file_id"),
|
|
],
|
|
)
|
|
async def test_create__missing_required_param_is_400(harness, body, missing_param):
|
|
set_body(harness, body)
|
|
|
|
with pytest.raises(ProxyException) as exc_info:
|
|
await call_create(harness)
|
|
|
|
assert exc_info.value.code == "400"
|
|
assert exc_info.value.type == "invalid_request_error"
|
|
assert exc_info.value.param == missing_param
|
|
assert exc_info.value.message == f"/batches: Missing required parameter: '{missing_param}'."
|
|
harness.litellm_acreate.assert_not_called()
|
|
harness.router_acreate.assert_not_called()
|
|
|
|
|
|
def _raw_batches_request(body: Dict[str, Any]) -> MagicMock:
|
|
"""A request that reaches the real pre-call logic, which the `harness` fixture
|
|
mocks out. Metadata validation lives there, so it cannot be seen through the seam."""
|
|
request = MagicMock(spec=Request)
|
|
request.url = MagicMock()
|
|
request.url.__str__.return_value = "http://localhost/v1/batches"
|
|
request.url.path = "/v1/batches"
|
|
request.method = "POST"
|
|
request.query_params = {}
|
|
request.headers = {"Content-Type": "application/json"}
|
|
request.client = MagicMock()
|
|
request.client.host = "127.0.0.1"
|
|
request.body = AsyncMock(return_value=json.dumps(body).encode())
|
|
return request
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("field", ["metadata", "litellm_metadata"])
|
|
async def test_create__non_object_metadata_is_400(field):
|
|
"""A non-object metadata field is rejected with a 400 naming it (#37147), rather
|
|
than being coerced to None and silently dropped, or crashing behind a 500 with
|
|
"'str' object has no attribute 'update'"."""
|
|
body = {
|
|
"input_file_id": "file-abc",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
field: "abc",
|
|
}
|
|
|
|
with pytest.raises(ProxyException) as exc_info:
|
|
await endpoints.create_batch(
|
|
request=_raw_batches_request(body),
|
|
fastapi_response=Response(),
|
|
provider=None,
|
|
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
|
|
)
|
|
|
|
assert exc_info.value.code == "400"
|
|
assert exc_info.value.param == field
|
|
assert "has no attribute 'update'" not in exc_info.value.message
|
|
|
|
|
|
# =========================================================================== #
|
|
# Team-level batch expiry enforcement (independent of routing).
|
|
# =========================================================================== #
|
|
|
|
|
|
def _user_with_expiry(expiry: Any) -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth(
|
|
api_key="sk-test",
|
|
team_metadata={"enforced_batch_output_expires_after": expiry},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__team_expiry_injected(harness, openai_env_creds):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
await call_create(harness, user=_user_with_expiry({"anchor": "created_at", "seconds": 3600}))
|
|
|
|
assert harness.acreate_kwargs()["output_expires_after"] == {
|
|
"anchor": "created_at",
|
|
"seconds": 3600,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__no_team_expiry_not_injected(harness, openai_env_creds):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
await call_create(harness, user=UserAPIKeyAuth(api_key="sk-test"))
|
|
|
|
assert "output_expires_after" not in harness.acreate_kwargs()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"expiry",
|
|
[
|
|
{"seconds": 3600}, # missing anchor
|
|
{"anchor": "created_at"}, # missing seconds
|
|
{"anchor": "completed_at", "seconds": 3600}, # wrong anchor
|
|
],
|
|
)
|
|
async def test_create__team_expiry_malformed_500(harness, expiry):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_create(harness, user=_user_with_expiry(expiry))
|
|
|
|
assert exc.value.code == "500"
|
|
harness.litellm_acreate.assert_not_called()
|
|
harness.router_acreate.assert_not_called()
|
|
|
|
|
|
# =========================================================================== #
|
|
# Cross-cutting: enrichment route_type, metadata sanitization, failure hook.
|
|
# =========================================================================== #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__uses_acreate_batch_route_type(harness, openai_env_creds):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
await call_create(harness)
|
|
|
|
assert harness.pre_call.call_args.kwargs["route_type"] == "acreate_batch"
|
|
|
|
|
|
def install_managed_files_hook(harness: Harness) -> AsyncMock:
|
|
prisma_client = AsyncMock()
|
|
managed_files = _PROXY_LiteLLMManagedFiles(MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client)
|
|
harness.logging.post_call_success_hook = AsyncMock(side_effect=managed_files.async_post_call_success_hook)
|
|
harness.router.model_list = []
|
|
return prisma_client
|
|
|
|
|
|
TEAM_A_KEY = UserAPIKeyAuth(api_key="sk-team-a", user_id="user_a", team_id="team_a")
|
|
|
|
|
|
def assert_ownership_registered_for_team_a(prisma_client: AsyncMock, batch_id: str) -> None:
|
|
upsert = prisma_client.db.litellm_managedobjecttable.upsert
|
|
upsert.assert_awaited_once()
|
|
assert upsert.await_args.kwargs["where"] == {"unified_object_id": batch_id}
|
|
created = upsert.await_args.kwargs["data"]["create"]
|
|
assert created["created_by"] == "user_a"
|
|
assert created["team_id"] == "team_a"
|
|
prisma_client.db.litellm_managedobjecttable.update_many.assert_not_awaited()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"body",
|
|
[
|
|
{"input_file_id": AZURE_FILE_ID},
|
|
{"input_file_id": "file-plain", "model": "vertex-model"},
|
|
{"input_file_id": "file-plain"},
|
|
],
|
|
ids=["model_encoded_file_id", "model_param", "provider_fallback"],
|
|
)
|
|
async def test_create__registers_ownership_for_creator(harness, openai_env_creds, body):
|
|
set_body(harness, {**body, "endpoint": "/v1/chat/completions", "completion_window": "24h"})
|
|
prisma_client = install_managed_files_hook(harness)
|
|
|
|
resp = await call_create(harness, user=TEAM_A_KEY)
|
|
|
|
assert_ownership_registered_for_team_a(prisma_client, resp.id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__unified_file_id_registers_ownership_for_creator(harness):
|
|
unified_input_file_id = base64.urlsafe_b64encode(
|
|
b"litellm_proxy:application/octet-stream;unified_id,input-uuid;target_model_names,gpt-4o-mini"
|
|
).decode()
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": unified_input_file_id,
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
prisma_client = install_managed_files_hook(harness)
|
|
|
|
resp = await call_create(harness, user=TEAM_A_KEY)
|
|
|
|
assert harness.router_acreate.call_count == 1
|
|
assert_ownership_registered_for_team_a(prisma_client, resp.id)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__metadata_sanitized_before_forwarding(harness, openai_env_creds):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"metadata": {"user_key": "user_val", "spend_logs_metadata": {"x": 1}},
|
|
},
|
|
)
|
|
|
|
await call_create(harness)
|
|
|
|
# provider-internal-only key dropped, string key kept (real sanitize runs)
|
|
assert harness.acreate_kwargs()["metadata"] == {"user_key": "user_val"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__exception_calls_failure_hook(harness, openai_env_creds):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
harness.litellm_acreate.side_effect = ValueError("provider boom")
|
|
|
|
with pytest.raises(ProxyException):
|
|
await call_create(harness)
|
|
|
|
harness.logging.post_call_failure_hook.assert_called_once()
|
|
assert harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom"
|
|
|
|
|
|
async def test_create__exception_carries_the_litellm_call_id(harness, openai_env_creds, caplog):
|
|
call_id = "lit7836-batch-call-id"
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"litellm_call_id": call_id,
|
|
},
|
|
)
|
|
harness.litellm_acreate.side_effect = ValueError("provider boom")
|
|
|
|
with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised:
|
|
await call_create(harness)
|
|
|
|
assert raised.value.headers["x-litellm-call-id"] == call_id
|
|
record = next(r for r in caplog.records if "Exception occured" in r.getMessage())
|
|
assert record.litellm_call_id == call_id
|
|
assert call_id in record.getMessage()
|
|
|
|
|
|
# =========================================================================== #
|
|
# #
|
|
# GET /v1/batches/{batch_id} - retrieve_batch routing-contract tests #
|
|
# #
|
|
# Same discipline as create_batch above. retrieve_batch has more seams: it #
|
|
# first consults the ManagedObjectTable (get_batch_from_database) and may #
|
|
# short-circuit on a terminal-status row WITHOUT ever calling a provider, #
|
|
# then on a miss/non-terminal row routes to one of three downstream seams #
|
|
# (litellm.aretrieve_batch via model creds, llm_router.aretrieve_batch, or #
|
|
# litellm.aretrieve_batch via env-var provider) and writes the fresh state #
|
|
# back via update_batch_in_database. Every test below locks exactly which #
|
|
# of those seams fired and asserts the siblings did NOT, so a reordered or #
|
|
# negated branch - or a dropped DB short-circuit / write-back - fails loud. #
|
|
# #
|
|
# The DB seams (get_batch_from_database, update_batch_in_database, #
|
|
# resolve_*_to_unified) are true prisma I/O boundaries and are mocked. The #
|
|
# encode/decode/credential-merge helpers run for real, so payload and id #
|
|
# round-trip assertions reflect production exactly. #
|
|
# =========================================================================== #
|
|
|
|
|
|
# A real model-encoded BATCH id: decodes to "azure/gpt-4o", strips to
|
|
# "batch_orig123". Distinct from AZURE_FILE_ID so retrieve tests can't pass by
|
|
# accidentally reusing the create fixture's value.
|
|
AZURE_BATCH_ID = encode_file_id_with_model("batch_orig123", "azure/gpt-4o", id_type="batch")
|
|
|
|
# A realistic decoded unified batch id (what _is_base64_encoded_unified_file_id
|
|
# returns). model_id / llm_batch_id are parsed out of this by the real helpers.
|
|
UNIFIED_BATCH_ID = "litellm_proxy;model_id:gpt-4o-mini;llm_batch_id:batch-raw-xyz"
|
|
|
|
|
|
@dataclass
|
|
class RetrieveHarness:
|
|
"""Seams for retrieve_batch. `data['data']` is the dict pre-call enrichment
|
|
returns; the routing branches mutate it, so it is reset per call."""
|
|
|
|
data: Dict[str, Any]
|
|
pre_call: AsyncMock
|
|
get_headers: MagicMock
|
|
provider_from_headers: MagicMock
|
|
provider_from_query: MagicMock
|
|
litellm_aretrieve: AsyncMock
|
|
router: MagicMock
|
|
logging: MagicMock
|
|
creds_resolver: MagicMock
|
|
get_batch_from_db: AsyncMock
|
|
update_batch_in_db: AsyncMock
|
|
ensure_managed_files: AsyncMock
|
|
|
|
@property
|
|
def router_aretrieve(self) -> AsyncMock:
|
|
return self.router.aretrieve_batch
|
|
|
|
def aretrieve_kwargs(self) -> Dict[str, Any]:
|
|
"""Exact kwargs forwarded to litellm.aretrieve_batch."""
|
|
assert self.litellm_aretrieve.call_count == 1
|
|
return dict(self.litellm_aretrieve.call_args.kwargs)
|
|
|
|
def router_kwargs(self) -> Dict[str, Any]:
|
|
assert self.router_aretrieve.call_count == 1
|
|
return dict(self.router_aretrieve.call_args.kwargs)
|
|
|
|
|
|
@pytest.fixture
|
|
def retrieve_harness():
|
|
data_holder: Dict[str, Any] = {"data": {}}
|
|
logging = MagicMock(spec=ProxyLogging)
|
|
logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"])
|
|
logging.post_call_failure_hook = AsyncMock()
|
|
logging.update_request_status = AsyncMock()
|
|
logging.get_proxy_hook = MagicMock(return_value=None)
|
|
|
|
router = MagicMock(spec=Router)
|
|
router.aretrieve_batch = AsyncMock(return_value=make_batch())
|
|
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
|
|
|
pre_call = AsyncMock(side_effect=lambda **kw: (data_holder["data"], MagicMock()))
|
|
get_headers = MagicMock(return_value={})
|
|
provider_from_headers = MagicMock(return_value=None)
|
|
provider_from_query = MagicMock(return_value=None)
|
|
litellm_aretrieve = AsyncMock(return_value=make_batch())
|
|
# Default: DB miss -> always fall through to provider routing.
|
|
get_batch_from_db = AsyncMock(return_value=(None, None))
|
|
update_batch_in_db = AsyncMock(return_value=None)
|
|
ensure_managed_files = AsyncMock(return_value=None)
|
|
|
|
with ExitStack() as stack:
|
|
stack.enter_context(
|
|
patch.object(
|
|
ProxyBaseLLMRequestProcessing,
|
|
"common_processing_pre_call_logic",
|
|
pre_call,
|
|
)
|
|
)
|
|
stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
|
|
stack.enter_context(
|
|
patch.object(
|
|
endpoints,
|
|
"get_custom_llm_provider_from_request_headers",
|
|
provider_from_headers,
|
|
)
|
|
)
|
|
stack.enter_context(
|
|
patch.object(
|
|
endpoints,
|
|
"get_custom_llm_provider_from_request_query",
|
|
provider_from_query,
|
|
)
|
|
)
|
|
stack.enter_context(patch.object(endpoints, "get_batch_from_database", get_batch_from_db))
|
|
stack.enter_context(patch.object(endpoints, "update_batch_in_database", update_batch_in_db))
|
|
stack.enter_context(patch.object(endpoints, "ensure_batch_response_managed_file_ids", ensure_managed_files))
|
|
stack.enter_context(patch.object(litellm, "aretrieve_batch", litellm_aretrieve))
|
|
stack.enter_context(patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False))
|
|
stack.enter_context(patch.object(proxy_server, "llm_router", router))
|
|
stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging))
|
|
stack.enter_context(patch.object(proxy_server, "general_settings", {}))
|
|
stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock()))
|
|
stack.enter_context(patch.object(proxy_server, "version", "test-version"))
|
|
stack.enter_context(patch.object(proxy_server, "prisma_client", MagicMock()))
|
|
|
|
yield RetrieveHarness(
|
|
data=data_holder,
|
|
pre_call=pre_call,
|
|
get_headers=get_headers,
|
|
provider_from_headers=provider_from_headers,
|
|
provider_from_query=provider_from_query,
|
|
litellm_aretrieve=litellm_aretrieve,
|
|
router=router,
|
|
logging=logging,
|
|
creds_resolver=router.get_deployment_credentials_with_provider,
|
|
get_batch_from_db=get_batch_from_db,
|
|
update_batch_in_db=update_batch_in_db,
|
|
ensure_managed_files=ensure_managed_files,
|
|
)
|
|
|
|
|
|
async def call_retrieve(
|
|
harness: RetrieveHarness,
|
|
batch_id: str,
|
|
*,
|
|
provider: Optional[str] = None,
|
|
user: Optional[UserAPIKeyAuth] = None,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
query: Optional[Dict[str, str]] = None,
|
|
):
|
|
# Mirror the real flow: data starts as RetrieveBatchRequest(batch_id=...).
|
|
harness.data["data"] = {"batch_id": batch_id}
|
|
return await endpoints.retrieve_batch(
|
|
request=FakeRequest(headers=headers, query=query),
|
|
fastapi_response=Response(),
|
|
user_api_key_dict=user or UserAPIKeyAuth(api_key="sk-test"),
|
|
provider=provider,
|
|
batch_id=batch_id,
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# SCENARIO 1 - batch id encoded with model. litellm.aretrieve_batch via the
|
|
# model's resolved credentials; response ids re-encoded for the client.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__model_encoded_id(retrieve_harness):
|
|
resp = await call_retrieve(retrieve_harness, AZURE_BATCH_ID)
|
|
|
|
# 1. DISPATCH - model-credential path fired via litellm, router did not.
|
|
assert retrieve_harness.litellm_aretrieve.call_count == 1
|
|
retrieve_harness.router_aretrieve.assert_not_called()
|
|
|
|
# 2. CREDENTIALS - resolved for the model decoded FROM the batch id.
|
|
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
|
|
|
# 3. SEAM PAYLOAD - exact, whole dict forwarded to the provider call.
|
|
# Note `model` is the DECODED model, not the deployment from creds: the
|
|
# endpoint overrides it (provider-config providers like bedrock need it).
|
|
assert retrieve_harness.aretrieve_kwargs() == {
|
|
"custom_llm_provider": "azure",
|
|
"batch_id": "batch_orig123", # encoding stripped by this layer
|
|
"api_key": "sk-azure",
|
|
"api_base": "https://azure.test",
|
|
"model": "azure/gpt-4o",
|
|
}
|
|
|
|
# 4. OUTPUT SHAPE - ids re-encoded with the model for the round-trip.
|
|
assert resp.id == encode_file_id_with_model("batch-provider-id", "azure/gpt-4o", id_type="batch")
|
|
|
|
# write-back to the managed-object table happened, tagged as a retrieve.
|
|
assert retrieve_harness.update_batch_in_db.call_count == 1
|
|
assert retrieve_harness.update_batch_in_db.call_args.kwargs["operation"] == "retrieve"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__model_encoded_id__forwards_decoded_model_not_deployment(
|
|
retrieve_harness,
|
|
):
|
|
"""Regression guard for the line-483 override: the model forwarded to the
|
|
provider must be the decoded model id, never the deployment name that the
|
|
credential merge pulled in. Dropping the override silently 400s bedrock."""
|
|
await call_retrieve(retrieve_harness, AZURE_BATCH_ID)
|
|
|
|
assert retrieve_harness.aretrieve_kwargs()["model"] == "azure/gpt-4o"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__model_encoded_id__encodes_output_and_error_ids(
|
|
retrieve_harness,
|
|
):
|
|
retrieve_harness.litellm_aretrieve.return_value = make_batch(
|
|
id="batch-xyz",
|
|
output_file_id="file-out-raw",
|
|
error_file_id="file-err-raw",
|
|
)
|
|
|
|
resp = await call_retrieve(retrieve_harness, AZURE_BATCH_ID)
|
|
|
|
assert resp.output_file_id == encode_file_id_with_model("file-out-raw", "azure/gpt-4o")
|
|
assert resp.error_file_id == encode_file_id_with_model("file-err-raw", "azure/gpt-4o")
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__model_encoded_beats_loadbalancing(retrieve_harness):
|
|
"""Precedence: model-encoded id is checked before the loadbalancing/unified
|
|
elif, so it wins even with loadbalancing enabled."""
|
|
with patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True):
|
|
await call_retrieve(retrieve_harness, AZURE_BATCH_ID)
|
|
|
|
assert retrieve_harness.litellm_aretrieve.call_count == 1
|
|
retrieve_harness.router_aretrieve.assert_not_called()
|
|
retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Unified managed batch id -> llm_router.aretrieve_batch. model_id is parsed
|
|
# out of the unified id and stamped onto hidden params; raw file ids on the
|
|
# response are resolved back to unified ids.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__unified_batch_id_routes_to_router(retrieve_harness):
|
|
with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
|
|
resp = await call_retrieve(retrieve_harness, "batch-unified-blob")
|
|
|
|
# DISPATCH - router fired, direct litellm did not.
|
|
assert retrieve_harness.router_aretrieve.call_count == 1
|
|
retrieve_harness.litellm_aretrieve.assert_not_called()
|
|
|
|
# Credentials are resolved for the deployment behind the unified id so the batch's
|
|
# output file can be read for cost accounting. This id resolves to nothing here, and
|
|
# the retrieve must still serve the batch rather than fail on the lookup.
|
|
retrieve_harness.creds_resolver.assert_called_once_with(model_id="gpt-4o-mini")
|
|
|
|
# router receives the (still-encoded) batch id verbatim - this layer does
|
|
# not decode it for the unified path.
|
|
assert retrieve_harness.router_kwargs() == {"batch_id": "batch-unified-blob"}
|
|
|
|
# hidden params: unified id passed through, model_id parsed from it.
|
|
assert resp._hidden_params["unified_batch_id"] == UNIFIED_BATCH_ID
|
|
assert resp._hidden_params["model_id"] == "gpt-4o-mini"
|
|
|
|
# raw provider file ids on the response are normalized to managed unified ids.
|
|
retrieve_harness.ensure_managed_files.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__loadbalancing_raw_id_routes_to_router(retrieve_harness):
|
|
"""Loadbalancing on + a plain (non-encoded, non-unified) batch id routes to
|
|
the router. Locks the current dispatch contract of the shared elif."""
|
|
with patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True):
|
|
resp = await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
assert retrieve_harness.router_aretrieve.call_count == 1
|
|
retrieve_harness.litellm_aretrieve.assert_not_called()
|
|
assert retrieve_harness.router_kwargs() == {"batch_id": "batch-raw-xyz"}
|
|
# not a unified id -> hidden param reflects that, no model_id stamped.
|
|
assert resp._hidden_params["unified_batch_id"] is False
|
|
assert "model_id" not in resp._hidden_params
|
|
# not a unified id -> no managed-file normalization.
|
|
retrieve_harness.ensure_managed_files.assert_not_called()
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# SCENARIO 3 - fallback to custom_llm_provider (env-var creds). MUST NOT touch
|
|
# the credential resolver or the router.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__fallback_default_openai(retrieve_harness, openai_env_creds):
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
assert retrieve_harness.litellm_aretrieve.call_count == 1
|
|
retrieve_harness.router_aretrieve.assert_not_called()
|
|
retrieve_harness.creds_resolver.assert_not_called() # inverse-bug guard
|
|
assert retrieve_harness.aretrieve_kwargs() == {
|
|
"custom_llm_provider": "openai",
|
|
"batch_id": "batch-raw-xyz",
|
|
}
|
|
assert retrieve_harness.update_batch_in_db.call_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__fallback_no_creds_404(retrieve_harness, no_openai_creds):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
assert exc.value.code == "404"
|
|
assert exc.value.type == "invalid_request_error"
|
|
assert exc.value.param is None
|
|
assert exc.value.message == "No batch found with id 'batch-raw-xyz'."
|
|
retrieve_harness.litellm_aretrieve.assert_not_called()
|
|
retrieve_harness.router_aretrieve.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__fallback_explicit_provider_bypasses_not_found_gate(retrieve_harness, no_openai_creds):
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz", provider="anthropic")
|
|
|
|
assert retrieve_harness.aretrieve_kwargs()["custom_llm_provider"] == "anthropic"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__fallback_env_key_alone_forwards(retrieve_harness, monkeypatch):
|
|
monkeypatch.setattr(litellm, "api_key", None)
|
|
monkeypatch.setattr(litellm, "openai_key", None)
|
|
monkeypatch.setenv("OPENAI_API_KEY", "sk-env-openai")
|
|
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
assert retrieve_harness.aretrieve_kwargs() == {
|
|
"custom_llm_provider": "openai",
|
|
"batch_id": "batch-raw-xyz",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__fallback_provider_path_param(retrieve_harness):
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz", provider="anthropic")
|
|
|
|
retrieve_harness.creds_resolver.assert_not_called()
|
|
assert retrieve_harness.aretrieve_kwargs()["custom_llm_provider"] == "anthropic"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__fallback_provider_from_header(retrieve_harness):
|
|
retrieve_harness.provider_from_headers.return_value = "bedrock"
|
|
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
assert retrieve_harness.aretrieve_kwargs()["custom_llm_provider"] == "bedrock"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__fallback_provider_from_query(retrieve_harness):
|
|
retrieve_harness.provider_from_query.return_value = "vertex_ai"
|
|
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
assert retrieve_harness.aretrieve_kwargs()["custom_llm_provider"] == "vertex_ai"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__fallback_provider_precedence_path_over_header(
|
|
retrieve_harness,
|
|
):
|
|
"""provider path param beats the header-derived provider."""
|
|
retrieve_harness.provider_from_headers.return_value = "bedrock"
|
|
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz", provider="anthropic")
|
|
|
|
assert retrieve_harness.aretrieve_kwargs()["custom_llm_provider"] == "anthropic"
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# ManagedObjectTable short-circuit. A terminal-status DB row is returned
|
|
# immediately - no provider call, no write-back. A non-terminal row falls
|
|
# through to a provider sync.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("status", ["completed", "complete", "failed", "cancelled", "expired"])
|
|
async def test_retrieve__db_terminal_state_short_circuits(retrieve_harness, status):
|
|
# "complete" is the DB-normalized alias of "completed"; it is not a valid
|
|
# constructor literal but reaches the endpoint via a stored row, so set it
|
|
# post-construction to exercise that exact branch.
|
|
db_response = make_batch(id="batch-from-db", status="completed")
|
|
db_response.status = status
|
|
retrieve_harness.get_batch_from_db.return_value = (MagicMock(), db_response)
|
|
|
|
resp = await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
# No provider seam fired, and no write-back (the row is already terminal).
|
|
retrieve_harness.litellm_aretrieve.assert_not_called()
|
|
retrieve_harness.router_aretrieve.assert_not_called()
|
|
retrieve_harness.update_batch_in_db.assert_not_called()
|
|
# The DB object is what the client gets back.
|
|
assert resp is db_response
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__db_terminal_unified_resolves_file_ids(retrieve_harness):
|
|
db_response = make_batch(id="batch-from-db", status="completed")
|
|
db_batch_object = MagicMock()
|
|
retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response)
|
|
|
|
with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
|
|
await call_retrieve(retrieve_harness, "batch-unified-blob")
|
|
|
|
# Terminal short-circuit still registers/normalizes raw provider file ids.
|
|
retrieve_harness.ensure_managed_files.assert_called_once()
|
|
ensure_kwargs = retrieve_harness.ensure_managed_files.call_args.kwargs
|
|
assert ensure_kwargs["response"] is db_response
|
|
assert ensure_kwargs["db_batch_object"] is db_batch_object
|
|
assert ensure_kwargs["unified_batch_id"] == UNIFIED_BATCH_ID
|
|
retrieve_harness.litellm_aretrieve.assert_not_called()
|
|
retrieve_harness.router_aretrieve.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harness, openai_env_creds):
|
|
"""A non-terminal DB row must NOT short-circuit; the endpoint syncs with the
|
|
provider to refresh state."""
|
|
db_response = make_batch(id="batch-from-db", status="validating")
|
|
retrieve_harness.get_batch_from_db.return_value = (MagicMock(), db_response)
|
|
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
# Provider sync happened despite the DB hit.
|
|
assert retrieve_harness.litellm_aretrieve.call_count == 1
|
|
assert retrieve_harness.update_batch_in_db.call_count == 1
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Cross-cutting: enrichment route_type and failure-hook on provider error.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__uses_aretrieve_batch_route_type(retrieve_harness, openai_env_creds):
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
assert retrieve_harness.pre_call.call_args.kwargs["route_type"] == "aretrieve_batch"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__exception_calls_failure_hook(retrieve_harness, openai_env_creds):
|
|
retrieve_harness.litellm_aretrieve.side_effect = ValueError("provider boom")
|
|
|
|
with pytest.raises(ProxyException):
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
retrieve_harness.logging.post_call_failure_hook.assert_called_once()
|
|
assert (
|
|
retrieve_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0]
|
|
== "provider boom"
|
|
)
|
|
|
|
|
|
# =========================================================================== #
|
|
# #
|
|
# GET /v1/batches - list_batches routing-contract tests #
|
|
# #
|
|
# Branch order (first match wins): #
|
|
# 1. managed_files hook present -> managed_files_obj.list_user_batches #
|
|
# 2. model from body/query/header -> litellm.alist_batches + id encode #
|
|
# 3. target_model_names (param or body) -> llm_router.alist_batches #
|
|
# 4. fallback -> litellm.alist_batches via env-var custom_llm_provider #
|
|
# #
|
|
# llm_router is required; absence is a 500 before any branch runs. #
|
|
# =========================================================================== #
|
|
|
|
|
|
class FakeListPage:
|
|
"""Stand-in for the SyncCursorPage[Batch] that alist_batches returns. The
|
|
endpoint only touches `.data` (to encode ids) and `._hidden_params`."""
|
|
|
|
def __init__(self, data: Any):
|
|
self.data = data
|
|
self._hidden_params: Dict[str, Any] = {}
|
|
|
|
|
|
@dataclass
|
|
class ListHarness:
|
|
body: Dict[str, Any]
|
|
read_body: AsyncMock
|
|
pre_call: AsyncMock
|
|
get_headers: MagicMock
|
|
provider_from_headers: MagicMock
|
|
provider_from_query: MagicMock
|
|
litellm_alist: AsyncMock
|
|
router: MagicMock
|
|
logging: MagicMock
|
|
creds_resolver: MagicMock
|
|
|
|
@property
|
|
def router_alist(self) -> AsyncMock:
|
|
return self.router.alist_batches
|
|
|
|
def set_managed_files(self, page: Any) -> AsyncMock:
|
|
"""Install a managed_files hook exposing list_user_batches -> page."""
|
|
hook = MagicMock()
|
|
hook.list_user_batches = AsyncMock(return_value=page)
|
|
self.logging.get_proxy_hook = MagicMock(return_value=hook)
|
|
return hook.list_user_batches
|
|
|
|
def alist_kwargs(self) -> Dict[str, Any]:
|
|
assert self.litellm_alist.call_count == 1
|
|
return dict(self.litellm_alist.call_args.kwargs)
|
|
|
|
def router_kwargs(self) -> Dict[str, Any]:
|
|
assert self.router_alist.call_count == 1
|
|
return dict(self.router_alist.call_args.kwargs)
|
|
|
|
|
|
@pytest.fixture
|
|
def list_harness():
|
|
body_holder: Dict[str, Any] = {"body": {}}
|
|
logging = MagicMock(spec=ProxyLogging)
|
|
logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"])
|
|
logging.post_call_failure_hook = AsyncMock()
|
|
logging.update_request_status = AsyncMock()
|
|
# Default: no managed_files hook -> branches 2/3/4 are reachable.
|
|
logging.get_proxy_hook = MagicMock(return_value=None)
|
|
|
|
router = MagicMock(spec=Router)
|
|
router.alist_batches = AsyncMock(return_value=FakeListPage([]))
|
|
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
|
|
|
read_body = AsyncMock(side_effect=lambda request: body_holder["body"])
|
|
pre_call = AsyncMock(side_effect=lambda **kw: (body_holder["body"], MagicMock()))
|
|
get_headers = MagicMock(return_value={})
|
|
provider_from_headers = MagicMock(return_value=None)
|
|
provider_from_query = MagicMock(return_value=None)
|
|
litellm_alist = AsyncMock(return_value=FakeListPage([]))
|
|
|
|
with ExitStack() as stack:
|
|
stack.enter_context(patch.object(endpoints, "_read_request_body", read_body))
|
|
stack.enter_context(
|
|
patch.object(
|
|
ProxyBaseLLMRequestProcessing,
|
|
"common_processing_pre_call_logic",
|
|
pre_call,
|
|
)
|
|
)
|
|
stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
|
|
stack.enter_context(
|
|
patch.object(
|
|
endpoints,
|
|
"get_custom_llm_provider_from_request_headers",
|
|
provider_from_headers,
|
|
)
|
|
)
|
|
stack.enter_context(
|
|
patch.object(
|
|
endpoints,
|
|
"get_custom_llm_provider_from_request_query",
|
|
provider_from_query,
|
|
)
|
|
)
|
|
stack.enter_context(patch.object(litellm, "alist_batches", litellm_alist))
|
|
stack.enter_context(patch.object(proxy_server, "llm_router", router))
|
|
stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging))
|
|
stack.enter_context(patch.object(proxy_server, "general_settings", {}))
|
|
stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock()))
|
|
stack.enter_context(patch.object(proxy_server, "version", "test-version"))
|
|
|
|
yield ListHarness(
|
|
body=body_holder,
|
|
read_body=read_body,
|
|
pre_call=pre_call,
|
|
get_headers=get_headers,
|
|
provider_from_headers=provider_from_headers,
|
|
provider_from_query=provider_from_query,
|
|
litellm_alist=litellm_alist,
|
|
router=router,
|
|
logging=logging,
|
|
creds_resolver=router.get_deployment_credentials_with_provider,
|
|
)
|
|
|
|
|
|
async def call_list(
|
|
harness: ListHarness,
|
|
*,
|
|
provider: Optional[str] = None,
|
|
limit: Optional[int] = None,
|
|
after: Optional[str] = None,
|
|
target_model_names: Optional[str] = None,
|
|
user: Optional[UserAPIKeyAuth] = None,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
query: Optional[Dict[str, str]] = None,
|
|
body: Optional[Dict[str, Any]] = None,
|
|
):
|
|
harness.body["body"] = body if body is not None else {}
|
|
return await endpoints.list_batches(
|
|
request=FakeRequest(headers=headers, query=query),
|
|
fastapi_response=Response(),
|
|
provider=provider,
|
|
limit=limit,
|
|
after=after,
|
|
user_api_key_dict=user or UserAPIKeyAuth(api_key="sk-test"),
|
|
target_model_names=target_model_names,
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Branch 1 - ManagedObjectTable listing. This is the default production path
|
|
# (the managed_files hook is registered) and wins over every other branch.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__managed_files_path(list_harness):
|
|
page = FakeListPage([make_batch(id="batch-1")])
|
|
list_user_batches = list_harness.set_managed_files(page)
|
|
|
|
user = UserAPIKeyAuth(api_key="sk-test")
|
|
resp = await call_list(
|
|
list_harness,
|
|
user=user,
|
|
limit=7,
|
|
after="batch-cursor",
|
|
provider="openai",
|
|
target_model_names="m1,m2",
|
|
)
|
|
|
|
# DISPATCH - managed-files seam fired, neither provider seam did.
|
|
list_user_batches.assert_called_once_with(
|
|
user_api_key_dict=user,
|
|
limit=7,
|
|
after="batch-cursor",
|
|
provider="openai",
|
|
target_model_names="m1,m2",
|
|
llm_router=list_harness.router,
|
|
)
|
|
list_harness.litellm_alist.assert_not_called()
|
|
list_harness.router_alist.assert_not_called()
|
|
assert resp is page
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__managed_files_beats_model_param(list_harness):
|
|
"""Branch 1 is checked before the model branch: a model in the body does not
|
|
divert away from managed-files listing."""
|
|
page = FakeListPage([])
|
|
list_user_batches = list_harness.set_managed_files(page)
|
|
|
|
await call_list(list_harness, body={"model": "azure/gpt-4o"})
|
|
|
|
list_user_batches.assert_called_once()
|
|
list_harness.litellm_alist.assert_not_called()
|
|
list_harness.router_alist.assert_not_called()
|
|
list_harness.creds_resolver.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
"limit, expected_message, expected_openai_code",
|
|
[
|
|
(
|
|
-1,
|
|
"Invalid 'limit': integer below minimum value. Expected a value >= 0, but got -1 instead.",
|
|
"integer_below_min_value",
|
|
),
|
|
(
|
|
101,
|
|
"Invalid 'limit': integer above maximum value. Expected a value <= 100, but got 101 instead.",
|
|
"integer_above_max_value",
|
|
),
|
|
(
|
|
1000,
|
|
"Invalid 'limit': integer above maximum value. Expected a value <= 100, but got 1000 instead.",
|
|
"integer_above_max_value",
|
|
),
|
|
],
|
|
)
|
|
async def test_list__out_of_range_limit_rejected_with_400(list_harness, limit, expected_message, expected_openai_code):
|
|
"""OpenAI parity: GET /v1/batches rejects limit < 0 and limit > 100 with an
|
|
OpenAI-shaped 400 before any listing branch runs (issue #37149)."""
|
|
list_user_batches = list_harness.set_managed_files(FakeListPage([]))
|
|
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_list(list_harness, limit=limit)
|
|
|
|
assert exc.value.code == "400"
|
|
assert exc.value.param == "limit"
|
|
assert exc.value.type == "invalid_request_error"
|
|
assert exc.value.openai_code == expected_openai_code
|
|
assert exc.value.message == expected_message
|
|
list_user_batches.assert_not_called()
|
|
list_harness.litellm_alist.assert_not_called()
|
|
list_harness.router_alist.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("limit", [None, 0, 1, 100])
|
|
async def test_list__in_range_limit_dispatches(list_harness, limit):
|
|
"""OpenAI parity: live OpenAI accepts limit=0 (empty page) and 1..100, so
|
|
those values must keep flowing through to the listing branch untouched."""
|
|
page = FakeListPage([])
|
|
list_user_batches = list_harness.set_managed_files(page)
|
|
|
|
resp = await call_list(list_harness, limit=limit)
|
|
|
|
assert resp is page
|
|
assert list_user_batches.call_args.kwargs["limit"] == limit
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Branch 2 - model from body/query/header. The endpoint resolves credentials
|
|
# for the body model, forwards custom_llm_provider once (it pops it from data
|
|
# via prepare_data_with_credentials the way create/retrieve do), and encodes
|
|
# the response ids. Regression guard for the duplicate-kwarg
|
|
# "multiple values for keyword argument 'custom_llm_provider'" bug.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__model_from_body_routes_and_encodes(list_harness):
|
|
list_harness.litellm_alist.return_value = FakeListPage([make_batch(id="batch-1"), make_batch(id="batch-2")])
|
|
|
|
resp = await call_list(list_harness, body={"model": "azure/gpt-4o"})
|
|
|
|
assert list_harness.litellm_alist.call_count == 1
|
|
list_harness.router_alist.assert_not_called()
|
|
list_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
|
assert resp.data[0].id == encode_file_id_with_model("batch-1", "azure/gpt-4o", id_type="batch")
|
|
assert resp.data[1].id == encode_file_id_with_model("batch-2", "azure/gpt-4o", id_type="batch")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Branch 3 - target_model_names (function param or body) -> llm_router. Routes
|
|
# to the FIRST model in the comma list; `model` is stripped from data first.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__target_model_names_param_routes_to_router(list_harness):
|
|
await call_list(list_harness, target_model_names="m1,m2", limit=3, after="cur")
|
|
|
|
assert list_harness.router_alist.call_count == 1
|
|
list_harness.litellm_alist.assert_not_called()
|
|
list_harness.creds_resolver.assert_not_called()
|
|
# first model only; after/limit forwarded; nothing else (param not in data).
|
|
assert list_harness.router_kwargs() == {
|
|
"model": "m1",
|
|
"after": "cur",
|
|
"limit": 3,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__target_model_names_from_body(list_harness):
|
|
await call_list(list_harness, body={"target_model_names": "m1,m2"})
|
|
|
|
assert list_harness.router_alist.call_count == 1
|
|
list_harness.litellm_alist.assert_not_called()
|
|
kwargs = list_harness.router_kwargs()
|
|
assert kwargs["model"] == "m1"
|
|
# body-sourced target_model_names stays in the forwarded data.
|
|
assert kwargs["target_model_names"] == "m1,m2"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__target_model_names_takes_first_only(list_harness):
|
|
"""Locks the current behavior: with multiple target models, only the first
|
|
is routed to (silently, unlike create which 400s on >1). A change here -
|
|
intentional or not - must update this test."""
|
|
await call_list(list_harness, target_model_names="alpha,beta,gamma")
|
|
|
|
assert list_harness.router_kwargs()["model"] == "alpha"
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Branch 4 - fallback to custom_llm_provider (env-var creds). MUST NOT touch
|
|
# the credential resolver or the router.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__fallback_default_openai(list_harness, no_openai_creds):
|
|
"""list stays ungated by the no-creds 404 guard: it answers about a
|
|
collection, not a specific id, so there is nothing to 404 about."""
|
|
await call_list(list_harness)
|
|
|
|
assert list_harness.litellm_alist.call_count == 1
|
|
list_harness.router_alist.assert_not_called()
|
|
list_harness.creds_resolver.assert_not_called() # inverse-bug guard
|
|
assert list_harness.alist_kwargs() == {
|
|
"custom_llm_provider": "openai",
|
|
"after": None,
|
|
"limit": None,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__fallback_provider_path_param(list_harness):
|
|
await call_list(list_harness, provider="anthropic")
|
|
|
|
list_harness.creds_resolver.assert_not_called()
|
|
assert list_harness.alist_kwargs()["custom_llm_provider"] == "anthropic"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__fallback_provider_from_header(list_harness):
|
|
list_harness.provider_from_headers.return_value = "bedrock"
|
|
|
|
await call_list(list_harness)
|
|
|
|
assert list_harness.alist_kwargs()["custom_llm_provider"] == "bedrock"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__fallback_provider_from_query(list_harness):
|
|
list_harness.provider_from_query.return_value = "vertex_ai"
|
|
|
|
await call_list(list_harness)
|
|
|
|
assert list_harness.alist_kwargs()["custom_llm_provider"] == "vertex_ai"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__fallback_after_and_limit_forwarded(list_harness):
|
|
await call_list(list_harness, after="cursor-9", limit=42)
|
|
|
|
kwargs = list_harness.alist_kwargs()
|
|
assert kwargs["after"] == "cursor-9"
|
|
assert kwargs["limit"] == 42
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Cross-cutting: router requirement, route_type, failure hook.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__no_router_raises_500(list_harness):
|
|
with patch.object(proxy_server, "llm_router", None):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_list(list_harness)
|
|
|
|
assert exc.value.code == "500"
|
|
list_harness.litellm_alist.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__uses_alist_batches_route_type(list_harness):
|
|
await call_list(list_harness)
|
|
|
|
assert list_harness.pre_call.call_args.kwargs["route_type"] == "alist_batches"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__exception_calls_failure_hook(list_harness):
|
|
list_harness.litellm_alist.side_effect = ValueError("provider boom")
|
|
|
|
with pytest.raises(ProxyException):
|
|
await call_list(list_harness)
|
|
|
|
list_harness.logging.post_call_failure_hook.assert_called_once()
|
|
assert list_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__failure_hook_and_response_share_the_request_litellm_call_id(list_harness):
|
|
call_id = "lit7836-list-batches-call-id"
|
|
list_harness.pre_call.side_effect = lambda **kw: (
|
|
{**list_harness.body["body"], "litellm_call_id": call_id},
|
|
MagicMock(),
|
|
)
|
|
list_harness.litellm_alist.side_effect = ValueError("provider boom")
|
|
|
|
with pytest.raises(ProxyException) as raised:
|
|
await call_list(list_harness, after="batch-0", limit=5)
|
|
|
|
failure_request_data = list_harness.logging.post_call_failure_hook.call_args.kwargs["request_data"]
|
|
assert failure_request_data["litellm_call_id"] == call_id
|
|
assert (failure_request_data["after"], failure_request_data["limit"]) == ("batch-0", 5)
|
|
assert raised.value.headers["x-litellm-call-id"] == call_id
|
|
|
|
|
|
# =========================================================================== #
|
|
# #
|
|
# POST /v1/batches/{batch_id}/cancel - cancel_batch routing-contract tests #
|
|
# #
|
|
# Three branches, first match wins: #
|
|
# 1. model-encoded batch id -> litellm.acancel_batch via model creds #
|
|
# 2. unified batch id -> llm_router.acancel_batch (model+batch_id #
|
|
# parsed out of the unified id) #
|
|
# 3. fallback -> litellm.acancel_batch via env-var provider #
|
|
# Every branch then writes state back via update_batch_in_database( #
|
|
# operation="cancel"). There is NO ManagedObjectTable read short-circuit #
|
|
# here (unlike retrieve). #
|
|
# #
|
|
# These tests pin CURRENT behavior so a refactor can't silently change it. #
|
|
# Two current-behavior quirks are locked deliberately and noted inline: #
|
|
# - SCENARIO 1 forwards the DEPLOYMENT model from creds, not the decoded #
|
|
# model (retrieve overrides it; cancel does not). #
|
|
# - SCENARIO 3 rebuilds a CancelBatchRequest and forwards only #
|
|
# {custom_llm_provider, batch_id}, dropping enrichment keys. #
|
|
# =========================================================================== #
|
|
|
|
|
|
@dataclass
|
|
class CancelHarness:
|
|
data: Dict[str, Any]
|
|
pre_call: AsyncMock
|
|
add_data: AsyncMock
|
|
get_headers: MagicMock
|
|
provider_from_headers: MagicMock
|
|
provider_from_query: MagicMock
|
|
litellm_acancel: AsyncMock
|
|
router: MagicMock
|
|
logging: MagicMock
|
|
creds_resolver: MagicMock
|
|
update_batch_in_db: AsyncMock
|
|
|
|
@property
|
|
def router_acancel(self) -> AsyncMock:
|
|
return self.router.acancel_batch
|
|
|
|
def acancel_kwargs(self) -> Dict[str, Any]:
|
|
assert self.litellm_acancel.call_count == 1
|
|
return dict(self.litellm_acancel.call_args.kwargs)
|
|
|
|
def router_kwargs(self) -> Dict[str, Any]:
|
|
assert self.router_acancel.call_count == 1
|
|
return dict(self.router_acancel.call_args.kwargs)
|
|
|
|
|
|
@pytest.fixture
|
|
def cancel_harness():
|
|
data_holder: Dict[str, Any] = {"data": {}}
|
|
logging = MagicMock(spec=ProxyLogging)
|
|
logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"])
|
|
logging.post_call_failure_hook = AsyncMock()
|
|
logging.update_request_status = AsyncMock()
|
|
logging.get_proxy_hook = MagicMock(return_value=None)
|
|
|
|
router = MagicMock(spec=Router)
|
|
router.acancel_batch = AsyncMock(return_value=make_batch())
|
|
router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
|
|
|
|
pre_call = AsyncMock(side_effect=lambda **kw: (data_holder["data"], MagicMock()))
|
|
# add_litellm_data_to_request is a passthrough that returns the data it got.
|
|
add_data = AsyncMock(side_effect=lambda **kw: kw["data"])
|
|
get_headers = MagicMock(return_value={})
|
|
provider_from_headers = MagicMock(return_value=None)
|
|
provider_from_query = MagicMock(return_value=None)
|
|
litellm_acancel = AsyncMock(return_value=make_batch())
|
|
update_batch_in_db = AsyncMock(return_value=None)
|
|
|
|
with ExitStack() as stack:
|
|
stack.enter_context(
|
|
patch.object(
|
|
ProxyBaseLLMRequestProcessing,
|
|
"common_processing_pre_call_logic",
|
|
pre_call,
|
|
)
|
|
)
|
|
stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
|
|
stack.enter_context(
|
|
patch.object(
|
|
endpoints,
|
|
"get_custom_llm_provider_from_request_headers",
|
|
provider_from_headers,
|
|
)
|
|
)
|
|
stack.enter_context(
|
|
patch.object(
|
|
endpoints,
|
|
"get_custom_llm_provider_from_request_query",
|
|
provider_from_query,
|
|
)
|
|
)
|
|
stack.enter_context(patch.object(endpoints, "update_batch_in_database", update_batch_in_db))
|
|
stack.enter_context(patch.object(litellm, "acancel_batch", litellm_acancel))
|
|
stack.enter_context(patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False))
|
|
stack.enter_context(patch.object(proxy_server, "llm_router", router))
|
|
stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging))
|
|
stack.enter_context(patch.object(proxy_server, "general_settings", {}))
|
|
stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock()))
|
|
stack.enter_context(patch.object(proxy_server, "version", "test-version"))
|
|
stack.enter_context(patch.object(proxy_server, "prisma_client", MagicMock()))
|
|
stack.enter_context(patch.object(proxy_server, "add_litellm_data_to_request", add_data))
|
|
|
|
yield CancelHarness(
|
|
data=data_holder,
|
|
pre_call=pre_call,
|
|
add_data=add_data,
|
|
get_headers=get_headers,
|
|
provider_from_headers=provider_from_headers,
|
|
provider_from_query=provider_from_query,
|
|
litellm_acancel=litellm_acancel,
|
|
router=router,
|
|
logging=logging,
|
|
creds_resolver=router.get_deployment_credentials_with_provider,
|
|
update_batch_in_db=update_batch_in_db,
|
|
)
|
|
|
|
|
|
async def call_cancel(
|
|
harness: CancelHarness,
|
|
batch_id: str,
|
|
*,
|
|
provider: Optional[str] = None,
|
|
user: Optional[UserAPIKeyAuth] = None,
|
|
headers: Optional[Dict[str, str]] = None,
|
|
query: Optional[Dict[str, str]] = None,
|
|
data_extra: Optional[Dict[str, Any]] = None,
|
|
):
|
|
harness.data["data"] = {"batch_id": batch_id, **(data_extra or {})}
|
|
return await endpoints.cancel_batch(
|
|
request=FakeRequest(headers=headers, query=query),
|
|
batch_id=batch_id,
|
|
fastapi_response=Response(),
|
|
provider=provider,
|
|
user_api_key_dict=user or UserAPIKeyAuth(api_key="sk-test"),
|
|
)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# SCENARIO 1 - model-encoded batch id -> litellm.acancel_batch via model creds.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__model_encoded_id(cancel_harness):
|
|
resp = await call_cancel(cancel_harness, AZURE_BATCH_ID)
|
|
|
|
# DISPATCH - model-credential path via litellm; router untouched.
|
|
assert cancel_harness.litellm_acancel.call_count == 1
|
|
cancel_harness.router_acancel.assert_not_called()
|
|
|
|
# CREDENTIALS - resolved for the model decoded from the batch id.
|
|
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
|
|
|
# SEAM PAYLOAD - exact dict. NOTE current behavior: `model` is the
|
|
# DEPLOYMENT name from creds, NOT the decoded model (cancel, unlike
|
|
# retrieve, does not override it). Locking this guards the difference.
|
|
assert cancel_harness.acancel_kwargs() == {
|
|
"custom_llm_provider": "azure",
|
|
"batch_id": "batch_orig123", # decoded/stripped original id
|
|
"api_key": "sk-azure",
|
|
"api_base": "https://azure.test",
|
|
"model": "azure/gpt-4o-deployment",
|
|
}
|
|
|
|
# OUTPUT SHAPE - response id re-encoded with the DECODED model.
|
|
assert resp.id == encode_file_id_with_model("batch-provider-id", "azure/gpt-4o", id_type="batch")
|
|
|
|
# write-back tagged as a cancel.
|
|
assert cancel_harness.update_batch_in_db.call_count == 1
|
|
assert cancel_harness.update_batch_in_db.call_args.kwargs["operation"] == "cancel"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__model_encoded_id_forwards_deployment_model(cancel_harness):
|
|
"""Pin the current contract: cancel forwards the creds' deployment model.
|
|
If someone adds a decoded-model override (as retrieve has), this flips and
|
|
must be reviewed."""
|
|
await call_cancel(cancel_harness, AZURE_BATCH_ID)
|
|
|
|
assert cancel_harness.acancel_kwargs()["model"] == "azure/gpt-4o-deployment"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__model_encoded_beats_unified(cancel_harness):
|
|
with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
|
|
await call_cancel(cancel_harness, AZURE_BATCH_ID)
|
|
|
|
assert cancel_harness.litellm_acancel.call_count == 1
|
|
cancel_harness.router_acancel.assert_not_called()
|
|
cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# SCENARIO 2 - unified batch id -> llm_router.acancel_batch. model and batch_id
|
|
# are parsed out of the unified id; hidden params stamped.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__unified_batch_id_routes_to_router(cancel_harness):
|
|
with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
|
|
resp = await call_cancel(cancel_harness, "batch-unified-blob")
|
|
|
|
# DISPATCH - router fired, litellm did not, no creds lookup.
|
|
assert cancel_harness.router_acancel.call_count == 1
|
|
cancel_harness.litellm_acancel.assert_not_called()
|
|
cancel_harness.creds_resolver.assert_not_called()
|
|
|
|
# model + batch_id are extracted from the unified id and forwarded.
|
|
assert cancel_harness.router_kwargs() == {
|
|
"batch_id": "batch-raw-xyz",
|
|
"model": "gpt-4o-mini",
|
|
}
|
|
|
|
# hidden params: unified id passed through, model_id stamped from data.
|
|
assert resp._hidden_params["unified_batch_id"] == UNIFIED_BATCH_ID
|
|
assert resp._hidden_params["model_id"] == "gpt-4o-mini"
|
|
|
|
assert cancel_harness.update_batch_in_db.call_args.kwargs["operation"] == "cancel"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__db_write_receives_caller_auth(cancel_harness):
|
|
"""update_batch_in_database can only mint managed IDs for a cancelled batch's
|
|
output files when it has an auth context, so cancel must forward the caller's."""
|
|
caller = UserAPIKeyAuth(api_key="sk-test", user_id="user-cancel-1")
|
|
with patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
|
|
await call_cancel(cancel_harness, "batch-unified-blob", user=caller)
|
|
|
|
assert cancel_harness.update_batch_in_db.call_args.kwargs["user_api_key_dict"] is caller
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__unified_missing_model_id_400(cancel_harness):
|
|
# unified id with no model_id segment -> get_model_id returns None -> 400.
|
|
with patch.object(
|
|
endpoints,
|
|
"_is_base64_encoded_unified_file_id",
|
|
return_value="litellm_proxy;llm_batch_id:batch-xyz",
|
|
):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_cancel(cancel_harness, "batch-unified-blob")
|
|
|
|
assert exc.value.code == "400"
|
|
cancel_harness.router_acancel.assert_not_called()
|
|
cancel_harness.litellm_acancel.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__unified_no_router_500(cancel_harness):
|
|
with (
|
|
patch.object(proxy_server, "llm_router", None),
|
|
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID),
|
|
):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_cancel(cancel_harness, "batch-unified-blob")
|
|
|
|
assert exc.value.code == "500"
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# SCENARIO 3 - fallback to custom_llm_provider. Rebuilds a CancelBatchRequest
|
|
# and forwards only {custom_llm_provider, batch_id}.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__fallback_default_openai(cancel_harness, openai_env_creds):
|
|
await call_cancel(cancel_harness, "batch-raw-xyz")
|
|
|
|
assert cancel_harness.litellm_acancel.call_count == 1
|
|
cancel_harness.router_acancel.assert_not_called()
|
|
cancel_harness.creds_resolver.assert_not_called() # inverse-bug guard
|
|
# current behavior: enrichment keys dropped; only these two forwarded.
|
|
assert cancel_harness.acancel_kwargs() == {
|
|
"custom_llm_provider": "openai",
|
|
"batch_id": "batch-raw-xyz",
|
|
}
|
|
assert cancel_harness.update_batch_in_db.call_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__fallback_no_creds_404(cancel_harness, no_openai_creds):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_cancel(cancel_harness, "batch-raw-xyz")
|
|
|
|
assert exc.value.code == "404"
|
|
assert exc.value.type == "invalid_request_error"
|
|
assert exc.value.param is None
|
|
assert exc.value.message == "No batch found with id 'batch-raw-xyz'."
|
|
cancel_harness.litellm_acancel.assert_not_called()
|
|
cancel_harness.router_acancel.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__fallback_explicit_provider_bypasses_not_found_gate(cancel_harness, no_openai_creds):
|
|
await call_cancel(cancel_harness, "batch-raw-xyz", provider="anthropic")
|
|
|
|
assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "anthropic"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__fallback_env_key_alone_forwards(cancel_harness, monkeypatch):
|
|
monkeypatch.setattr(litellm, "api_key", None)
|
|
monkeypatch.setattr(litellm, "openai_key", None)
|
|
monkeypatch.setenv("OPENAI_API_KEY", "sk-env-openai")
|
|
|
|
await call_cancel(cancel_harness, "batch-raw-xyz")
|
|
|
|
assert cancel_harness.acancel_kwargs() == {
|
|
"custom_llm_provider": "openai",
|
|
"batch_id": "batch-raw-xyz",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__fallback_provider_path_param(cancel_harness):
|
|
await call_cancel(cancel_harness, "batch-raw-xyz", provider="anthropic")
|
|
|
|
cancel_harness.creds_resolver.assert_not_called()
|
|
assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "anthropic"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__fallback_provider_from_data_body(cancel_harness):
|
|
await call_cancel(cancel_harness, "batch-raw-xyz", data_extra={"custom_llm_provider": "bedrock"})
|
|
|
|
assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "bedrock"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__fallback_provider_from_header(cancel_harness):
|
|
cancel_harness.provider_from_headers.return_value = "vertex_ai"
|
|
|
|
await call_cancel(cancel_harness, "batch-raw-xyz")
|
|
|
|
assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "vertex_ai"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__fallback_provider_from_query(cancel_harness):
|
|
cancel_harness.provider_from_query.return_value = "azure"
|
|
|
|
await call_cancel(cancel_harness, "batch-raw-xyz")
|
|
|
|
assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "azure"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__fallback_provider_precedence_path_over_body(cancel_harness):
|
|
"""Intended contract: provider path param beats a body custom_llm_provider.
|
|
Regression guard: the body value is popped from data before the fallback
|
|
chain, so it never collides with the explicit kwarg."""
|
|
await call_cancel(
|
|
cancel_harness,
|
|
"batch-raw-xyz",
|
|
provider="anthropic",
|
|
data_extra={"custom_llm_provider": "bedrock"},
|
|
)
|
|
|
|
assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "anthropic"
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Cross-cutting: enrichment route_type and failure-hook on provider error.
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__uses_acancel_batch_route_type(cancel_harness, openai_env_creds):
|
|
await call_cancel(cancel_harness, "batch-raw-xyz")
|
|
|
|
assert cancel_harness.pre_call.call_args.kwargs["route_type"] == "acancel_batch"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__exception_calls_failure_hook(cancel_harness, openai_env_creds):
|
|
cancel_harness.litellm_acancel.side_effect = ValueError("provider boom")
|
|
|
|
with pytest.raises(ProxyException):
|
|
await call_cancel(cancel_harness, "batch-raw-xyz")
|
|
|
|
cancel_harness.logging.post_call_failure_hook.assert_called_once()
|
|
assert (
|
|
cancel_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom"
|
|
)
|
|
|
|
|
|
# =========================================================================== #
|
|
# Router-required 500 guards (one per endpoint branch that calls the router).
|
|
# These pin the defensive checks that fire when llm_router is unset.
|
|
# =========================================================================== #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__loadbalancing_no_router_500(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"model": "lb-model",
|
|
},
|
|
)
|
|
harness.is_known_model.return_value = True
|
|
with (
|
|
patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True),
|
|
patch.object(proxy_server, "llm_router", None),
|
|
):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_create(harness)
|
|
|
|
assert exc.value.code == "500"
|
|
harness.router_acreate.assert_not_called()
|
|
harness.litellm_acreate.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__unified_no_router_500(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "litellm_proxy_unified_id",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
with (
|
|
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
|
|
patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gpt-4o-mini"]),
|
|
patch.object(proxy_server, "llm_router", None),
|
|
):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_create(harness)
|
|
|
|
assert exc.value.code == "500"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__unified_no_router_500(retrieve_harness):
|
|
with (
|
|
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID),
|
|
patch.object(proxy_server, "llm_router", None),
|
|
):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_retrieve(retrieve_harness, "batch-unified-blob")
|
|
|
|
assert exc.value.code == "500"
|
|
retrieve_harness.router_aretrieve.assert_not_called()
|
|
retrieve_harness.litellm_aretrieve.assert_not_called()
|
|
|
|
|
|
# =========================================================================== #
|
|
# SCENARIO 3 + configured deployments: a provider-only call (custom-llm-provider
|
|
# header, no model anywhere) must resolve the gateway/team deployment's named
|
|
# credential for that provider and attach it to the provider call kwargs,
|
|
# instead of silently falling through to the host environment's default
|
|
# credentials (regression: vertex batch jobs landing in the hosting env's GCP
|
|
# project because litellm_credential_name never reached the call).
|
|
# =========================================================================== #
|
|
|
|
VERTEX_NAMED_CREDENTIAL = CredentialItem(
|
|
credential_name="vertex-named-cred",
|
|
credential_info={},
|
|
credential_values={
|
|
"vertex_project": "customer-project",
|
|
"vertex_location": "us-central1",
|
|
"vertex_credentials": "/creds/customer-sa.json",
|
|
},
|
|
)
|
|
|
|
|
|
def vertex_named_credential_router() -> Router:
|
|
return Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gemini-2.5-pro",
|
|
"litellm_params": {
|
|
"model": "vertex_ai/gemini-2.5-pro",
|
|
"litellm_credential_name": "vertex-named-cred",
|
|
},
|
|
}
|
|
]
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__provider_only_resolves_named_vertex_credentials(harness):
|
|
"""Provider-only create must attach the configured named credential, and must
|
|
NOT turn the call into a model-routed one (no model kwarg injected)."""
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
harness.provider_from_headers.return_value = "vertex_ai"
|
|
|
|
with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
|
|
with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
|
|
await call_create(harness)
|
|
|
|
assert harness.acreate_kwargs() == {
|
|
"custom_llm_provider": "vertex_ai",
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"metadata": None,
|
|
"vertex_project": "customer-project",
|
|
"vertex_location": "us-central1",
|
|
"vertex_credentials": "/creds/customer-sa.json",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__provider_only_ignores_other_provider_deployments(harness):
|
|
"""A provider-only vertex call must not pick up credentials from deployments
|
|
of a different provider; with no vertex deployment the payload is exactly the
|
|
pre-fix env-var fallback."""
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
harness.provider_from_headers.return_value = "vertex_ai"
|
|
openai_only_router = Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-4o",
|
|
"litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-global-openai"},
|
|
}
|
|
]
|
|
)
|
|
|
|
with patch.object(proxy_server, "llm_router", openai_only_router):
|
|
await call_create(harness)
|
|
|
|
assert harness.acreate_kwargs() == {
|
|
"custom_llm_provider": "vertex_ai",
|
|
"input_file_id": "file-plain",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
"metadata": None,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__provider_only_resolves_named_vertex_credentials(retrieve_harness):
|
|
retrieve_harness.provider_from_headers.return_value = "vertex_ai"
|
|
|
|
with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
|
|
with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
assert retrieve_harness.aretrieve_kwargs() == {
|
|
"custom_llm_provider": "vertex_ai",
|
|
"batch_id": "batch-raw-xyz",
|
|
"vertex_project": "customer-project",
|
|
"vertex_location": "us-central1",
|
|
"vertex_credentials": "/creds/customer-sa.json",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_list__provider_only_resolves_named_vertex_credentials(list_harness):
|
|
list_harness.provider_from_headers.return_value = "vertex_ai"
|
|
|
|
with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
|
|
with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
|
|
await call_list(list_harness)
|
|
|
|
assert list_harness.alist_kwargs() == {
|
|
"custom_llm_provider": "vertex_ai",
|
|
"after": None,
|
|
"limit": None,
|
|
"vertex_project": "customer-project",
|
|
"vertex_location": "us-central1",
|
|
"vertex_credentials": "/creds/customer-sa.json",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__provider_only_resolves_named_vertex_credentials(cancel_harness):
|
|
cancel_harness.provider_from_headers.return_value = "vertex_ai"
|
|
|
|
with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
|
|
with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
|
|
await call_cancel(cancel_harness, "batch-raw-xyz")
|
|
|
|
assert cancel_harness.acancel_kwargs() == {
|
|
"custom_llm_provider": "vertex_ai",
|
|
"batch_id": "batch-raw-xyz",
|
|
"vertex_project": "customer-project",
|
|
"vertex_location": "us-central1",
|
|
"vertex_credentials": "/creds/customer-sa.json",
|
|
}
|
|
|
|
|
|
# =========================================================================== #
|
|
# require_managed_files - raw provider ids must not reach the provider. #
|
|
# #
|
|
# Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to #
|
|
# these routes is forwarded under the shared provider credentials with no #
|
|
# tenant check, so any caller who learns another tenant's id can read its #
|
|
# batch, reuse its file as batch input, or cancel its job. These lock the #
|
|
# guard on every batches route that accepts a caller-supplied id. #
|
|
# =========================================================================== #
|
|
|
|
|
|
def _unified_batch_id(model_id: str = "azure/gpt-4o", batch_id: str = "batch-provider-id") -> str:
|
|
import base64
|
|
|
|
from litellm.types.utils import SpecialEnums
|
|
|
|
unified = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id)
|
|
return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")
|
|
|
|
|
|
def _unified_file_id() -> str:
|
|
import base64
|
|
|
|
from litellm.types.utils import SpecialEnums
|
|
|
|
unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
|
|
"application/json", "managed-id", "gpt-4o-mini", "file-provider-id", "gpt-4o-mini-id"
|
|
)
|
|
return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ManagedResourceAccessCheckerStub:
|
|
file_access: bool = True
|
|
object_access: bool = True
|
|
|
|
async def can_user_call_unified_file_id(
|
|
self,
|
|
unified_file_id: str,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
) -> bool:
|
|
return self.file_access
|
|
|
|
async def can_user_call_unified_object_id(
|
|
self,
|
|
unified_object_id: str,
|
|
user_api_key_dict: UserAPIKeyAuth,
|
|
) -> bool:
|
|
return self.object_access
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__raw_input_file_id_rejected_when_managed_files_required(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-victim-abc123",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
with patch.object(litellm, "require_managed_files", True):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_create(harness)
|
|
|
|
assert exc.value.code == "400"
|
|
harness.litellm_acreate.assert_not_called()
|
|
harness.router_acreate.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__model_encoded_input_file_id_rejected_when_managed_files_required(harness):
|
|
"""A model-encoded id is client-forgeable and has no ownership row, so it is
|
|
not a managed file id and must be rejected like any other raw id."""
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": AZURE_FILE_ID,
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
with patch.object(litellm, "require_managed_files", True):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_create(harness)
|
|
|
|
assert exc.value.code == "400"
|
|
harness.litellm_acreate.assert_not_called()
|
|
harness.router_acreate.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__raw_input_file_id_allowed_when_managed_files_not_required(harness, openai_env_creds):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": "file-victim-abc123",
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
|
|
with patch.object(litellm, "require_managed_files", False):
|
|
await call_create(harness)
|
|
|
|
assert harness.acreate_kwargs()["input_file_id"] == "file-victim-abc123"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_create__other_teams_unified_input_file_id_rejected(harness):
|
|
set_body(
|
|
harness,
|
|
{
|
|
"input_file_id": _unified_file_id(),
|
|
"endpoint": "/v1/chat/completions",
|
|
"completion_window": "24h",
|
|
},
|
|
)
|
|
harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False)
|
|
|
|
with patch.object(litellm, "require_managed_files", True):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_create(harness)
|
|
|
|
assert exc.value.code == "403"
|
|
harness.litellm_acreate.assert_not_called()
|
|
harness.router_acreate.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__raw_batch_id_rejected_when_managed_files_required(retrieve_harness):
|
|
with patch.object(litellm, "require_managed_files", True):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_retrieve(retrieve_harness, "batch-victim-abc123")
|
|
|
|
assert exc.value.code == "400"
|
|
retrieve_harness.litellm_aretrieve.assert_not_called()
|
|
retrieve_harness.router_aretrieve.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__unified_batch_id_allowed_when_managed_files_required(retrieve_harness):
|
|
retrieve_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub()
|
|
|
|
with patch.object(litellm, "require_managed_files", True):
|
|
await call_retrieve(retrieve_harness, _unified_batch_id())
|
|
|
|
assert retrieve_harness.router_aretrieve.call_count == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__raw_batch_id_rejected_when_managed_files_required(cancel_harness):
|
|
with patch.object(litellm, "require_managed_files", True):
|
|
with pytest.raises(ProxyException) as exc:
|
|
await call_cancel(cancel_harness, "batch-victim-abc123")
|
|
|
|
assert exc.value.code == "400"
|
|
cancel_harness.litellm_acancel.assert_not_called()
|
|
cancel_harness.router_acancel.assert_not_called()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cancel__unified_batch_id_allowed_when_managed_files_required(cancel_harness):
|
|
cancel_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub()
|
|
|
|
with patch.object(litellm, "require_managed_files", True):
|
|
await call_cancel(cancel_harness, _unified_batch_id())
|
|
|
|
assert cancel_harness.router_acancel.call_count == 1
|
|
|
|
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__managed_batch_defers_cost_to_the_poller_when_it_is_running(retrieve_harness):
|
|
with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)):
|
|
await call_retrieve(retrieve_harness, _unified_batch_id())
|
|
|
|
assert retrieve_harness.router.aretrieve_batch.await_count == 1
|
|
metadata = retrieve_harness.router.aretrieve_batch.await_args.kwargs.get("litellm_metadata") or {}
|
|
assert metadata.get("batch_ignore_default_logging") is True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__managed_batch_still_accounts_inline_without_a_poller(retrieve_harness):
|
|
with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=False)):
|
|
await call_retrieve(retrieve_harness, _unified_batch_id())
|
|
|
|
assert retrieve_harness.router.aretrieve_batch.await_count == 1
|
|
metadata = retrieve_harness.router.aretrieve_batch.await_args.kwargs.get("litellm_metadata") or {}
|
|
assert metadata.get("batch_ignore_default_logging") is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retrieve__raw_batch_id_is_untouched_by_the_poller_handoff(retrieve_harness, openai_env_creds):
|
|
with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)):
|
|
await call_retrieve(retrieve_harness, "batch-raw-xyz")
|
|
|
|
metadata = retrieve_harness.litellm_aretrieve.await_args.kwargs.get("litellm_metadata") or {}
|
|
assert metadata.get("batch_ignore_default_logging") is None
|