From 241f32b2a330e6cbe7e224c1a75cac3824ea1619 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 18 Aug 2025 22:09:07 -0700 Subject: [PATCH 1/2] fix(vertex_and_google_ai_studio_gemini.py): adjust 'text token' value to be the non-cached tokens - for accurate cost tracking --- .../vertex_and_google_ai_studio_gemini.py | 55 +++++++++++-------- 1 file changed, 33 insertions(+), 22 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 37a4ab84dd..c5c4f46c92 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -305,9 +305,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return None for tool in value: - openai_function_object: Optional[ - ChatCompletionToolParamFunctionChunk - ] = None + openai_function_object: Optional[ChatCompletionToolParamFunctionChunk] = ( + None + ) if "function" in tool: # tools list _openai_function_object = ChatCompletionToolParamFunctionChunk( # type: ignore **tool["function"] @@ -597,14 +597,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif param == "seed": optional_params["seed"] = value elif param == "reasoning_effort" and isinstance(value, str): - optional_params[ - "thinkingConfig" - ] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(value) + optional_params["thinkingConfig"] = ( + VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(value) + ) elif param == "thinking": - optional_params[ - "thinkingConfig" - ] = VertexGeminiConfig._map_thinking_param( - cast(AnthropicThinkingParam, value) + optional_params["thinkingConfig"] = ( + VertexGeminiConfig._map_thinking_param( + cast(AnthropicThinkingParam, value) + ) ) elif param == "modalities" and isinstance(value, list): response_modalities = self.map_response_modalities(value) @@ -1000,6 +1000,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): GenerateContentResponseBody, BidiGenerateContentServerMessage ], ) -> Usage: + if ( completion_response is not None and "usageMetadata" not in completion_response @@ -1038,6 +1039,16 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): text_tokens = detail.get("tokenCount", 0) if "thoughtsTokenCount" in usage_metadata: reasoning_tokens = usage_metadata["thoughtsTokenCount"] + + ## adjust 'text_tokens' to subtract cached tokens + if ( + (audio_tokens is None or audio_tokens == 0) + and text_tokens is not None + and text_tokens > 0 + and cached_tokens is not None + ): + text_tokens = text_tokens - cached_tokens + prompt_tokens_details = PromptTokensDetailsWrapper( cached_tokens=cached_tokens, audio_tokens=audio_tokens, @@ -1344,28 +1355,28 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ## ADD METADATA TO RESPONSE ## setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) - model_response._hidden_params[ - "vertex_ai_grounding_metadata" - ] = grounding_metadata + model_response._hidden_params["vertex_ai_grounding_metadata"] = ( + grounding_metadata + ) setattr( model_response, "vertex_ai_url_context_metadata", url_context_metadata ) - model_response._hidden_params[ - "vertex_ai_url_context_metadata" - ] = url_context_metadata + model_response._hidden_params["vertex_ai_url_context_metadata"] = ( + url_context_metadata + ) setattr(model_response, "vertex_ai_safety_results", safety_ratings) - model_response._hidden_params[ - "vertex_ai_safety_results" - ] = safety_ratings # older approach - maintaining to prevent regressions + model_response._hidden_params["vertex_ai_safety_results"] = ( + safety_ratings # older approach - maintaining to prevent regressions + ) ## ADD CITATION METADATA ## setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) - model_response._hidden_params[ - "vertex_ai_citation_metadata" - ] = citation_metadata # older approach - maintaining to prevent regressions + model_response._hidden_params["vertex_ai_citation_metadata"] = ( + citation_metadata # older approach - maintaining to prevent regressions + ) except Exception as e: raise VertexAIError( From f7f1a0d0b77a528172f35c18b4e9f8a7de43283e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 18 Aug 2025 22:32:36 -0700 Subject: [PATCH 2/2] test: add unit test --- tests/test_litellm/test_cost_calculator.py | 66 ++++++++++++++++++++++ 1 file changed, 66 insertions(+) diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 4efba133f1..9bff907700 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -480,3 +480,69 @@ def test_gemini_25_implicit_caching_cost(): ), f"Expected cost {expected_cost}, but got {result}" print(f"✓ Gemini 2.5 implicit caching cost calculation is correct: ${result:.8f}") + + +def test_gemini_25_explicit_caching_cost_direct_usage(): + """ + Test that Gemini 2.5 models correctly calculate costs with explicit caching. + + This test reproduces the issue from #11156 where cached tokens should receive + a 75% discount. + """ + from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token + from litellm.types.utils import ( + CompletionTokensDetailsWrapper, + PromptTokensDetailsWrapper, + Usage, + ) + from litellm.utils import get_model_info + + model_info = get_model_info(model="gemini-2.5-pro", custom_llm_provider="gemini") + + usage = Usage( + completion_tokens=2522, + prompt_tokens=42001, + total_tokens=44523, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=None, + audio_tokens=None, + reasoning_tokens=1908, + rejected_prediction_tokens=None, + text_tokens=614, + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=None, cached_tokens=40938, text_tokens=1063, image_tokens=None + ), + ) + + input_cost, output_cost = generic_cost_per_token( + model="gemini/gemini-2.5-pro", + usage=usage, + custom_llm_provider="gemini", + ) + + total_cost = input_cost + output_cost + + expected_higher_than_actual_cost = ( + model_info["input_cost_per_token"] * usage.prompt_tokens + + model_info["output_cost_per_token"] * usage.completion_tokens + ) + + print(f"expected_higher_than_actual_cost: {expected_higher_than_actual_cost}") + + assert expected_higher_than_actual_cost > total_cost + + expected_actual_cost = ( + model_info["input_cost_per_token"] * usage.prompt_tokens_details.text_tokens + + model_info["cache_read_input_token_cost"] + * usage.prompt_tokens_details.cached_tokens + + model_info["output_cost_per_token"] * usage.completion_tokens + ) + + print( + f"model_info['input_cost_per_token']: {model_info['input_cost_per_token']}, usage.prompt_tokens_details.text_tokens: {usage.prompt_tokens_details.text_tokens}, model_info['cache_read_input_token_cost']: {model_info['cache_read_input_token_cost']}, model_info['output_cost_per_token']: {model_info['output_cost_per_token']}" + ) + + print(f"Expected actual cost: {expected_actual_cost}") + + assert expected_actual_cost == total_cost