scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(proxy/batches): resolve managed unified input_file_id to storage_url with ownership check before dispatch (#34474)

Browse source 
* fix: resolve unified_file_id to real storage_url before dispatching batch create

litellm.create_batch() against a Vertex AI-backed model crashes with an
opaque error when the input file was uploaded as a LiteLLM-managed
'unified file' (multi-model file upload). The base64-encoded
unified_file_id token is a LiteLLM-internal identifier, not a real
provider-side file reference, but the batches_endpoints create_batch
handler forwards it unchanged to llm_router.acreate_batch() /
litellm.acreate_batch() for the unified_file_id branch. Provider-specific
code that expects a real file location (e.g. Vertex AI's batch
transformation, which parses a 'publishers/' segment out of the GCS URI)
then fails on the opaque token.

Resolve the unified_file_id to its real backend location
(LiteLLM_ManagedFileTable.storage_url) before dispatch, mirroring the
same lookup already used by the files retrieve/download endpoints for
managed files. Falls back to the previous (unchanged) behavior if no
managed-file record exists.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* fix(proxy/batches): null-guard await on find_first for sync MagicMock test harnesses

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

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

* fix(proxy/batches): fail closed when the managed file ownership lookup errors

A lookup exception previously fell back to dispatching the original
unified id with the ownership gate unexecuted; the managed-files
deployment hook maps unified ids from cache without re-checking
ownership, so a database outage let a caller dispatch another tenant's
file. Raise a clear 503 instead and lock the behavior with a regression
test. No-database and no-row cases still fall back unchanged

* test(proxy/batches): default harness prisma_client to None

The batch routing harness left proxy_server.prisma_client at its module
global, which a sibling test in the same shard can leave as a MagicMock.
The unified-file rows that do not opt into managed-file resolution then
entered the resolver and awaited a non-awaitable mock, surfacing as a
503. Patch prisma_client to None by default so those rows stay a no-op;
resolution tests still override it explicitly

* fix(proxy/batches): keep unified resolution in its own branch and fail closed on missing row

Cursor flagged that hoisting the storage_url substitution above the
load-balanced dispatch branch broke two things on that path: the
model_file_id_mapping deployment filter keys on the original unified id,
and the response returned the internal storage_url instead of the
unified id. Move the resolution back inside the unified branch and
exclude unified ids from the load-balanced branch so a managed file
always takes the resolving path (which restores input_file_id and the
unified_file_id hidden param on the response), and a load-balanced batch
keeps the original id for deployment filtering.

Also fail closed with a 404 when a unified id has no managed-file row
while a database is present: the id cannot be ownership-verified, and
dispatching it would both bypass the gate and hit the Vertex
publishers-segment IndexError. Owned rows without a storage_url (legacy)
still dispatch the original id

* fix(proxy/batches): do not divert unified files off the load-balanced branch

Excluding unified ids from the load-balanced branch (and not
unified_file_id) regressed a path that works on the base revision: a
multi-model managed file dispatched with an explicit router model under
load balancing was routed into the unified branch, which raises a 400
for anything other than exactly one target model. Verified live against
base (200, managed-files deployment hook remaps the unified id per
model) versus the guarded branch (400 Expected 1 model, got 2).

Restore the original three-condition load-balanced branch so that path
keeps working unchanged. Unified-file storage_url resolution and the
ownership 404 still apply on the non-load-balanced unified branch, which
is the common managed-batch flow; the load-balanced managed path retains
its existing behavior and its pre-existing enterprise-hook ownership gap,
unchanged from base

* refactor(proxy/batches): scope managed-file handling to resolution, drop ownership check

Narrow this PR to its one problem: resolving a managed unified input_file_id
to its backend storage_url so provider batch handlers (Vertex parses a
publishers/ segment) receive a real location instead of the opaque token,
and failing closed with a 404 when the token has no backing row so it is
never dispatched into the provider crash.

Remove the cross-tenant ownership check (can_access_resource) added earlier.
Batch-create had no ownership enforcement before this PR, and the gap spans
every managed-file call type, so it belongs in the enterprise managed-files
pre-call hook (its acreate_batch branch) where files, batches and
fine-tuning are covered uniformly, not partially in this one endpoint. Filed
as a follow-up. This also removes the load-balanced-path ownership
inconsistency the bots flagged, since there is no ownership branch to skip.

Drop the inline comments flagged against the no-comments rule; behavior is
documented in the helper docstring and the test docstrings

* fix(proxy/batches): fail closed with 503 when the managed-file lookup errors

A lookup exception previously fell back to dispatching the unresolved
unified token, which defeats the fail-closed guarantee: the token still
reaches the provider and can hit the same publishers-segment IndexError
the resolution prevents. Treat a lookup error like the missing-row case
and fail closed, but with a retryable 503 since the condition is
transient. No-database and no-storage_url rows still fall back unchanged

---------

Co-authored-by: htourinho-clgx <htourinho@cotality.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
yucheng-berri 2026-07-24 17:29:38 -07:00 • committed by GitHub
parent 76b0b10908
commit 5677bc237c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 282 additions and 172 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -38,11 +38,44 @@ 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) -> "str | None":
"""Resolve a managed (unified) input_file_id to its backend storage_url.
Provider batch handlers (e.g. Vertex AI, which parses a `publishers/`
segment out of the file URI) need a real storage location; the opaque
unified token crashes them. Returns None only when there is no database or
the row has no storage_url yet, so callers fall back to the original id
(which the managed-files deployment hook can still map). Fails closed
rather than dispatch a token that cannot be resolved: 404 when no
managed-file row exists, 503 when the lookup itself errors so the caller
can retry.
"""
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)
raise HTTPException(
status_code=503,
detail={"error": "Could not resolve managed file; please retry"},
)
if db_file is None:
raise HTTPException(
status_code=404,
detail={"error": f"Managed file not found: {input_file_id}"},
)
return db_file.storage_url or None
@router.post(
"/{provider}/v1/batches",
dependencies=[Depends(user_api_key_auth)],
@ -224,6 +257,11 @@ async def create_batch(
)
model = target_model_names[0]
_create_batch_data["model"] = model
resolved_storage_url = await _resolve_managed_input_file_storage_url(input_file_id)
if resolved_storage_url is not None:
_create_batch_data["input_file_id"] = resolved_storage_url
if llm_router is None:
raise HTTPException(
status_code=500,

416
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,14 +192,13 @@ 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", {}))
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,
@ -283,9 +274,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 +296,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 +343,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 +478,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 +504,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)
@ -535,6 +516,182 @@ async def test_create__unified_file_id_not_exactly_one_model_400(harness, models
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_fails_closed_503(harness):
"""A lookup error leaves the token unresolved, so it fails closed with a
retryable 503 rather than dispatching the opaque id into the provider crash
it cannot parse. Nothing is dispatched."""
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),
):
with pytest.raises(ProxyException) as exc:
await call_create(harness)
assert exc.value.code == "503"
harness.router_acreate.assert_not_called()
harness.litellm_acreate.assert_not_called()
@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_fails_closed_404(harness):
"""With a database present, a managed unified id that has no row cannot be
resolved to a real storage location, so it fails closed with a 404 rather
than dispatching the opaque token, which would hit the Vertex
publishers-segment IndexError this PR exists to prevent."""
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),
):
with pytest.raises(ProxyException) as exc:
await call_create(harness)
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_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
@ -547,10 +704,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)
@ -579,9 +735,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()
@ -630,9 +784,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",
@ -738,12 +890,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"
# =========================================================================== #
@ -770,9 +917,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.
@ -823,9 +968,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={})
@ -846,11 +989,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,
@ -865,24 +1004,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", {}))
@ -956,9 +1083,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
@ -989,12 +1114,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
@ -1018,9 +1139,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.
@ -1125,9 +1244,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
@ -1151,9 +1268,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.
@ -1186,9 +1301,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
@ -1200,9 +1313,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"
)
@ -1275,9 +1386,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()))
@ -1295,11 +1404,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,
@ -1432,21 +1537,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")
# --------------------------------------------------------------------------- #
@ -1577,12 +1676,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"
# =========================================================================== #
@ -1645,9 +1739,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.
@ -1666,11 +1758,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,
@ -1685,22 +1773,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,
@ -1765,9 +1847,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
@ -1786,9 +1866,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
@ -1804,9 +1882,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.
@ -1845,8 +1921,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")
@ -1885,9 +1962,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"
@ -1954,10 +2029,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"
)
@ -1979,9 +2051,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)
@ -2000,12 +2073,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)
@ -2015,9 +2086,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")