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>
This commit is contained in:
htourinho-clgx 2026-07-22 18:15:37 +02:00 committed by Yucheng Zhu
parent 2798f39f5a
commit 0e9f4ca09c
2 changed files with 105 additions and 0 deletions

View file

@ -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

View file

@ -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