Merge pull request #23718 from BerriAI/litellm_fix_vertex_ai_batch

Fix: Vertex ai Batch Output File Download Fails with 500

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)
                         

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(

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():
     """

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"
+    )