From a75964f040da411cecd81af70e8a28a01c423007 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Sun, 15 Mar 2026 15:27:32 -0700 Subject: [PATCH] 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 --- litellm/proxy/batches_endpoints/endpoints.py | 22 ++++--------------- .../openai_files_endpoints/common_utils.py | 22 +++++++++++++++++++ 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 04358f136a..9fa2a51fa8 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -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( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 5d546733b7..b75b2f4640 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -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]],