scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-19 00:01:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

fix(proxy/batches): enforce ownership and correct lookup key when resolving managed input_file_id

Browse source 
The adopted resolution queried LiteLLM_ManagedFileTable with the decoded
litellm_proxy string, but the unified_file_id column stores the raw base64
file id (see schema.prisma and the enterprise managed-files hook), so the
lookup never matched in production and silently fell back to the opaque id.
Query with the raw id instead and lock the key with a regression test.

Move the resolution above the dispatch branches so the load-balanced router
path receives the resolved storage_url too, enforce managed-file ownership
with the same can_access_resource semantics the files retrieve and download
endpoints use (404 on denial), and downgrade database failures to a logged
fallback instead of aborting batch creation. Unresolved ids still dispatch
unchanged because the managed-files deployment hook can map them via
model_file_id_mapping
This commit is contained in:
Yucheng Zhu 2026-07-23 23:50:22 -07:00
parent 411b525a83
commit f8fc629286
2 changed files with 246 additions and 222 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

68
litellm/proxy/batches_endpoints/endpoints.py
View file

@ -12,6 +12,7 @@ from fastapi import APIRouter, Depends, HTTPException, Path, Request, Response
import litellm
from litellm._logging import verbose_proxy_logger
from litellm.batches.main import CancelBatchRequest, RetrieveBatchRequest
from litellm.llms.base_llm.managed_resources.isolation import can_access_resource
from litellm.proxy._types import *
from litellm.proxy.common_utils.callback_utils import sanitize_openai_provider_metadata
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
@ -38,11 +39,47 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
update_batch_in_database,
)
from litellm.proxy.utils import handle_exception_on_proxy, is_known_model
from litellm.repositories.table_repositories import ManagedFileRepository
from litellm.types.llms.openai import LiteLLMBatchCreateRequest
router = APIRouter()
async def _resolve_managed_input_file_storage_url(
input_file_id: str,
user_api_key_dict: UserAPIKeyAuth,
) -> "str | None":
"""Resolve a managed (unified) input_file_id to its backend storage_url.
Returns None when the proxy has no database, the lookup fails, no managed
file row exists, or the row has no storage_url; callers fall back to
dispatching the original id so the managed-files deployment hook can still
map it via model_file_id_mapping. Raises a 404 when the caller does not
own the managed file.
"""
from litellm.proxy.proxy_server import prisma_client
if prisma_client is None:
return None
try:
db_file = await ManagedFileRepository(prisma_client).table.find_first(where={"unified_file_id": input_file_id})
except Exception as e:
verbose_proxy_logger.warning("create_batch: managed file lookup failed for %s: %s", input_file_id, e)
return None
if db_file is None:
return None
if not can_access_resource(
user_api_key_dict=user_api_key_dict,
created_by=db_file.created_by,
resource_team_id=db_file.team_id,
):
raise HTTPException(
status_code=404,
detail={"error": f"File not found: {input_file_id}"},
)
return db_file.storage_url or None
@router.post(
"/{provider}/v1/batches",
dependencies=[Depends(user_api_key_auth)],
@ -156,6 +193,14 @@ async def create_batch(
model_from_file_id = decode_model_from_file_id(input_file_id)
unified_file_id = _is_base64_encoded_unified_file_id(input_file_id)
if model_from_file_id is None and unified_file_id and input_file_id:
resolved_storage_url = await _resolve_managed_input_file_storage_url(
input_file_id=input_file_id,
user_api_key_dict=user_api_key_dict,
)
if resolved_storage_url is not None:
_create_batch_data["input_file_id"] = resolved_storage_url
# SCENARIO 1: File ID is encoded with model info
if model_from_file_id is not None and input_file_id:
credentials = get_credentials_for_model(
@ -230,29 +275,6 @@ async def create_batch(
detail={"error": "LLM Router not initialized. Ensure models added to proxy."},
)
# The base64-encoded unified_file_id is an opaque LiteLLM-internal
# token (unified_id + target_model_names), not a real
# provider-side file reference. Provider-specific handlers (e.g.
# Vertex AI's batch transformation, which parses a `publishers/`
# segment out of the file URI) require the actual backend storage
# location. Resolve it from LiteLLM_ManagedFileTable before
# dispatching, mirroring the same lookup already used by the
# files retrieve/download endpoints for managed files.
from litellm.proxy.proxy_server import prisma_client
if prisma_client is not None:
import inspect
from litellm.repositories.table_repositories import (
ManagedFileRepository,
)
db_file_res = ManagedFileRepository(prisma_client).table.find_first(
where={"unified_file_id": unified_file_id}
)
db_file = await db_file_res if inspect.isawaitable(db_file_res) else db_file_res
if db_file is not None and getattr(db_file, "storage_url", None):
_create_batch_data["input_file_id"] = db_file.storage_url
response = await llm_router.acreate_batch(**_create_batch_data)
response.input_file_id = input_file_id
response._hidden_params["unified_file_id"] = unified_file_id

400
tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
View file

@ -76,9 +76,7 @@ CREDS: Dict[str, Dict[str, str]] = {
}
# 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"
)
AZURE_FILE_ID = encode_file_id_with_model("file-original123", "azure/gpt-4o", id_type="file")
def make_batch(
@ -166,9 +164,7 @@ def harness():
router = MagicMock(spec=Router)
router.acreate_batch = AsyncMock(return_value=make_batch())
router.get_deployment_credentials_with_provider = MagicMock(
side_effect=_creds_lookup
)
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()))
@ -186,11 +182,7 @@ def harness():
pre_call,
)
)
stack.enter_context(
patch.object(
ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers
)
)
stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
stack.enter_context(
patch.object(
endpoints,
@ -200,9 +192,7 @@ def harness():
)
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(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", {}))
@ -283,9 +273,7 @@ async def test_create__model_encoded_file_id(harness):
}
# 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.id == encode_file_id_with_model("batch-provider-id", "azure/gpt-4o", id_type="batch")
assert resp.input_file_id == AZURE_FILE_ID
@ -307,12 +295,8 @@ async def test_create__model_encoded_file_id__encodes_output_and_error_ids(harne
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"
)
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
@ -358,9 +342,7 @@ async def test_create__model_from_body(harness):
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"
)
assert resp.id == encode_file_id_with_model("batch-provider-id", "vertex-model", id_type="batch")
@pytest.mark.asyncio
@ -495,10 +477,9 @@ async def test_create__unified_file_id_single_model(harness):
"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"]
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)
@ -522,10 +503,9 @@ async def test_create__unified_file_id_not_exactly_one_model_400(harness, models
"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 (
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)
@ -541,8 +521,13 @@ async def test_create__unified_file_id_resolves_real_storage_url(harness):
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 - the
unified id itself must never reach the router/provider call."""
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,
{
@ -553,43 +538,149 @@ async def test_create__unified_file_id_resolves_real_storage_url(harness):
)
fake_db_file = MagicMock(
storage_url="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.0/abc"
storage_url="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.0/abc",
created_by="user-1",
team_id=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)
fake_prisma_client = MagicMock()
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", fake_prisma_client
), patch(
"litellm.repositories.table_repositories.ManagedFileRepository",
fake_repo_cls,
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)
resp = await call_create(harness, user=UserAPIKeyAuth(api_key="sk-test", user_id="user-1"))
# The real storage_url - not the opaque unified id - must be what's
# forwarded to the router/provider.
assert harness.router_kwargs()["input_file_id"] == fake_db_file.storage_url
find_first.assert_awaited_once_with(where={"unified_file_id": "unified-xyz"})
# The lookup key is the RAW base64 id from the request, not the decoded string.
find_first.assert_awaited_once_with(where={"unified_file_id": "litellm_proxy_unified_id"})
# The unified id is still what's returned to the client.
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_other_tenant_gets_404(harness):
"""Ownership gate: a caller who does not own the managed file (different
user, different team, not an admin) must get a 404, and nothing may be
dispatched. Uses the real can_access_resource so the semantics cannot
drift from the files retrieve/download endpoints."""
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",
created_by="owner-user",
team_id="owner-team",
)
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),
):
with pytest.raises(ProxyException) as exc:
await call_create(harness, user=UserAPIKeyAuth(api_key="sk-test", user_id="intruder"))
assert exc.value.code == "404"
harness.router_acreate.assert_not_called()
harness.litellm_acreate.assert_not_called()
@pytest.mark.asyncio
async def test_create__unified_file_id_db_error_falls_back_to_raw_id(harness):
"""A database failure during resolution must not abort batch creation:
the original id is dispatched so the managed-files deployment hook can
still map it via model_file_id_mapping."""
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__unified_file_id_loadbalanced_path_gets_resolved_id(harness):
"""Resolution runs before dispatch branching, so the load-balanced router
branch also receives the resolved storage_url instead of the opaque
unified id."""
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
fake_db_file = MagicMock(
storage_url="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.0/abc",
created_by="user-1",
team_id=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(litellm, "enable_loadbalancing_on_batch_endpoints", True),
patch.object(endpoints, "_is_base64_encoded_unified_file_id", return_value="unified-xyz"),
patch.object(proxy_server, "prisma_client", MagicMock()),
patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
):
await call_create(harness, user=UserAPIKeyAuth(api_key="sk-test", user_id="user-1"))
assert harness.router_acreate.call_count == 1
assert harness.router_kwargs()["input_file_id"] == fake_db_file.storage_url
harness.litellm_acreate.assert_not_called()
@pytest.mark.asyncio
async def test_create__unified_file_id_no_managed_file_record_falls_back_to_raw_id(
harness,
):
"""If there's no LiteLLM_ManagedFileTable row (or it has no storage_url),
fall back to the previous behavior instead of raising - callers/providers
that don't need the resolved path (or legacy data) keep working."""
fall back to dispatching the original id instead of raising: the
managed-files deployment hook can still map it via model_file_id_mapping,
and legacy rows without storage_url keep working."""
set_body(
harness,
{
@ -604,15 +695,11 @@ async def test_create__unified_file_id_no_managed_file_record_falls_back_to_raw_
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(
"litellm.repositories.table_repositories.ManagedFileRepository",
fake_repo_cls,
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)
@ -631,10 +718,9 @@ async def test_create__model_encoded_beats_unified(harness):
"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"]
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)
@ -663,9 +749,7 @@ async def test_create__loadbalancing_routes_to_router(harness):
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
)
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()
@ -714,9 +798,7 @@ async def test_create__team_expiry_injected(harness):
},
)
await call_create(
harness, user=_user_with_expiry({"anchor": "created_at", "seconds": 3600})
)
await call_create(harness, user=_user_with_expiry({"anchor": "created_at", "seconds": 3600}))
assert harness.acreate_kwargs()["output_expires_after"] == {
"anchor": "created_at",
@ -822,12 +904,7 @@ async def test_create__exception_calls_failure_hook(harness):
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"
)
assert harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom"
# =========================================================================== #
@ -854,9 +931,7 @@ async def test_create__exception_calls_failure_hook(harness):
# 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"
)
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.
@ -907,9 +982,7 @@ def retrieve_harness():
router = MagicMock(spec=Router)
router.aretrieve_batch = AsyncMock(return_value=make_batch())
router.get_deployment_credentials_with_provider = MagicMock(
side_effect=_creds_lookup
)
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={})
@ -930,11 +1003,7 @@ def retrieve_harness():
pre_call,
)
)
stack.enter_context(
patch.object(
ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers
)
)
stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
stack.enter_context(
patch.object(
endpoints,
@ -949,24 +1018,12 @@ def retrieve_harness():
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, "resolve_input_file_id_to_unified", resolve_input)
)
stack.enter_context(
patch.object(
endpoints, "resolve_output_file_ids_to_unified", resolve_output
)
)
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, "resolve_input_file_id_to_unified", resolve_input))
stack.enter_context(patch.object(endpoints, "resolve_output_file_ids_to_unified", resolve_output))
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(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", {}))
@ -1040,9 +1097,7 @@ async def test_retrieve__model_encoded_id(retrieve_harness):
}
# 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"
)
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
@ -1073,12 +1128,8 @@ async def test_retrieve__model_encoded_id__encodes_output_and_error_ids(
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"
)
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
@ -1102,9 +1153,7 @@ async def test_retrieve__model_encoded_beats_loadbalancing(retrieve_harness):
@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
):
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.
@ -1209,9 +1258,7 @@ async def test_retrieve__fallback_provider_precedence_path_over_header(
@pytest.mark.asyncio
@pytest.mark.parametrize(
"status", ["completed", "complete", "failed", "cancelled", "expired"]
)
@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
@ -1235,9 +1282,7 @@ async def test_retrieve__db_terminal_unified_resolves_file_ids(retrieve_harness)
db_response = make_batch(id="batch-from-db", status="completed")
retrieve_harness.get_batch_from_db.return_value = (MagicMock(), db_response)
with patch.object(
endpoints, "_is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID
):
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 resolves raw provider file ids to unified.
@ -1270,9 +1315,7 @@ async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harn
async def test_retrieve__uses_aretrieve_batch_route_type(retrieve_harness):
await call_retrieve(retrieve_harness, "batch-raw-xyz")
assert (
retrieve_harness.pre_call.call_args.kwargs["route_type"] == "aretrieve_batch"
)
assert retrieve_harness.pre_call.call_args.kwargs["route_type"] == "aretrieve_batch"
@pytest.mark.asyncio
@ -1284,9 +1327,7 @@ async def test_retrieve__exception_calls_failure_hook(retrieve_harness):
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]
retrieve_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0]
== "provider boom"
)
@ -1359,9 +1400,7 @@ def list_harness():
router = MagicMock(spec=Router)
router.alist_batches = AsyncMock(return_value=FakeListPage([]))
router.get_deployment_credentials_with_provider = MagicMock(
side_effect=_creds_lookup
)
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()))
@ -1379,11 +1418,7 @@ def list_harness():
pre_call,
)
)
stack.enter_context(
patch.object(
ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers
)
)
stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
stack.enter_context(
patch.object(
endpoints,
@ -1516,21 +1551,15 @@ async def test_list__managed_files_beats_model_param(list_harness):
)
@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")]
)
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"
)
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")
# --------------------------------------------------------------------------- #
@ -1661,12 +1690,7 @@ async def test_list__exception_calls_failure_hook(list_harness):
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"
)
assert list_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom"
# =========================================================================== #
@ -1729,9 +1753,7 @@ def cancel_harness():
router = MagicMock(spec=Router)
router.acancel_batch = AsyncMock(return_value=make_batch())
router.get_deployment_credentials_with_provider = MagicMock(
side_effect=_creds_lookup
)
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.
@ -1750,11 +1772,7 @@ def cancel_harness():
pre_call,
)
)
stack.enter_context(
patch.object(
ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers
)
)
stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
stack.enter_context(
patch.object(
endpoints,
@ -1769,22 +1787,16 @@ def cancel_harness():
provider_from_query,
)
)
stack.enter_context(
patch.object(endpoints, "update_batch_in_database", update_batch_in_db)
)
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(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)
)
stack.enter_context(patch.object(proxy_server, "add_litellm_data_to_request", add_data))
yield CancelHarness(
data=data_holder,
@ -1849,9 +1861,7 @@ async def test_cancel__model_encoded_id(cancel_harness):
}
# 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"
)
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
@ -1870,9 +1880,7 @@ async def test_cancel__model_encoded_id_forwards_deployment_model(cancel_harness
@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
):
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
@ -1888,9 +1896,7 @@ async def test_cancel__model_encoded_beats_unified(cancel_harness):
@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
):
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.
@ -1929,8 +1935,9 @@ async def test_cancel__unified_missing_model_id_400(cancel_harness):
@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 (
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")
@ -1969,9 +1976,7 @@ async def test_cancel__fallback_provider_path_param(cancel_harness):
@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"}
)
await call_cancel(cancel_harness, "batch-raw-xyz", data_extra={"custom_llm_provider": "bedrock"})
assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "bedrock"
@ -2038,10 +2043,7 @@ async def test_cancel__exception_calls_failure_hook(cancel_harness):
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"
cancel_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom"
)
@ -2063,9 +2065,10 @@ async def test_create__loadbalancing_no_router_500(harness):
},
)
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 (
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)
@ -2084,12 +2087,10 @@ async def test_create__unified_no_router_500(harness):
"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 (
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)
@ -2099,9 +2100,10 @@ async def test_create__unified_no_router_500(harness):
@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 (
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")