Managed batches fixes for vertex

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -589,7 +589,14 @@ async def async_pre_call_deployment_hook(
             model_file_id_mapping = cast(
                 Optional[Dict[str, Dict[str, str]]], kwargs.get("model_file_id_mapping")
             )
+            # model_info may be at top-level or nested under litellm_metadata
+            # (batch/file operations use litellm_metadata)
             model_id = cast(Optional[str], kwargs.get("model_info", {}).get("id", None))
+            if model_id is None:
+                model_id = cast(
+                    Optional[str],
+                    kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
+                )
             mapped_file_id: Optional[str] = None
             if input_file_id and model_file_id_mapping and model_id:
                 mapped_file_id = model_file_id_mapping.get(input_file_id, {}).get(

litellm/batches/batch_utils.py

@@ -128,73 +128,58 @@ def calculate_vertex_ai_batch_cost_and_usage(
     model_name: Optional[str] = None,
 ) -> Tuple[float, Usage]:
     """
-    Calculate both cost and usage from Vertex AI batch responses
+    Calculate both cost and usage from Vertex AI batch responses.
+
+    Vertex AI batch output lines have format:
+      {"request": ..., "status": "", "response": {"candidates": [...], "usageMetadata": {...}}}
+
+    usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount.
     """
-    from litellm.litellm_core_utils.litellm_logging import Logging
-    from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
-        VertexGeminiConfig,
-    )
+    from litellm.cost_calculator import batch_cost_calculator
+
     total_cost = 0.0
     total_tokens = 0
     prompt_tokens = 0
     completion_tokens = 0
-
+    actual_model_name = model_name or "gemini-2.0-flash-001"
+
     for response in vertex_ai_batch_responses:
-        if response.get("status") == "JOB_STATE_SUCCEEDED":  # Check if response was successful
-            # Transform Vertex AI response to OpenAI format if needed
-
-            # Create required arguments for the transformation method
-            model_response = ModelResponse()
-
-            # Ensure model_name is not None
-            actual_model_name = model_name or "gemini-2.5-flash"
-
-            # Create a real LiteLLM logging object
-            logging_obj = Logging(
-                model=actual_model_name,
-                messages=[{"role": "user", "content": "batch_request"}],
-                stream=False,
-                call_type=CallTypes.aretrieve_batch,
-                start_time=time.time(),
-                litellm_call_id="batch_" + str(uuid.uuid4()),
-                function_id="batch_processing",
-                litellm_trace_id=str(uuid.uuid4()),
-                kwargs={"optional_params": {}}
-            )
-
-            # Add the optional_params attribute that the Vertex AI transformation expects
-            logging_obj.optional_params = {}
-            raw_response = httpx.Response(200)  # Mock response object
-
-            openai_format_response = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response(
-                completion_response=response["response"],
-                model_response=model_response,
+        response_body = response.get("response")
+        if response_body is None:
+            continue
+
+        usage_metadata = response_body.get("usageMetadata", {})
+        _prompt = usage_metadata.get("promptTokenCount", 0) or 0
+        _completion = usage_metadata.get("candidatesTokenCount", 0) or 0
+        _total = usage_metadata.get("totalTokenCount", 0) or (_prompt + _completion)
+
+        line_usage = Usage(
+            prompt_tokens=_prompt,
+            completion_tokens=_completion,
+            total_tokens=_total,
+        )
+
+        try:
+            p_cost, c_cost = batch_cost_calculator(
+                usage=line_usage,
                 model=actual_model_name,
-                logging_obj=logging_obj,
-                raw_response=raw_response,
-            )
-
-            # Calculate cost using existing function
-            cost = litellm.completion_cost(
-                completion_response=openai_format_response,
                 custom_llm_provider="vertex_ai",
-                call_type=CallTypes.aretrieve_batch.value,
             )
-            total_cost += cost
-            
-            # Extract usage from the transformed response
-            usage_obj = getattr(openai_format_response, 'usage', None)
-            if usage_obj:
-                usage = usage_obj
-            else:
-                # Fallback: create usage from response dict
-                response_dict = openai_format_response.dict() if hasattr(openai_format_response, 'dict') else {}
-                usage = _get_batch_job_usage_from_response_body(response_dict)
-            
-            total_tokens += usage.total_tokens
-            prompt_tokens += usage.prompt_tokens
-            completion_tokens += usage.completion_tokens
-    
+            total_cost += p_cost + c_cost
+        except Exception as e:
+            verbose_logger.debug(
+                "vertex_ai batch cost calculation error for line: %s", str(e)
+            )
+
+        prompt_tokens += _prompt
+        completion_tokens += _completion
+        total_tokens += _total
+
+    verbose_logger.info(
+        "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d",
+        total_cost, prompt_tokens, completion_tokens, total_tokens,
+    )
+
     return total_cost, Usage(
         total_tokens=total_tokens,
         prompt_tokens=prompt_tokens,

litellm/files/main.py

@@ -295,7 +295,7 @@ def create_file(
 @client
 async def afile_retrieve(
     file_id: str,
-    custom_llm_provider: Literal["openai", "azure", "gemini", "hosted_vllm", "manus"] = "openai",
+    custom_llm_provider: Literal["openai", "azure", "gemini", "vertex_ai", "hosted_vllm", "manus"] = "openai",
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
     **kwargs,

litellm/llms/vertex_ai/batches/handler.py

@@ -108,11 +108,18 @@ async def _async_create_batch(
         client = get_async_httpx_client(
             llm_provider=litellm.LlmProviders.VERTEX_AI,
         )
-        response = await client.post(
-            url=api_base,
-            headers=headers,
-            data=json.dumps(vertex_batch_request),
-        )
+        try:
+            response = await client.post(
+                url=api_base,
+                headers=headers,
+                data=json.dumps(vertex_batch_request),
+            )
+        except httpx.HTTPStatusError as e:
+            error_body = e.response.text if hasattr(e, 'response') else "N/A"
+            litellm.verbose_logger.error(
+                f"Vertex AI batch create failed: status={e.response.status_code}, body={error_body[:1000]}"
+            )
+            raise
         if response.status_code != 200:
             raise Exception(f"Error: {response.status_code} {response.text}")
 

litellm/llms/vertex_ai/batches/transformation.py

@@ -29,7 +29,7 @@ def transform_openai_batch_request_to_vertex_ai_batch_request(
         if input_file_id is None:
             raise ValueError("input_file_id is required, but not provided")
         input_config: InputConfig = InputConfig(
-            gcsSource=GcsSource(uris=input_file_id), instancesFormat="jsonl"
+            gcsSource=GcsSource(uris=[input_file_id]), instancesFormat="jsonl"
         )
         model: str = cls._get_model_from_gcs_file(input_file_id)
         output_config: OutputConfig = OutputConfig(

litellm/llms/vertex_ai/files/transformation.py

@@ -335,37 +335,84 @@ def get_error_class(
             status_code=status_code, message=error_message, headers=headers
         )
 
+    def _parse_gcs_uri(self, file_id: str) -> Tuple[str, str]:
+        """
+        Parse a GCS URI (gs://bucket/path/to/object) into (bucket, url-encoded-object-path).
+        Handles both raw and URL-encoded input.
+        """
+        import urllib.parse
+
+        decoded = urllib.parse.unquote(file_id)
+        if decoded.startswith("gs://"):
+            full_path = decoded[5:]
+        else:
+            full_path = decoded
+
+        if "/" in full_path:
+            bucket_name, object_path = full_path.split("/", 1)
+        else:
+            bucket_name = full_path
+            object_path = ""
+
+        encoded_object = urllib.parse.quote(object_path, safe="")
+        return bucket_name, encoded_object
+
     def transform_retrieve_file_request(
         self,
         file_id: str,
         optional_params: dict,
         litellm_params: dict,
     ) -> tuple[str, dict]:
-        raise NotImplementedError("VertexAIFilesConfig does not support file retrieval")
+        bucket, encoded_object = self._parse_gcs_uri(file_id)
+        url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}"
+        return url, {}
 
     def transform_retrieve_file_response(
         self,
         raw_response: Response,
         logging_obj: LiteLLMLoggingObj,
         litellm_params: dict,
     ) -> OpenAIFileObject:
-        raise NotImplementedError("VertexAIFilesConfig does not support file retrieval")
+        response_json = raw_response.json()
+        gcs_id = response_json.get("id", "")
+        gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else ""
+        return OpenAIFileObject(
+            id=f"gs://{gcs_id}",
+            bytes=int(response_json.get("size", 0)),
+            created_at=_convert_vertex_datetime_to_openai_datetime(
+                vertex_datetime=response_json.get("timeCreated", "")
+            ),
+            filename=response_json.get("name", ""),
+            object="file",
+            purpose=response_json.get("metadata", {}).get("purpose", "batch"),
+            status="processed",
+            status_details=None,
+        )
 
     def transform_delete_file_request(
         self,
         file_id: str,
         optional_params: dict,
         litellm_params: dict,
     ) -> tuple[str, dict]:
-        raise NotImplementedError("VertexAIFilesConfig does not support file deletion")
+        bucket, encoded_object = self._parse_gcs_uri(file_id)
+        url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}"
+        return url, {}
 
     def transform_delete_file_response(
         self,
         raw_response: Response,
         logging_obj: LiteLLMLoggingObj,
         litellm_params: dict,
     ) -> FileDeleted:
-        raise NotImplementedError("VertexAIFilesConfig does not support file deletion")
+        file_id = "deleted"
+        if hasattr(raw_response, "request") and raw_response.request:
+            url = str(raw_response.request.url)
+            if "/o/" in url:
+                import urllib.parse
+                encoded_name = url.split("/o/")[-1].split("?")[0]
+                file_id = f"gs://{urllib.parse.unquote(encoded_name)}"
+        return FileDeleted(id=file_id, deleted=True, object="file")
 
     def transform_list_files_request(
         self,
@@ -389,15 +436,18 @@ def transform_file_content_request(
         optional_params: dict,
         litellm_params: dict,
     ) -> tuple[str, dict]:
-        raise NotImplementedError("VertexAIFilesConfig does not support file content retrieval")
+        file_id = file_content_request.get("file_id", "")
+        bucket, encoded_object = self._parse_gcs_uri(file_id)
+        url = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}?alt=media"
+        return url, {}
 
     def transform_file_content_response(
         self,
         raw_response: Response,
         logging_obj: LiteLLMLoggingObj,
         litellm_params: dict,
     ) -> HttpxBinaryResponseContent:
-        raise NotImplementedError("VertexAIFilesConfig does not support file content retrieval")
+        return HttpxBinaryResponseContent(response=raw_response)
 
 
 class VertexAIJsonlFilesTransformation(VertexGeminiConfig):

litellm/types/llms/vertex_ai.py

@@ -560,7 +560,7 @@ class VertexAIBatchEmbeddingsResponseObject(TypedDict):
 
 
 class GcsSource(TypedDict):
-    uris: str
+    uris: List[str]
 
 
 class InputConfig(TypedDict):

tests/batches_tests/test_openai_batches_and_files.py

@@ -29,6 +29,7 @@
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.types.utils import StandardLoggingPayload
 import random
+import httpx
 from unittest.mock import patch, MagicMock
 
 
@@ -579,6 +580,48 @@ async def test_vertex_list_batches(monkeypatch):
         assert list_response["data"][1].id == "test-batch-id-789"
 
 
+@pytest.mark.asyncio
+async def test_vertex_async_create_batch_logs_error_body_on_http_error():
+    """
+    When Vertex AI returns an HTTP error (e.g. 400), _async_create_batch should
+    re-raise httpx.HTTPStatusError (not swallow it) and log the response body.
+
+    Before the fix the error body was lost because AsyncHTTPHandler.post()
+    calls raise_for_status() internally, raising before the handler's own
+    status-code check could log the body.
+    """
+    from litellm.llms.vertex_ai.batches.handler import VertexAIBatchPrediction
+
+    handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket")
+
+    error_body = '{"error": {"code": 400, "message": "Do not support publisher model gemini-2.0-flash"}}'
+
+    mock_response = MagicMock(spec=httpx.Response)
+    mock_response.status_code = 400
+    mock_response.text = error_body
+    mock_response.headers = {}
+
+    http_error = httpx.HTTPStatusError(
+        message="Bad Request",
+        request=httpx.Request("POST", "https://fake-vertex-url"),
+        response=mock_response,
+    )
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
+        side_effect=http_error,
+    ):
+        with pytest.raises(httpx.HTTPStatusError) as exc_info:
+            await handler._async_create_batch(
+                vertex_batch_request={},
+                api_base="https://us-central1-aiplatform.googleapis.com/v1/projects/test/locations/us-central1/batchPredictionJobs",
+                headers={"Authorization": "Bearer fake-token"},
+            )
+
+        assert exc_info.value.response.status_code == 400
+        assert "gemini-2.0-flash" in exc_info.value.response.text
+
+
 @pytest.mark.asyncio
 async def test_delete_batch_output_file():
     """