From c2efe9e422b6ce62f0001d847d578d1e7d7ea6e3 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 15 May 2026 17:17:02 +0530 Subject: [PATCH] fix(vertex-ai): fix zero cost/usage on completed Vertex AI batch jobs (#27912) * fix(vertex-ai): fix zero cost/usage on completed Vertex AI batch jobs Vertex batch jobs recorded 0 spend and 0 tokens after PR #25627 added automatic transformation of GCS predictions.jsonl to OpenAI format. Two bugs fixed: 1. batch_utils.py: the Vertex-specific cost/usage reader (calculate_vertex_ai_batch_cost_and_usage) was always invoked and reads raw usageMetadata fields that no longer exist in the OpenAI-shaped output. Now the reader is only used when disable_vertex_batch_output_transformation=True; otherwise the generic path handles the already-transformed OpenAI-shaped content. 2. cost_calculator.py: batch_cost_calculator skipped the global litellm.get_model_info() lookup when a model_info dict was passed in, even when that dict had no pricing fields (e.g. deployment metadata with only id/db_model). It now falls back to the global pricing table when the provided model_info has no pricing data. Co-authored-by: Cursor * Update litellm/cost_calculator.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix(cost-calculator): use not-any guard for pricing fallback in batch_cost_calculator Co-authored-by: Cursor * fix(cost-calculator): treat explicit zero batch pricing as set in model_info The fallback to litellm.get_model_info() used truthy checks on pricing fields, so 0.0 was treated as missing and replaced by global rates. Use `is not None` like elsewhere in cost calculation. Add regression test. Co-authored-by: Sameer Kankute --------- Co-authored-by: Cursor Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Sameer Kankute --- litellm/batches/batch_utils.py | 23 ++- litellm/cost_calculator.py | 20 +++ .../test_batch_custom_pricing.py | 32 ++++ .../test_vertex_ai_batch_passthrough.py | 155 ++++++++++++++++++ 4 files changed, 223 insertions(+), 7 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index aaf083e75d..74e753b09e 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -113,8 +113,11 @@ def _batch_cost_calculator( """ Calculate the cost of a batch based on the output file id """ - # Handle Vertex AI with specialized method - if custom_llm_provider == "vertex_ai" and model_name: + if ( + custom_llm_provider == "vertex_ai" + and model_name + and getattr(litellm, "disable_vertex_batch_output_transformation", False) + ): batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage( file_content_dictionary, model_name ) @@ -136,10 +139,13 @@ 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 raw Vertex AI batch responses. - Vertex AI batch output lines have format: - {"request": ..., "status": "", "response": {"candidates": [...], "usageMetadata": {...}}} + Used only when ``litellm.disable_vertex_batch_output_transformation = True``. + In that case the GCS predictions.jsonl is returned as-is, with each line in + the native Vertex format: + + {"request": ..., "response": {"candidates": [...], "usageMetadata": {...}}} usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount. """ @@ -362,8 +368,11 @@ def _get_batch_job_total_usage_from_file_content( """ Get the tokens of a batch job from the file content """ - # Handle Vertex AI with specialized method - if custom_llm_provider == "vertex_ai" and model_name: + if ( + custom_llm_provider == "vertex_ai" + and model_name + and getattr(litellm, "disable_vertex_batch_output_transformation", False) + ): _, batch_usage = calculate_vertex_ai_batch_cost_and_usage( file_content_dictionary, model_name ) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 9b4dd80265..c0ec148d03 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -2120,6 +2120,26 @@ def batch_cost_calculator( ) except Exception: model_info = None + elif not any( + model_info.get(k) is not None + for k in ( + "input_cost_per_token_batches", + "input_cost_per_token", + "output_cost_per_token_batches", + "output_cost_per_token", + ) + ): + # model_info was provided (e.g. deployment metadata with only id/db_model) + # but carries no pricing fields. Fall back to the global pricing table so + # that standard model pricing is used instead of silently returning $0. + try: + global_info = litellm.get_model_info( + model=model, custom_llm_provider=custom_llm_provider + ) + if global_info: + model_info = global_info + except Exception: + pass if not model_info: return 0.0, 0.0 diff --git a/tests/batches_tests/test_batch_custom_pricing.py b/tests/batches_tests/test_batch_custom_pricing.py index f4e84b46be..46870f1227 100644 --- a/tests/batches_tests/test_batch_custom_pricing.py +++ b/tests/batches_tests/test_batch_custom_pricing.py @@ -8,6 +8,7 @@ are ignored by the batch cost pipeline because they are never threaded through to `batch_cost_calculator`. """ +import litellm import pytest from litellm.batches.batch_utils import ( @@ -60,6 +61,37 @@ CUSTOM_MODEL_INFO = { # --- tests --- +def test_batch_cost_calculator_explicit_zero_pricing_not_overridden_by_global( + monkeypatch, +): + """ + Explicit ``0`` / ``0.0`` pricing must count as present so we do not fall back + to the global pricing table (truthiness would treat zero as missing). + """ + usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + def fake_get_model_info(*args, **kwargs): + return { + "input_cost_per_token_batches": 1e-3, + "output_cost_per_token_batches": 2e-3, + } + + monkeypatch.setattr(litellm, "get_model_info", fake_get_model_info) + + prompt_cost, completion_cost = batch_cost_calculator( + usage=usage, + model="any-model", + custom_llm_provider="openai", + model_info={ + "input_cost_per_token_batches": 0.0, + "output_cost_per_token_batches": 0.0, + }, + ) + + assert prompt_cost == 0.0 + assert completion_cost == 0.0 + + def test_batch_cost_calculator_uses_custom_model_info(): """batch_cost_calculator should use model_info override when provided.""" usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15) 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 efa26a61bf..044827e287 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 @@ -589,3 +589,158 @@ class TestVertexAIBatchCostCalculation: assert usage.prompt_tokens == 0 assert usage.completion_tokens == 0 assert usage.total_tokens == 0 + + def test_openai_shaped_output_records_nonzero_cost_and_usage(self): + """ + Regression test for the bug where Vertex batch cost/usage was always 0. + + After PR #25627 (transform_file_content_response), the GCS predictions.jsonl + is rewritten into OpenAI batch shape before the cost-tracking path sees it. + With disable_vertex_batch_output_transformation=False (default), the content + is OpenAI-shaped, so _batch_cost_calculator must fall through to the generic + path rather than calling calculate_vertex_ai_batch_cost_and_usage (which only + reads raw usageMetadata fields). + """ + import litellm + from litellm.batches.batch_utils import ( + _batch_cost_calculator, + _get_batch_job_total_usage_from_file_content, + ) + + openai_shaped_responses = [ + { + "id": "batch_req_abc123", + "custom_id": "request-1", + "response": { + "status_code": 200, + "request_id": "chatcmpl-xyz", + "body": { + "id": "chatcmpl-xyz", + "object": "chat.completion", + "model": "gemini-2.0-flash-001", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "Hello!"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15, + }, + }, + }, + "error": None, + }, + { + "id": "batch_req_def456", + "custom_id": "request-2", + "response": { + "status_code": 200, + "request_id": "chatcmpl-uvw", + "body": { + "id": "chatcmpl-uvw", + "object": "chat.completion", + "model": "gemini-2.0-flash-001", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "World!"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 8, + "completion_tokens": 3, + "total_tokens": 11, + }, + }, + }, + "error": None, + }, + ] + + original_flag = getattr( + litellm, "disable_vertex_batch_output_transformation", False + ) + try: + litellm.disable_vertex_batch_output_transformation = False + + cost = _batch_cost_calculator( + file_content_dictionary=openai_shaped_responses, + custom_llm_provider="vertex_ai", + model_name="gemini-2.0-flash-001", + ) + usage = _get_batch_job_total_usage_from_file_content( + file_content_dictionary=openai_shaped_responses, + custom_llm_provider="vertex_ai", + model_name="gemini-2.0-flash-001", + ) + finally: + litellm.disable_vertex_batch_output_transformation = original_flag + + assert ( + usage.prompt_tokens == 18 + ), f"expected 18 prompt tokens, got {usage.prompt_tokens}" + assert ( + usage.completion_tokens == 8 + ), f"expected 8 completion tokens, got {usage.completion_tokens}" + assert ( + usage.total_tokens == 26 + ), f"expected 26 total tokens, got {usage.total_tokens}" + assert ( + cost > 0 + ), f"expected non-zero cost for completed Vertex batch, got {cost}" + + def test_raw_vertex_output_still_works_when_transformation_disabled(self): + """ + When disable_vertex_batch_output_transformation=True the GCS file is returned + as raw Vertex predictions.jsonl; the specialized reader must be used. + """ + import litellm + from litellm.batches.batch_utils import ( + _batch_cost_calculator, + _get_batch_job_total_usage_from_file_content, + ) + + raw_vertex_responses = [ + { + "request": {"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}, + "status": "", + "response": { + "candidates": [{"content": {"parts": [{"text": "Hello!"}]}}], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, + }, + "processed_time": "2026-01-01T00:00:00Z", + }, + ] + + original_flag = getattr( + litellm, "disable_vertex_batch_output_transformation", False + ) + try: + litellm.disable_vertex_batch_output_transformation = True + + cost = _batch_cost_calculator( + file_content_dictionary=raw_vertex_responses, + custom_llm_provider="vertex_ai", + model_name="gemini-2.0-flash-001", + ) + usage = _get_batch_job_total_usage_from_file_content( + file_content_dictionary=raw_vertex_responses, + custom_llm_provider="vertex_ai", + model_name="gemini-2.0-flash-001", + ) + finally: + litellm.disable_vertex_batch_output_transformation = original_flag + + assert usage.prompt_tokens == 10 + assert usage.completion_tokens == 5 + assert usage.total_tokens == 15 + assert cost > 0, "raw Vertex shape should also produce non-zero cost"