From 0e9f4ca09c1fa37b20df49fcba0b7a2d0684b1f3 Mon Sep 17 00:00:00 2001 From: htourinho-clgx Date: Wed, 22 Jul 2026 18:15:37 +0200 Subject: [PATCH] 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> --- litellm/proxy/batches_endpoints/endpoints.py | 21 +++++ .../proxy/batches_endpoints/test_endpoints.py | 84 +++++++++++++++++++ 2 files changed, 105 insertions(+) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index fffa0bf86d2..2fd2e5e5fee 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -230,6 +230,27 @@ 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: + from litellm.repositories.table_repositories import ( + ManagedFileRepository, + ) + + db_file = await ManagedFileRepository(prisma_client).table.find_first( + where={"unified_file_id": unified_file_id} + ) + if db_file is not None and db_file.storage_url: + _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 diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py index 26c654cd154..4564f1ecd55 100644 --- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py @@ -535,6 +535,90 @@ 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 - the + unified id itself must never reach the router/provider call.""" + 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) + + 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, + ): + resp = await call_create(harness) + + # 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 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_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.""" + 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( + "litellm.repositories.table_repositories.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