diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index cbd17845c5f..674db49a259 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -504,7 +504,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if retrieve_file_id else False ) - if potential_file_id: + if potential_file_id and "llm_output_file_id," in potential_file_id: model_id = self.get_model_id_from_unified_file_id(potential_file_id) if model_id: data["model"] = model_id @@ -1058,7 +1058,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return file_id.split("llm_output_file_model_id,")[1].split(";")[0] def get_output_file_id_from_unified_file_id(self, file_id: str) -> str: - return file_id.split("llm_output_file_id,")[1].split(";")[0] + marker = "llm_output_file_id," + if marker not in file_id: + raise ValueError( + f"Unified id does not contain {marker!r}: {file_id[:80]!r}" + ) + return file_id.split(marker, 1)[1].split(";")[0] async def async_post_call_success_hook( self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes @@ -1102,13 +1107,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): decoded_unified_file_id = _is_base64_encoded_unified_file_id( file_id_value ) - if decoded_unified_file_id: + if ( + decoded_unified_file_id + and "llm_output_file_id," in decoded_unified_file_id + ): provider_file_id = ( self.get_output_file_id_from_unified_file_id( decoded_unified_file_id ) ) unified_file_id = file_id_value + elif decoded_unified_file_id: + verbose_logger.warning( + f"Skipping {file_attr}={file_id_value!r}: " + "unified id is not a managed file output id" + ) + continue else: provider_file_id = file_id_value unified_file_id = self.get_unified_output_file_id( diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 7ca86d24494..1336490a344 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -222,3 +222,36 @@ async def test_should_not_double_wrap_already_unified_output_file_id(): already_unified + "=" * (-len(already_unified) % 4) ).decode() assert decoded.count(f"llm_output_file_id,{provider_file_id}") == 1 + + +@pytest.mark.asyncio +async def test_should_skip_non_file_unified_id_on_output_file_id(): + """Batch-style unified ids lack llm_output_file_id; must not IndexError or re-wrap.""" + import base64 + + managed_files = _make_managed_files_instance() + batch_unified = ( + base64.urlsafe_b64encode( + b"litellm_proxy;model_id:openai/openai/gpt-5.5-batch;llm_batch_id:batch_abc" + ) + .decode() + .rstrip("=") + ) + + batch_response = _make_batch_response( + model_id="openai/openai/gpt-5.5-batch", + model_name="openai/openai/gpt-5.5-batch", + output_file_id=batch_unified, + ) + user_api_key_dict = _make_user_api_key_dict() + + with patch("litellm.afile_retrieve", AsyncMock()) as mock_afile_retrieve: + await managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=user_api_key_dict, + response=batch_response, + ) + + assert batch_response.output_file_id == batch_unified + mock_afile_retrieve.assert_not_called() + managed_files.store_unified_file_id.assert_not_awaited()