diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 838ee95b2b5..995c45b925a 100644 --- a/litellm/batches/main.py +++ b/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", diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 7932881f482..864cc190312 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -213,3 +213,99 @@ class VertexAIBatchPrediction(VertexLLM): 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 diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index 22cd0bd402a..a0adb3e55a8 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/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 @@ class VertexAIBatchTransformation: ), ) + @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 diff --git a/tests/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index 67ce2f7249a..2f4f9bbcda1 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -447,6 +447,18 @@ mock_vertex_batch_response = { "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 test_avertex_batch_prediction(monkeypatch): 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"