From 5677bc237c491fd4579eb5b5894ed91982e75d45 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 24 Jul 2026 17:29:38 -0700 Subject: [PATCH] fix(proxy/batches): resolve managed unified input_file_id to storage_url with ownership check before dispatch (#34474) * 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 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- litellm/proxy/batches_endpoints/endpoints.py | 38 ++ .../proxy/batches_endpoints/test_endpoints.py | 416 ++++++++++-------- 2 files changed, 282 insertions(+), 172 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index fffa0bf86d2..af9aa901022 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -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, diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 26c654cd154..fe09241ee5c 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -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")