scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

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

Browse source 
fix: support Vertex AI batch listing in LiteLLM proxy
This commit is contained in:
Sameer Kankute 2025-11-27 16:50:17 +05:30 • committed by GitHub
commit 1f6d8fc151
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 198 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
litellm/batches/main.py
View file

@ -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",

96
litellm/llms/vertex_ai/batches/handler.py
View file

@ -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

29
litellm/llms/vertex_ai/batches/transformation.py
View file

@ -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

44
tests/batches_tests/test_openai_batches_and_files.py
View file

@ -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"