From 35f6fd4223065eb2c1710964ba43c60b25f9edff Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Sat, 28 Feb 2026 14:42:24 -0500 Subject: [PATCH 1/2] Managed batches fixes for Gemini/Vertex --- .../proxy/hooks/managed_files.py | 7 ++ litellm/batches/batch_utils.py | 101 ++++++++---------- litellm/files/main.py | 2 +- litellm/llms/vertex_ai/batches/handler.py | 17 ++- .../llms/vertex_ai/batches/transformation.py | 2 +- .../llms/vertex_ai/files/transformation.py | 62 +++++++++-- litellm/types/llms/vertex_ai.py | 2 +- 7 files changed, 121 insertions(+), 72 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 4fa050a84a..37ca341fdf 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -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( diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 29bd99c2a6..80351664df 100644 --- a/litellm/batches/batch_utils.py +++ b/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 - - 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, diff --git a/litellm/files/main.py b/litellm/files/main.py index 78e41bb5a6..66d3a97468 100644 --- a/litellm/files/main.py +++ b/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, diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index 36f5e65e7a..ba3b5fb7a2 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -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}") diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index a0adb3e55a..7cb06fea9e 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -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( diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 2470c59bba..f0493cd6be 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -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): diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 049a5010c7..190e680b7b 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -560,7 +560,7 @@ class VertexAIBatchEmbeddingsResponseObject(TypedDict): class GcsSource(TypedDict): - uris: str + uris: List[str] class InputConfig(TypedDict): From b16397ae1ab6d91472323673f990db6a0f3bf80f Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Sat, 28 Feb 2026 20:45:16 -0500 Subject: [PATCH 2/2] Managed batches fixes for Gemini/Vertex --- .../test_openai_batches_and_files.py | 43 ++++ .../proxy/hooks/test_managed_files.py | 104 +++++++++ .../test_vertex_ai_files_transformation.py | 184 +++++++++++++++ .../test_vertex_ai_batch_passthrough.py | 219 ++++++------------ 4 files changed, 405 insertions(+), 145 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py diff --git a/tests/batches_tests/test_openai_batches_and_files.py b/tests/batches_tests/test_openai_batches_and_files.py index 055af02494..641590ad04 100644 --- a/tests/batches_tests/test_openai_batches_and_files.py +++ b/tests/batches_tests/test_openai_batches_and_files.py @@ -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(): """ diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 58efa854e7..58fbd9e64b 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -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()) diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py new file mode 100644 index 0000000000..6f1d753484 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -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 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py index 66c063d47d..c2f6d3fd53 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_ai_batch_passthrough.py @@ -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