Fix batch retrieve double-encoding: resolve output_file_id via DB lookup instead of hooks

The previous approach populated _hidden_params to trigger the managed files
hook, but the hook also re-encodes response.id (batch ID), causing double-
encoding when the DB already stores unified IDs. Instead, resolve raw
output_file_id/error_file_id to unified IDs via a direct DB lookup (same
pattern as resolve_input_file_id_to_unified), which avoids the hook entirely.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
yuneng-jiang
2026-03-15 15:27:40 -07:00
co-authored by Claude Opus 4.6
parent ff869e91b0
commit a75964f040
2 changed files with 26 additions and 18 deletions
+4 -18
View File
@@ -32,6 +32,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (
get_original_file_id,
prepare_data_with_credentials,
resolve_input_file_id_to_unified,
resolve_output_file_ids_to_unified,
update_batch_in_database,
)
from litellm.proxy.utils import handle_exception_on_proxy, is_known_model
@@ -412,31 +413,16 @@ async def retrieve_batch( # noqa: PLR0915
"cancelled",
"expired",
]:
# Populate _hidden_params so managed files hook can translate IDs.
# When the response comes from the DB, _hidden_params is empty.
# Only set if output_file_id is still a raw provider ID (not yet unified).
# The DB may store unified IDs after a previous hook run; setting
# _hidden_params in that case would cause double-encoding.
if unified_batch_id:
_output_fid = getattr(response, "output_file_id", None)
_needs_translation = _output_fid and not _is_base64_encoded_unified_file_id(_output_fid)
if _needs_translation:
response._hidden_params["unified_batch_id"] = unified_batch_id
model_id_from_batch = get_model_id_from_unified_batch_id(
unified_batch_id
)
if model_id_from_batch:
response._hidden_params["model_id"] = model_id_from_batch
# Call hooks and return
response = await proxy_logging_obj.post_call_success_hook(
data=data, user_api_key_dict=user_api_key_dict, response=response
)
# async_post_call_success_hook replaces batch.id and output_file_id with unified IDs
# but not input_file_id. Resolve raw provider ID to unified ID.
# The DB may store raw provider file IDs (before hooks translate them).
# Resolve any raw input/output/error file IDs to unified IDs.
if unified_batch_id:
await resolve_input_file_id_to_unified(response, prisma_client)
await resolve_output_file_ids_to_unified(response, prisma_client)
asyncio.create_task(
proxy_logging_obj.update_request_status(
@@ -697,6 +697,28 @@ async def resolve_input_file_id_to_unified(response, prisma_client) -> None:
pass
async def resolve_output_file_ids_to_unified(response, prisma_client) -> None:
"""
If the batch response contains raw provider output_file_id or error_file_id
(not already unified IDs), look up the corresponding unified file IDs from
the managed file table and replace them in-place.
"""
if not prisma_client:
return
for attr in ("output_file_id", "error_file_id"):
raw_id = getattr(response, attr, None)
if not raw_id or _is_base64_encoded_unified_file_id(raw_id):
continue
try:
managed_file = await prisma_client.db.litellm_managedfiletable.find_first(
where={"flat_model_file_ids": {"has": raw_id}}
)
if managed_file:
setattr(response, attr, managed_file.unified_file_id)
except Exception:
pass
async def get_batch_from_database(
batch_id: str,
unified_batch_id: Union[str, Literal[False]],