tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-02 10:21:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #22464 from Point72/ephrimstanley/batch-fixes-feb27

Browse Source 
Managed batches fixes for vertex
This commit is contained in:
Sameer Kankute
2026-03-03 18:53:53 +05:30
committed by GitHub
parent 67f90254ed b16397ae1a
commit 9ffbd9e30e
11 changed files with 526 additions and 217 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
+7
View File
@@ -589,7 +589,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
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
+43 -58
View File
@@ -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
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
actual_model_name = model_name or "gemini-2.0-flash-001"
# 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(
for response in vertex_ai_batch_responses:
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,
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,
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
+1 -1
View File
@@ -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
+12 -5
View File
@@ -108,11 +108,18 @@ class VertexAIBatchPrediction(VertexLLM):
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
+1 -1
View File
@@ -29,7 +29,7 @@ class VertexAIBatchTransformation:
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
+56 -6
View File
@@ -335,13 +335,37 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
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,
@@ -349,7 +373,21 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
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,
@@ -357,7 +395,9 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
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,
@@ -365,7 +405,14 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
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,7 +436,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
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,
@@ -397,7 +447,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
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
+1 -1
View File
@@ -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
+43
View File
@@ -29,6 +29,7 @@ verbose_logger.setLevel(logging.DEBUG)
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():
"""
tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py
+104
View File
@@ -6,6 +6,7 @@ from fastapi import HTTPException
from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
from litellm.caching import DualCache
from litellm.proxy._types import CallTypes
from litellm.proxy.openai_files_endpoints.common_utils import (
_is_base64_encoded_unified_file_id,
)
@@ -61,6 +62,109 @@ async def test_async_pre_call_hook_batch_retrieve():
assert response["model"] == "my-general-azure-deployment"
@pytest.mark.asyncio
async def test_async_pre_call_deployment_hook_resolves_model_id_from_litellm_metadata():
"""
For batch operations the router stores model_info under
kwargs["litellm_metadata"]["model_info"] (not top-level kwargs["model_info"]).
async_pre_call_deployment_hook must check both locations so the managed
file ID is resolved to the provider-specific file ID.
"""
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=MagicMock()
)
managed_file_id = "managed-file-abc"
model_id = "deployment-xyz"
provider_file_id = "gs://bucket/path/to/file.jsonl"
# model_info is nested under litellm_metadata (batch path)
kwargs = {
"input_file_id": managed_file_id,
"model_file_id_mapping": {
managed_file_id: {model_id: provider_file_id},
},
"litellm_metadata": {
"model_info": {"id": model_id},
},
}
result = await proxy_managed_files.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.acreate_batch
)
assert result["input_file_id"] == provider_file_id, (
f"Expected provider file ID '{provider_file_id}', got '{result['input_file_id']}'"
)
@pytest.mark.asyncio
async def test_async_pre_call_deployment_hook_prefers_top_level_model_info():
"""
When model_info exists at top-level kwargs, async_pre_call_deployment_hook
should use it without falling back to litellm_metadata.
"""
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=MagicMock()
)
managed_file_id = "managed-file-abc"
top_level_model_id = "deployment-top"
nested_model_id = "deployment-nested"
top_level_provider_file = "file-top-123"
nested_provider_file = "file-nested-456"
kwargs = {
"input_file_id": managed_file_id,
"model_file_id_mapping": {
managed_file_id: {
top_level_model_id: top_level_provider_file,
nested_model_id: nested_provider_file,
},
},
"model_info": {"id": top_level_model_id},
"litellm_metadata": {
"model_info": {"id": nested_model_id},
},
}
result = await proxy_managed_files.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.acreate_batch
)
assert result["input_file_id"] == top_level_provider_file, (
"Should prefer top-level model_info over litellm_metadata"
)
@pytest.mark.asyncio
async def test_async_pre_call_deployment_hook_no_model_info_leaves_file_id_unchanged():
"""
When model_info is absent from both top-level and litellm_metadata,
the managed file ID should remain unchanged.
"""
proxy_managed_files = _PROXY_LiteLLMManagedFiles(
DualCache(), prisma_client=MagicMock()
)
managed_file_id = "managed-file-abc"
kwargs = {
"input_file_id": managed_file_id,
"model_file_id_mapping": {
managed_file_id: {"some-model": "provider-file-xyz"},
},
}
result = await proxy_managed_files.async_pre_call_deployment_hook(
kwargs=kwargs, call_type=CallTypes.acreate_batch
)
assert result["input_file_id"] == managed_file_id, (
"File ID should remain unchanged when model_info is not available"
)
# def test_list_managed_files():
# proxy_managed_files = _PROXY_LiteLLMManagedFiles(DualCache())
tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py
+184
View File
@@ -0,0 +1,184 @@
"""
Tests for VertexAIFilesConfig transformation methods (Issues 5-7).
"""
import json
import urllib.parse
import httpx
import pytest
from unittest.mock import MagicMock
from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig
from litellm.types.llms.openai import OpenAIFileObject, HttpxBinaryResponseContent
from openai.types.file_deleted import FileDeleted
@pytest.fixture
def config():
return VertexAIFilesConfig()
class TestParseGcsUri:
"""Tests for the _parse_gcs_uri helper used by retrieve / content / delete."""
def test_should_parse_standard_gs_uri(self, config):
bucket, encoded = config._parse_gcs_uri(
"gs://my-bucket/path/to/object.jsonl"
)
assert bucket == "my-bucket"
assert encoded == urllib.parse.quote("path/to/object.jsonl", safe="")
def test_should_parse_uri_with_nested_publisher_path(self, config):
uri = "gs://litellm-local/litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123"
bucket, encoded = config._parse_gcs_uri(uri)
assert bucket == "litellm-local"
expected_path = "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc-123"
assert encoded == urllib.parse.quote(expected_path, safe="")
def test_should_handle_url_encoded_input(self, config):
encoded_uri = urllib.parse.quote("gs://my-bucket/some/path", safe="")
bucket, encoded = config._parse_gcs_uri(encoded_uri)
assert bucket == "my-bucket"
assert encoded == urllib.parse.quote("some/path", safe="")
def test_should_handle_bucket_only(self, config):
bucket, encoded = config._parse_gcs_uri("gs://my-bucket")
assert bucket == "my-bucket"
assert encoded == ""
def test_should_handle_no_gs_prefix(self, config):
bucket, encoded = config._parse_gcs_uri("my-bucket/object.txt")
assert bucket == "my-bucket"
assert encoded == "object.txt"
class TestTransformRetrieveFile:
def test_should_build_correct_gcs_metadata_url(self, config):
file_id = "gs://my-bucket/path/to/file.jsonl"
url, params = config.transform_retrieve_file_request(
file_id=file_id, optional_params={}, litellm_params={}
)
expected_encoded = urllib.parse.quote("path/to/file.jsonl", safe="")
assert url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{expected_encoded}"
assert params == {}
def test_should_return_openai_file_object_from_gcs_response(self, config):
gcs_json = {
"id": "my-bucket/path/to/file.jsonl/123456",
"name": "path/to/file.jsonl",
"size": "4096",
"timeCreated": "2025-02-15T10:00:00.000Z",
"metadata": {"purpose": "batch"},
}
raw_response = MagicMock(spec=httpx.Response)
raw_response.json.return_value = gcs_json
result = config.transform_retrieve_file_response(
raw_response=raw_response,
logging_obj=MagicMock(),
litellm_params={},
)
assert isinstance(result, OpenAIFileObject)
assert result.id == "gs://my-bucket/path/to/file.jsonl"
assert result.filename == "path/to/file.jsonl"
assert result.bytes == 4096
assert result.object == "file"
assert result.status == "processed"
assert result.purpose == "batch"
def test_should_default_purpose_to_batch_when_metadata_missing(self, config):
gcs_json = {
"id": "bucket/obj/999",
"name": "obj",
"size": "0",
"timeCreated": "2025-01-01T00:00:00.000Z",
}
raw_response = MagicMock(spec=httpx.Response)
raw_response.json.return_value = gcs_json
result = config.transform_retrieve_file_response(
raw_response=raw_response,
logging_obj=MagicMock(),
litellm_params={},
)
assert result.purpose == "batch"
class TestTransformFileContent:
def test_should_build_gcs_media_download_url(self, config):
file_id = "gs://my-bucket/path/to/file.jsonl"
url, params = config.transform_file_content_request(
file_content_request={"file_id": file_id},
optional_params={},
litellm_params={},
)
encoded = urllib.parse.quote("path/to/file.jsonl", safe="")
assert url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}?alt=media"
assert params == {}
def test_should_return_binary_response_content(self, config):
raw_response = httpx.Response(
status_code=200,
content=b'{"line": 1}\n{"line": 2}\n',
headers={"content-type": "application/octet-stream"},
request=httpx.Request("GET", "https://example.com"),
)
result = config.transform_file_content_response(
raw_response=raw_response,
logging_obj=MagicMock(),
litellm_params={},
)
assert isinstance(result, HttpxBinaryResponseContent)
assert result.response.content == b'{"line": 1}\n{"line": 2}\n'
class TestTransformDeleteFile:
def test_should_build_correct_gcs_delete_url(self, config):
file_id = "gs://my-bucket/path/to/file.jsonl"
url, params = config.transform_delete_file_request(
file_id=file_id, optional_params={}, litellm_params={}
)
encoded = urllib.parse.quote("path/to/file.jsonl", safe="")
assert url == f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded}"
assert params == {}
def test_should_return_file_deleted_with_reconstructed_id(self, config):
raw_response = MagicMock(spec=httpx.Response)
mock_request = MagicMock()
encoded_name = urllib.parse.quote(
"litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc", safe=""
)
mock_request.url = (
f"https://storage.googleapis.com/storage/v1/b/my-bucket/o/{encoded_name}"
)
raw_response.request = mock_request
result = config.transform_delete_file_response(
raw_response=raw_response,
logging_obj=MagicMock(),
litellm_params={},
)
assert isinstance(result, FileDeleted)
assert result.deleted is True
assert result.object == "file"
assert "litellm-vertex-files/publishers/google/models/gemini-2.0-flash-001/abc" in result.id
def test_should_fallback_to_deleted_id_when_no_request(self, config):
raw_response = MagicMock(spec=httpx.Response)
raw_response.request = None
result = config.transform_delete_file_response(
raw_response=raw_response,
logging_obj=MagicMock(),
litellm_params={},
)
assert isinstance(result, FileDeleted)
assert result.id == "deleted"
assert result.deleted is True
tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py
+74 -145
View File
@@ -227,52 +227,29 @@ class TestVertexAIBatchPassthroughHandler:
mock_managed_files_hook.store_unified_object_id.assert_called_once()
def test_batch_cost_calculation_integration(self):
"""Test integration with batch cost calculation"""
"""Single Vertex AI response → non-zero cost with correct token counts."""
from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage
# Mock Vertex AI batch responses
vertex_ai_batch_responses = [
{
"status": "JOB_STATE_SUCCEEDED",
"response": {
"candidates": [
{
"content": {
"parts": [
{"text": "Hello, world!"}
]
}
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 5,
"totalTokenCount": 15
"totalTokenCount": 15,
}
}
}
]
with patch('litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexGeminiConfig') as mock_config:
with patch('litellm.completion_cost') as mock_completion_cost:
# Setup mocks
mock_config.return_value._transform_google_generate_content_to_openai_model_response.return_value = Mock(
usage=Mock(total_tokens=15, prompt_tokens=10, completion_tokens=5)
)
mock_completion_cost.return_value = 0.001
# Test the cost calculation
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses,
model_name="gemini-1.5-flash"
)
# Verify results
assert total_cost == 0.001
assert usage.total_tokens == 15
assert usage.prompt_tokens == 10
assert usage.completion_tokens == 5
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses, model_name="gemini-1.5-flash-001"
)
assert usage.total_tokens == 15
assert usage.prompt_tokens == 10
assert usage.completion_tokens == 5
assert total_cost > 0, "batch_cost_calculator should return a non-zero cost"
def test_batch_response_transformation(self):
"""Test transformation of Vertex AI batch responses to OpenAI format"""
@@ -385,155 +362,107 @@ class TestVertexAIBatchPassthroughHandler:
class TestVertexAIBatchCostCalculation:
"""Test cases for Vertex AI batch cost calculation functionality"""
"""Test cases for Vertex AI batch cost calculation functionality.
def test_calculate_vertex_ai_batch_cost_and_usage_success(self):
"""Test successful batch cost and usage calculation"""
The function under test (calculate_vertex_ai_batch_cost_and_usage) extracts
usageMetadata directly from Vertex AI response dicts and calls
batch_cost_calculator — no VertexGeminiConfig transformation involved.
"""
def test_should_aggregate_cost_and_usage_across_responses(self):
"""Two successful responses → costs and token counts are summed."""
from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage
# Mock successful batch responses
vertex_ai_batch_responses = [
responses = [
{
"status": "JOB_STATE_SUCCEEDED",
"response": {
"candidates": [
{
"content": {
"parts": [
{"text": "Hello, world!"}
]
}
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 5,
"totalTokenCount": 15
"totalTokenCount": 15,
}
}
},
{
"status": "JOB_STATE_SUCCEEDED",
"response": {
"candidates": [
{
"content": {
"parts": [
{"text": "How are you?"}
]
}
}
],
"usageMetadata": {
"promptTokenCount": 8,
"candidatesTokenCount": 3,
"totalTokenCount": 11
"totalTokenCount": 11,
}
}
}
},
]
with patch('litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexGeminiConfig') as mock_config:
with patch('litellm.completion_cost') as mock_completion_cost:
# Setup mocks
mock_model_response = Mock()
mock_model_response.usage = Mock(total_tokens=15, prompt_tokens=10, completion_tokens=5)
mock_config.return_value._transform_google_generate_content_to_openai_model_response.return_value = mock_model_response
mock_completion_cost.return_value = 0.001
# Test the calculation
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses,
model_name="gemini-1.5-flash"
)
# Verify results
assert total_cost == 0.002 # 2 responses * 0.001 each
assert usage.total_tokens == 30 # 15 + 15
assert usage.prompt_tokens == 20 # 10 + 10
assert usage.completion_tokens == 10 # 5 + 5
def test_calculate_vertex_ai_batch_cost_and_usage_with_failed_responses(self):
"""Test batch cost calculation with some failed responses"""
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
responses, model_name="gemini-1.5-flash-001"
)
assert usage.prompt_tokens == 18
assert usage.completion_tokens == 8
assert usage.total_tokens == 26
assert total_cost > 0, "batch_cost_calculator should return a non-zero cost"
def test_should_skip_responses_with_null_response_body(self):
"""Failed lines (response: None) are skipped without error."""
from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage
# Mock batch responses with some failures
vertex_ai_batch_responses = [
responses = [
{
"status": "JOB_STATE_SUCCEEDED",
"response": {
"candidates": [
{
"content": {
"parts": [
{"text": "Hello, world!"}
]
}
}
],
"usageMetadata": {
"promptTokenCount": 10,
"candidatesTokenCount": 5,
"totalTokenCount": 15
"totalTokenCount": 15,
}
}
},
{"status": "JOB_STATE_FAILED", "response": None},
{
"status": "JOB_STATE_FAILED", # Failed response
"response": None
},
{
"status": "JOB_STATE_SUCCEEDED",
"response": {
"candidates": [
{
"content": {
"parts": [
{"text": "How are you?"}
]
}
}
],
"usageMetadata": {
"promptTokenCount": 8,
"candidatesTokenCount": 3,
"totalTokenCount": 11
"totalTokenCount": 11,
}
}
}
},
]
with patch('litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini.VertexGeminiConfig') as mock_config:
with patch('litellm.completion_cost') as mock_completion_cost:
# Setup mocks
mock_model_response = Mock()
mock_model_response.usage = Mock(total_tokens=15, prompt_tokens=10, completion_tokens=5)
mock_config.return_value._transform_google_generate_content_to_openai_model_response.return_value = mock_model_response
mock_completion_cost.return_value = 0.001
# Test the calculation
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses,
model_name="gemini-1.5-flash"
)
# Verify results - should only process successful responses
assert total_cost == 0.002 # 2 successful responses * 0.001 each
assert usage.total_tokens == 30 # 15 + 15
assert usage.prompt_tokens == 20 # 10 + 10
assert usage.completion_tokens == 10 # 5 + 5
def test_calculate_vertex_ai_batch_cost_and_usage_empty_responses(self):
"""Test batch cost calculation with empty response list"""
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
responses, model_name="gemini-1.5-flash-001"
)
assert usage.prompt_tokens == 18
assert usage.completion_tokens == 8
assert usage.total_tokens == 26
assert total_cost > 0
def test_should_return_zeros_for_empty_response_list(self):
"""Empty input → zero cost and zero usage."""
from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage
# Test with empty list
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage([], model_name="gemini-1.5-flash")
# Verify results
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
[], model_name="gemini-1.5-flash-001"
)
assert total_cost == 0.0
assert usage.total_tokens == 0
assert usage.prompt_tokens == 0
assert usage.completion_tokens == 0
def test_should_handle_missing_usage_metadata_gracefully(self):
"""Response without usageMetadata → 0 tokens, 0 cost for that line."""
from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage
responses = [
{"response": {"candidates": [{"content": {"parts": [{"text": "hi"}]}}]}},
]
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
responses, model_name="gemini-1.5-flash-001"
)
assert usage.prompt_tokens == 0
assert usage.completion_tokens == 0
assert usage.total_tokens == 0