From a0d4d0b304b52fd0121da1468b14e926cdfd23e4 Mon Sep 17 00:00:00 2001 From: Dima-Mediator Date: Fri, 21 Nov 2025 22:59:24 -0500 Subject: [PATCH] Gemini models: capture image_tokens and support cost_per_output_image_token in costs calculations (#16912) --- .../litellm_core_utils/llm_cost_calc/utils.py | 23 +++++ .../vertex_and_google_ai_studio_gemini.py | 28 ++++++ litellm/types/utils.py | 38 ++++++-- litellm/utils.py | 1 + .../llm_cost_calc/test_llm_cost_calc_utils.py | 92 +++++++++++++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 80 ++++++++++++++++ tests/test_litellm/test_utils.py | 1 + tests/test_litellm/types/test_types_utils.py | 3 +- 8 files changed, 257 insertions(+), 9 deletions(-) diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index eff5376e49..9717f442b8 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -408,6 +408,7 @@ class CompletionTokensDetailsResult(TypedDict): audio_tokens: int text_tokens: int reasoning_tokens: int + image_tokens: int def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResult: @@ -432,11 +433,19 @@ def _parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsRes ) or 0 ) + image_tokens = ( + cast( + Optional[int], + getattr(usage.completion_tokens_details, "image_tokens", 0), + ) + or 0 + ) return CompletionTokensDetailsResult( audio_tokens=audio_tokens, text_tokens=text_tokens, reasoning_tokens=reasoning_tokens, + image_tokens=image_tokens, ) @@ -565,12 +574,14 @@ def generic_cost_per_token( text_tokens = 0 audio_tokens = 0 reasoning_tokens = 0 + image_tokens = 0 is_text_tokens_total = False if usage.completion_tokens_details is not None: completion_tokens_details = _parse_completion_tokens_details(usage) audio_tokens = completion_tokens_details["audio_tokens"] text_tokens = completion_tokens_details["text_tokens"] reasoning_tokens = completion_tokens_details["reasoning_tokens"] + image_tokens = completion_tokens_details["image_tokens"] if text_tokens == 0: text_tokens = usage.completion_tokens @@ -585,6 +596,9 @@ def generic_cost_per_token( _output_cost_per_reasoning_token = _get_cost_per_unit( model_info, "output_cost_per_reasoning_token", None ) + _output_cost_per_image_token = _get_cost_per_unit( + model_info, "output_cost_per_image_token", None + ) ## AUDIO COST if not is_text_tokens_total and audio_tokens is not None and audio_tokens > 0: @@ -604,6 +618,15 @@ def generic_cost_per_token( ) completion_cost += float(reasoning_tokens) * _output_cost_per_reasoning_token + ## IMAGE COST + if not is_text_tokens_total and image_tokens and image_tokens > 0: + _output_cost_per_image_token = ( + _output_cost_per_image_token + if _output_cost_per_image_token is not None + else completion_base_cost + ) + completion_cost += float(image_tokens) * _output_cost_per_image_token + return prompt_cost, completion_cost 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 47e6af10ee..02e4682b3a 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 @@ -1379,6 +1379,30 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): response_tokens_details.audio_tokens = detail.get("tokenCount", 0) ######################################################### + ## CANDIDATES TOKEN DETAILS (e.g., for image generation models) ## + if "candidatesTokensDetails" in usage_metadata: + if response_tokens_details is None: + response_tokens_details = CompletionTokensDetailsWrapper() + for detail in usage_metadata["candidatesTokensDetails"]: + modality = detail.get("modality") + token_count = detail.get("tokenCount", 0) + if modality == "TEXT": + response_tokens_details.text_tokens = token_count + elif modality == "AUDIO": + response_tokens_details.audio_tokens = token_count + elif modality == "IMAGE": + response_tokens_details.image_tokens = token_count + + # Calculate text_tokens if not explicitly provided in candidatesTokensDetails + # candidatesTokenCount includes all modalities, so: text = total - (image + audio) + if response_tokens_details.text_tokens is None: + candidates_token_count = usage_metadata.get("candidatesTokenCount", 0) + image_tokens = response_tokens_details.image_tokens or 0 + audio_tokens_candidate = response_tokens_details.audio_tokens or 0 + calculated_text_tokens = candidates_token_count - image_tokens - audio_tokens_candidate + response_tokens_details.text_tokens = calculated_text_tokens + ######################################################### + if "promptTokensDetails" in usage_metadata: for detail in usage_metadata["promptTokensDetails"]: if detail["modality"] == "AUDIO": @@ -1387,6 +1411,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): text_tokens = detail.get("tokenCount", 0) if "thoughtsTokenCount" in usage_metadata: reasoning_tokens = usage_metadata["thoughtsTokenCount"] + # Also add reasoning tokens to response_tokens_details + if response_tokens_details is None: + response_tokens_details = CompletionTokensDetailsWrapper() + response_tokens_details.reasoning_tokens = reasoning_tokens ## adjust 'text_tokens' to subtract cached tokens if ( diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a751adf542..c868af85d2 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -162,6 +162,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): float ] # only for vertex ai models output_cost_per_image: Optional[float] + output_cost_per_image_token: Optional[float] output_vector_size: Optional[int] output_cost_per_reasoning_token: Optional[float] output_cost_per_video_per_second: Optional[float] # only for vertex ai models @@ -945,6 +946,9 @@ class CompletionTokensDetailsWrapper( text_tokens: Optional[int] = None """Text tokens generated by the model.""" + image_tokens: Optional[int] = None + """Image tokens generated by the model.""" + class CacheCreationTokenDetails(BaseModel): ephemeral_5m_input_tokens: Optional[int] = None @@ -1033,15 +1037,8 @@ class Usage(CompletionUsage): ): # handle reasoning_tokens _completion_tokens_details: Optional[CompletionTokensDetailsWrapper] = None - if reasoning_tokens: - text_tokens = ( - completion_tokens - reasoning_tokens if completion_tokens else None - ) - completion_tokens_details = CompletionTokensDetailsWrapper( - reasoning_tokens=reasoning_tokens, text_tokens=text_tokens - ) - # Ensure completion_tokens_details is properly handled + # First, handle existing completion_tokens_details if completion_tokens_details: if isinstance(completion_tokens_details, dict): _completion_tokens_details = CompletionTokensDetailsWrapper( @@ -1050,6 +1047,30 @@ class Usage(CompletionUsage): elif isinstance(completion_tokens_details, CompletionTokensDetails): _completion_tokens_details = completion_tokens_details + # Handle reasoning_tokens and auto-calculate text_tokens if needed + if reasoning_tokens: + # Ensure we have a details object to work with + if _completion_tokens_details is None: + _completion_tokens_details = CompletionTokensDetailsWrapper() + + # Set reasoning_tokens if not already set by provider + if _completion_tokens_details.reasoning_tokens is None: + _completion_tokens_details.reasoning_tokens = reasoning_tokens + + # Auto-calculate text_tokens only if provider didn't set it explicitly + # Formula: text_tokens = completion_tokens - reasoning_tokens - image_tokens - audio_tokens + if _completion_tokens_details.text_tokens is None and completion_tokens is not None: + calculated_text_tokens = completion_tokens - reasoning_tokens + + # Subtract other modality tokens if present + if _completion_tokens_details.image_tokens: + calculated_text_tokens -= _completion_tokens_details.image_tokens + if _completion_tokens_details.audio_tokens: + calculated_text_tokens -= _completion_tokens_details.audio_tokens + + # Prevent negative token counts from inconsistent data + _completion_tokens_details.text_tokens = max(0, calculated_text_tokens) + # handle prompt_tokens_details _prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None @@ -2370,6 +2391,7 @@ class CustomPricingLiteLLMParams(BaseModel): output_cost_per_token_above_200k_tokens: Optional[float] = None output_cost_per_character_above_128k_tokens: Optional[float] = None output_cost_per_image: Optional[float] = None + output_cost_per_image_token: Optional[float] = None output_cost_per_reasoning_token: Optional[float] = None output_cost_per_video_per_second: Optional[float] = None output_cost_per_audio_per_second: Optional[float] = None diff --git a/litellm/utils.py b/litellm/utils.py index 474f5d57b1..d7474387cd 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5025,6 +5025,7 @@ def _get_model_info_helper( # noqa: PLR0915 "output_cost_per_video_per_second", None ), output_cost_per_image=_model_info.get("output_cost_per_image", None), + output_cost_per_image_token=_model_info.get("output_cost_per_image_token", None), output_vector_size=_model_info.get("output_vector_size", None), citation_cost_per_token=_model_info.get( "citation_cost_per_token", None diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 664ecb896b..3d3dd98f49 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -115,6 +115,98 @@ def test_reasoning_tokens_gemini(): ) +def test_image_tokens_with_custom_pricing(): + """Test that image_tokens in completion are properly costed with output_cost_per_image_token.""" + from unittest.mock import patch + + # Mock model info with image token pricing + mock_model_info = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "output_cost_per_image_token": 5e-6, # Custom pricing for image tokens in output + } + + usage = Usage( + completion_tokens=1720, # text_tokens (600) + image_tokens (1120) + prompt_tokens=14, + total_tokens=1734, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=None, + audio_tokens=None, + reasoning_tokens=0, + rejected_prediction_tokens=None, + text_tokens=600, + image_tokens=1120, + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=None, cached_tokens=None, text_tokens=14, image_tokens=None + ), + ) + + with patch( + "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", + return_value=mock_model_info, + ): + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", usage=usage, custom_llm_provider="gemini" + ) + + # Expected costs: + # Prompt: 14 * 1e-6 + # Completion: (600 * 2e-6) + (1120 * 5e-6) + expected_prompt_cost = 14 * 1e-6 + expected_completion_cost = (600 * 2e-6) + (1120 * 5e-6) + + assert round(prompt_cost, 12) == round(expected_prompt_cost, 12) + assert round(completion_cost, 12) == round(expected_completion_cost, 12) + + +def test_image_tokens_fallback_to_base_cost(): + """Test that image_tokens fall back to base cost when output_cost_per_image_token is not set.""" + from unittest.mock import patch + + # Mock model info without image token pricing + mock_model_info = { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + # No output_cost_per_image_token defined + } + + usage = Usage( + completion_tokens=1720, + prompt_tokens=14, + total_tokens=1734, + completion_tokens_details=CompletionTokensDetailsWrapper( + accepted_prediction_tokens=None, + audio_tokens=None, + reasoning_tokens=0, + rejected_prediction_tokens=None, + text_tokens=600, + image_tokens=1120, + ), + prompt_tokens_details=PromptTokensDetailsWrapper( + audio_tokens=None, cached_tokens=None, text_tokens=14, image_tokens=None + ), + ) + + with patch( + "litellm.litellm_core_utils.llm_cost_calc.utils.get_model_info", + return_value=mock_model_info, + ): + prompt_cost, completion_cost = generic_cost_per_token( + model="test-model", usage=usage, custom_llm_provider="gemini" + ) + + # Expected costs: + # Prompt: 14 * 1e-6 + # Completion: (600 * 2e-6) + (1120 * 2e-6) # image_tokens use base cost + expected_prompt_cost = 14 * 1e-6 + expected_completion_cost = (600 * 2e-6) + (1120 * 2e-6) + + assert round(prompt_cost, 12) == round(expected_prompt_cost, 12) + assert round(completion_cost, 12) == round(expected_completion_cost, 12) + + def test_generic_cost_per_token_above_200k_tokens(): model = "gemini-2.5-pro-exp-03-25" custom_llm_provider = "vertex_ai" diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index d7194033c6..35e0cd4bec 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -477,6 +477,86 @@ def test_vertex_ai_usage_metadata_response_token_count(): assert result.completion_tokens_details.text_tokens == 74 +def test_vertex_ai_usage_metadata_with_image_tokens(): + """Test candidatesTokensDetails with IMAGE modality (e.g., Imagen models) + + This test simulates the case where candidatesTokenCount is EXCLUSIVE of thoughtsTokenCount. + Gemini API returns: totalTokenCount = promptTokenCount + candidatesTokenCount + thoughtsTokenCount + """ + v = VertexGeminiConfig() + usage_metadata = { + "promptTokenCount": 14, + "candidatesTokenCount": 1442, # Does NOT include thoughtsTokenCount + "totalTokenCount": 1614, # 14 + 1442 + 158 + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 14}], + "candidatesTokensDetails": [ + {"modality": "IMAGE", "tokenCount": 1120}, + {"modality": "TEXT", "tokenCount": 322} # 1442 - 1120 = 322 + ], + "thoughtsTokenCount": 158 + } + usage_metadata = UsageMetadata(**usage_metadata) + result = v._calculate_usage(completion_response={"usageMetadata": usage_metadata}) + print("result", result) + + # Verify basic token counts + assert result.prompt_tokens == 14 + # completion_tokens = candidatesTokenCount + thoughtsTokenCount (when exclusive) + assert result.completion_tokens == 1600 # 1442 + 158 + assert result.total_tokens == 1614 + + # Verify detailed token breakdown + assert result.completion_tokens_details.image_tokens == 1120 + assert result.completion_tokens_details.text_tokens == 322 + assert result.completion_tokens_details.reasoning_tokens == 158 + + # Verify the math: completion_tokens = image + text + reasoning + # 1600 = 1120 (image) + 322 (text) + 158 (reasoning) + assert ( + result.completion_tokens_details.image_tokens + + result.completion_tokens_details.text_tokens + + result.completion_tokens_details.reasoning_tokens + == result.completion_tokens + ) + + +def test_vertex_ai_usage_metadata_with_image_tokens_auto_calculated_text(): + """Test that text_tokens is auto-calculated when only IMAGE modality is provided + + This test verifies the auto-calculation logic at line 1367-1372 in vertex_and_google_ai_studio_gemini.py + """ + v = VertexGeminiConfig() + usage_metadata = { + "promptTokenCount": 14, + "candidatesTokenCount": 1442, + "totalTokenCount": 1614, # 14 + 1442 + 158 (exclusive) + "promptTokensDetails": [{"modality": "TEXT", "tokenCount": 14}], + "candidatesTokensDetails": [ + {"modality": "IMAGE", "tokenCount": 1120} + # TEXT modality omitted - should be auto-calculated + ], + "thoughtsTokenCount": 158 + } + usage_metadata = UsageMetadata(**usage_metadata) + result = v._calculate_usage(completion_response={"usageMetadata": usage_metadata}) + print("result", result) + + # Verify basic token counts + assert result.prompt_tokens == 14 + assert result.completion_tokens == 1600 # 1442 + 158 + assert result.total_tokens == 1614 + + # Verify image_tokens is set + assert result.completion_tokens_details.image_tokens == 1120 + + # Verify text_tokens is auto-calculated: candidatesTokenCount - image_tokens + # Note: reasoning_tokens is NOT subtracted here because candidatesTokenCount is exclusive + # 1442 - 1120 = 322 + expected_text_tokens = 1442 - 1120 + assert result.completion_tokens_details.text_tokens == expected_text_tokens + assert result.completion_tokens_details.reasoning_tokens == 158 + + def test_vertex_ai_map_thinking_param_with_budget_tokens_0(): """ If budget_tokens is 0, do not set includeThoughts to True diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index b30b048995..c8db2c6c74 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -416,6 +416,7 @@ def validate_model_cost_values(model_data, exceptions=None): "input_cost_per_request", "input_cost_per_audio_token", "output_cost_per_audio_token", + "output_cost_per_image_token", "input_cost_per_audio_per_second", "input_cost_per_video_per_second", "input_cost_per_token_above_128k_tokens", diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index 71b9247518..a2dfb5268e 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -119,7 +119,8 @@ def test_usage_completion_tokens_details_text_tokens(): 'audio_tokens': None, 'reasoning_tokens': 65, 'rejected_prediction_tokens': None, - 'text_tokens': 12 + 'text_tokens': 12, + 'image_tokens': None } assert dump_result['completion_tokens_details'] == expected_completion_details