mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
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:
parent
2798f39f5a
commit
0e9f4ca09c
2 changed files with 105 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue