diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 37ca341fdf..5530054170 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -26,6 +26,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( get_batch_id_from_unified_batch_id, get_content_type_from_file_object, get_model_id_from_unified_batch_id, + get_models_from_unified_file_id, normalize_mime_type_for_provider, ) from litellm.types.llms.openai import ( @@ -904,6 +905,21 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) # managed batch id model_id = cast(Optional[str], response._hidden_params.get("model_id")) model_name = cast(Optional[str], response._hidden_params.get("model_name")) + resolved_model_name = model_name + + # Some providers (e.g. Vertex batch retrieve) do not set model_name on + # the response. In that case, recover target_model_names from the input + # managed file metadata so unified output IDs preserve routing metadata. + if not resolved_model_name and isinstance(unified_file_id, str): + decoded_unified_file_id = ( + _is_base64_encoded_unified_file_id(unified_file_id) + or unified_file_id + ) + target_model_names = get_models_from_unified_file_id( + decoded_unified_file_id + ) + if target_model_names: + resolved_model_name = ",".join(target_model_names) original_response_id = response.id if (unified_batch_id or unified_file_id) and model_id: @@ -919,7 +935,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): unified_file_id = self.get_unified_output_file_id( output_file_id=original_file_id, model_id=model_id, - model_name=model_name, + model_name=resolved_model_name, ) setattr(response, file_attr, unified_file_id) diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index 7cb06fea9e..86bdc2c7b5 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -1,6 +1,6 @@ -from litellm._uuid import uuid from typing import Any, Dict +from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( _convert_vertex_datetime_to_openai_datetime, ) @@ -144,9 +144,10 @@ class VertexAIBatchTransformation: output_file_id: str = ( response.get("outputInfo", OutputInfo()).get("gcsOutputDirectory", "") - + "/predictions.jsonl" ) - if output_file_id != "/predictions.jsonl": + if output_file_id: + output_file_id = output_file_id.rstrip("/") + "/predictions.jsonl" + if output_file_id and output_file_id != "/predictions.jsonl": return output_file_id output_config = response.get("outputConfig") @@ -158,7 +159,9 @@ class VertexAIBatchTransformation: return output_file_id output_uri_prefix = gcs_destination.get("outputUriPrefix", "") - return output_uri_prefix + if output_uri_prefix.endswith("/predictions.jsonl"): + return output_uri_prefix + return output_uri_prefix.rstrip("/") + "/predictions.jsonl" @classmethod def _get_batch_job_status_from_vertex_ai_batch_response( diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 58fbd9e64b..9f4ca4ed10 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1,4 +1,6 @@ +import base64 import json +from typing import cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -477,6 +479,81 @@ async def test_output_file_id_for_batch_retrieve(): assert not cast(LiteLLMBatch, response).output_file_id.startswith("file-") +@pytest.mark.asyncio +async def test_output_file_id_preserves_target_model_names_when_model_name_missing(): + """ + Regression test: when provider response does not include _hidden_params.model_name + (e.g. Vertex batch retrieve), unified output_file_id should still include + target_model_names from the managed input file ID. + """ + from openai.types.batch import BatchRequestCounts + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.llms.openai import OpenAIFileObject + from litellm.types.utils import LiteLLMBatch + + batch = LiteLLMBatch( + id="batch_123", + completion_window="24h", + created_at=1750883933, + endpoint="/v1/chat/completions", + input_file_id="file-input-provider-id", + object="batch", + status="completed", + output_file_id="file-provider-output-id", + request_counts=BatchRequestCounts(completed=1, failed=0, total=1), + usage=None, + ) + + # Build a valid managed input id string and base64 encode it. + managed_input_file_payload = ( + "litellm_proxy:application/octet-stream;" + "unified_id,test-uuid;" + "target_model_names,gemini-2.5-pro;" + "llm_output_file_id,file-input-1;" + "llm_output_file_model_id,model-id-1" + ) + managed_input_file_id = ( + base64.urlsafe_b64encode(managed_input_file_payload.encode()) + .decode() + .rstrip("=") + ) + + batch._hidden_params = { + "model_id": "model-id-1", + "unified_batch_id": "litellm_proxy;model_id:model-id-1;llm_batch_id:batch_123", + "unified_file_id": managed_input_file_id, + # Intentionally omit model_name to mimic Vertex issue. + } + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=AsyncMock() + ) + + provider_output_file = OpenAIFileObject( + id="file-provider-output-id", + object="file", + bytes=1, + created_at=1, + filename="predictions.jsonl", + purpose="batch_output", + ) + + with patch("litellm.afile_retrieve", new_callable=AsyncMock) as mock_retrieve: + mock_retrieve.return_value = provider_output_file + response = await proxy_managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), + response=batch, + ) + + decoded_output_file_id = _is_base64_encoded_unified_file_id( + cast(LiteLLMBatch, response).output_file_id + ) + assert decoded_output_file_id + assert "target_model_names,gemini-2.5-pro" in cast(str, decoded_output_file_id) + + @pytest.mark.asyncio async def test_error_file_id_for_failed_batch(): """ diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py new file mode 100644 index 0000000000..1aab74ddc2 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_batch_transformation.py @@ -0,0 +1,38 @@ +from litellm.llms.vertex_ai.batches.transformation import VertexAIBatchTransformation + + +def test_output_file_id_uses_predictions_jsonl_with_output_info(): + response = { + "outputInfo": { + "gcsOutputDirectory": "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-123" + } + } + + output_file_id = VertexAIBatchTransformation._get_output_file_id_from_vertex_ai_batch_response( + response + ) + + assert ( + output_file_id + == "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-123/predictions.jsonl" + ) + + +def test_output_file_id_falls_back_to_output_uri_prefix_with_predictions_jsonl(): + response = { + "outputInfo": {}, + "outputConfig": { + "gcsDestination": { + "outputUriPrefix": "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-456" + } + }, + } + + output_file_id = VertexAIBatchTransformation._get_output_file_id_from_vertex_ai_batch_response( + response + ) + + assert ( + output_file_id + == "gs://test-bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-pro/prediction-model-456/predictions.jsonl" + )