Merge pull request #17079 from naaa760/fix/vertex-batch-support

fix: support Vertex AI batch listing in LiteLLM proxy

Author: Sameerlite

litellm/batches/main.py

@@ -644,7 +644,7 @@ def retrieve_batch(
 async def alist_batches(
     after: Optional[str] = None,
     limit: Optional[int] = None,
-    custom_llm_provider: Literal["openai", "azure"] = "openai",
+    custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
     metadata: Optional[Dict[str, str]] = None,
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
@@ -687,7 +687,7 @@ async def alist_batches(
 def list_batches(
     after: Optional[str] = None,
     limit: Optional[int] = None,
-    custom_llm_provider: Literal["openai", "azure"] = "openai",
+    custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai",
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
     **kwargs,
@@ -784,9 +784,36 @@ def list_batches(
                 max_retries=optional_params.max_retries,
                 litellm_params=litellm_params,
             )
+        elif custom_llm_provider == "vertex_ai":
+            api_base = optional_params.api_base or ""
+            vertex_ai_project = (
+                optional_params.vertex_project
+                or litellm.vertex_project
+                or get_secret_str("VERTEXAI_PROJECT")
+            )
+            vertex_ai_location = (
+                optional_params.vertex_location
+                or litellm.vertex_location
+                or get_secret_str("VERTEXAI_LOCATION")
+            )
+            vertex_credentials = optional_params.vertex_credentials or get_secret_str(
+                "VERTEXAI_CREDENTIALS"
+            )
+
+            response = vertex_ai_batches_instance.list_batches(
+                _is_async=_is_async,
+                after=after,
+                limit=limit,
+                api_base=api_base,
+                vertex_project=vertex_ai_project,
+                vertex_location=vertex_ai_location,
+                vertex_credentials=vertex_credentials,
+                timeout=timeout,
+                max_retries=optional_params.max_retries,
+            )
         else:
             raise litellm.exceptions.BadRequestError(
-                message="LiteLLM doesn't support {} for 'list_batch'. Only 'openai' is supported.".format(
+                message="LiteLLM doesn't support {} for 'list_batch'. Supported providers: openai, azure, vertex_ai.".format(
                     custom_llm_provider
                 ),
                 model="n/a",

litellm/llms/vertex_ai/batches/handler.py

@@ -213,3 +213,99 @@ async def _async_retrieve_batch(
             response=_json_response
         )
         return vertex_batch_response
+
+    def list_batches(
+        self,
+        _is_async: bool,
+        after: Optional[str],
+        limit: Optional[int],
+        api_base: Optional[str],
+        vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES],
+        vertex_project: Optional[str],
+        vertex_location: Optional[str],
+        timeout: Union[float, httpx.Timeout],
+        max_retries: Optional[int],
+    ):
+        sync_handler = _get_httpx_client()
+
+        access_token, project_id = self._ensure_access_token(
+            credentials=vertex_credentials,
+            project_id=vertex_project,
+            custom_llm_provider="vertex_ai",
+        )
+
+        default_api_base = self.create_vertex_batch_url(
+            vertex_location=vertex_location or "us-central1",
+            vertex_project=vertex_project or project_id,
+        )
+
+        if len(default_api_base.split(":")) > 1:
+            endpoint = default_api_base.split(":")[-1]
+        else:
+            endpoint = ""
+
+        _, api_base = self._check_custom_proxy(
+            api_base=api_base,
+            custom_llm_provider="vertex_ai",
+            gemini_api_key=None,
+            endpoint=endpoint,
+            stream=None,
+            auth_header=None,
+            url=default_api_base,
+        )
+
+        headers = {
+            "Content-Type": "application/json; charset=utf-8",
+            "Authorization": f"Bearer {access_token}",
+        }
+
+        params: Dict[str, Any] = {}
+        if limit is not None:
+            params["pageSize"] = str(limit)
+        if after is not None:
+            params["pageToken"] = after
+
+        if _is_async is True:
+            return self._async_list_batches(
+                api_base=api_base,
+                headers=headers,
+                params=params,
+            )
+
+        response = sync_handler.get(
+            url=api_base,
+            headers=headers,
+            params=params,
+        )
+
+        if response.status_code != 200:
+            raise Exception(f"Error: {response.status_code} {response.text}")
+
+        _json_response = response.json()
+        vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_list_response_to_openai_list_response(
+            response=_json_response
+        )
+        return vertex_batch_response
+
+    async def _async_list_batches(
+        self,
+        api_base: str,
+        headers: Dict[str, str],
+        params: Dict[str, Any],
+    ):
+        client = get_async_httpx_client(
+            llm_provider=litellm.LlmProviders.VERTEX_AI,
+        )
+        response = await client.get(
+            url=api_base,
+            headers=headers,
+            params=params,
+        )
+        if response.status_code != 200:
+            raise Exception(f"Error: {response.status_code} {response.text}")
+
+        _json_response = response.json()
+        vertex_batch_response = VertexAIBatchTransformation.transform_vertex_ai_batch_list_response_to_openai_list_response(
+            response=_json_response
+        )
+        return vertex_batch_response

litellm/llms/vertex_ai/batches/transformation.py

@@ -1,5 +1,5 @@
 from litellm._uuid import uuid
-from typing import Dict
+from typing import Any, Dict
 
 from litellm.llms.vertex_ai.common_utils import (
     _convert_vertex_datetime_to_openai_datetime,
@@ -67,6 +67,33 @@ def transform_vertex_ai_batch_response_to_openai_batch_response(
             ),
         )
 
+    @classmethod
+    def transform_vertex_ai_batch_list_response_to_openai_list_response(
+        cls, response: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """
+        Transforms Vertex AI batch list response into OpenAI-compatible list response.
+        """
+
+        batch_jobs = response.get("batchPredictionJobs", []) or []
+        data = [
+            cls.transform_vertex_ai_batch_response_to_openai_batch_response(job)
+            for job in batch_jobs
+        ]
+
+        first_id = data[0].id if len(data) > 0 else None
+        last_id = data[-1].id if len(data) > 0 else None
+        next_page_token = response.get("nextPageToken")
+
+        return {
+            "object": "list",
+            "data": data,
+            "first_id": first_id,
+            "last_id": last_id,
+            "has_more": bool(next_page_token),
+            "next_page_token": next_page_token,
+        }
+
     @classmethod
     def _get_batch_id_from_vertex_ai_batch_response(
         cls, response: VertexBatchPredictionResponse

tests/batches_tests/test_openai_batches_and_files.py

@@ -447,6 +447,18 @@ async def test_async_create_batch(provider):
     "completionStats": {"successfulCount": 0, "failedCount": 0, "remainingCount": 100},
 }
 
+mock_vertex_list_response = {
+    "batchPredictionJobs": [
+        mock_vertex_batch_response,
+        {
+            **mock_vertex_batch_response,
+            "name": "projects/123456789/locations/us-central1/batchPredictionJobs/test-batch-id-789",
+            "state": "JOB_STATE_SUCCEEDED",
+        },
+    ],
+    "nextPageToken": "",
+}
+
 
 @pytest.mark.asyncio
 async def test_avertex_batch_prediction(monkeypatch):
@@ -533,3 +545,35 @@ async def mock_side_effect(*args, **kwargs):
             print("retrieved_batch=", retrieved_batch)
 
             assert retrieved_batch.id == "test-batch-id-456"
+
+
+@pytest.mark.asyncio
+async def test_vertex_list_batches(monkeypatch):
+    monkeypatch.setenv("GCS_BUCKET_NAME", "litellm-local")
+    monkeypatch.setenv("VERTEXAI_PROJECT", "litellm-test-project")
+    monkeypatch.setenv("VERTEXAI_LOCATION", "us-central1")
+
+    monkeypatch.setattr(
+        "litellm.llms.vertex_ai.batches.handler.VertexAIBatchPrediction._ensure_access_token",
+        lambda self, credentials, project_id, custom_llm_provider: ("mock-token", "litellm-test-project"),
+    )
+
+    with patch(
+        "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.get"
+    ) as mock_get:
+        mock_get_response = MagicMock()
+        mock_get_response.json.return_value = mock_vertex_list_response
+        mock_get_response.status_code = 200
+        mock_get_response.raise_for_status.return_value = None
+        mock_get.return_value = mock_get_response
+
+        list_response = await litellm.alist_batches(
+            custom_llm_provider="vertex_ai",
+            limit=2,
+        )
+
+        assert list_response["object"] == "list"
+        assert list_response["has_more"] is False
+        assert len(list_response["data"]) == 2
+        assert list_response["data"][0].id == "test-batch-id-456"
+        assert list_response["data"][1].id == "test-batch-id-789"