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 7a37812bcf..f65a19ac46 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 @@ -1535,9 +1535,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): f"usageMetadata not found in completion_response. Got={completion_response}" ) cached_tokens: Optional[int] = None - audio_tokens: Optional[int] = None - image_tokens: Optional[int] = None - text_tokens: Optional[int] = None + # Separate variables for prompt tokens by modality + prompt_audio_tokens: Optional[int] = None + prompt_image_tokens: Optional[int] = None + prompt_text_tokens: Optional[int] = None prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None reasoning_tokens: Optional[int] = None response_tokens: Optional[int] = None @@ -1580,10 +1581,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if response_tokens_details is None: response_tokens_details = CompletionTokensDetailsWrapper() if response_tokens_details.text_tokens is None: - image_tokens = response_tokens_details.image_tokens or 0 - audio_tokens_candidate = response_tokens_details.audio_tokens or 0 + completion_image_tokens = response_tokens_details.image_tokens or 0 + completion_audio_tokens = response_tokens_details.audio_tokens or 0 calculated_text_tokens = ( - candidates_token_count - image_tokens - audio_tokens_candidate + candidates_token_count - completion_image_tokens - completion_audio_tokens ) response_tokens_details.text_tokens = calculated_text_tokens ######################################################### @@ -1592,11 +1593,11 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): if "promptTokensDetails" in usage_metadata: for detail in usage_metadata["promptTokensDetails"]: if detail["modality"] == "AUDIO": - audio_tokens = detail.get("tokenCount", 0) + prompt_audio_tokens = detail.get("tokenCount", 0) elif detail["modality"] == "TEXT": - text_tokens = detail.get("tokenCount", 0) + prompt_text_tokens = detail.get("tokenCount", 0) elif detail["modality"] == "IMAGE": - image_tokens = detail.get("tokenCount", 0) + prompt_image_tokens = detail.get("tokenCount", 0) ## Parse cacheTokensDetails (breakdown of cached tokens by modality) ## When explicit caching is used, Gemini provides this field to show which modalities were cached @@ -1616,12 +1617,12 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ## Calculate non-cached tokens by subtracting cached from total (per modality) ## This is necessary because promptTokensDetails includes both cached and non-cached tokens ## See: https://github.com/BerriAI/litellm/issues/18750 - if cached_text_tokens is not None and text_tokens is not None: - text_tokens = text_tokens - cached_text_tokens - if cached_audio_tokens is not None and audio_tokens is not None: - audio_tokens = audio_tokens - cached_audio_tokens - if cached_image_tokens is not None and image_tokens is not None: - image_tokens = image_tokens - cached_image_tokens + if cached_text_tokens is not None and prompt_text_tokens is not None: + prompt_text_tokens = prompt_text_tokens - cached_text_tokens + if cached_audio_tokens is not None and prompt_audio_tokens is not None: + prompt_audio_tokens = prompt_audio_tokens - cached_audio_tokens + if cached_image_tokens is not None and prompt_image_tokens is not None: + prompt_image_tokens = prompt_image_tokens - cached_image_tokens if "thoughtsTokenCount" in usage_metadata: reasoning_tokens = usage_metadata["thoughtsTokenCount"] @@ -1632,9 +1633,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): prompt_tokens_details = PromptTokensDetailsWrapper( cached_tokens=cached_tokens, - audio_tokens=audio_tokens, - text_tokens=text_tokens, - image_tokens=image_tokens, + audio_tokens=prompt_audio_tokens, + text_tokens=prompt_text_tokens, + image_tokens=prompt_image_tokens, ) completion_tokens = response_tokens or completion_response["usageMetadata"].get( 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 c0b80ae771..f44e864010 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 @@ -2587,3 +2587,69 @@ def test_gemini_token_usage_standard_response(): assert result.completion_tokens == 50 assert result.completion_tokens_details.text_tokens == 40 assert result.completion_tokens_details.image_tokens == 10 + + +def test_gemini_image_gen_usage_metadata_prompt_vs_completion_separation(): + """ + Test that image generation models correctly separate prompt and completion token details. + + This is a regression test for the bug where prompt_tokens_details.image_tokens + was incorrectly set to the completion's image token count instead of 0. + + Scenario: Text-only prompt generates an image response + - Input: Text prompt (no images) + - Output: Generated image + text description + + Expected behavior: + - prompt_tokens_details.image_tokens should be 0 (text-only input) + - completion_tokens_details.image_tokens should be 1290 (generated image) + + Bug behavior (before fix): + - prompt_tokens_details.image_tokens was 1290 (incorrect!) + - completion_tokens_details.image_tokens was 1290 (correct) + + The bug was caused by reusing the same variables (image_tokens, audio_tokens, text_tokens) + for both prompt and completion token details. + """ + v = VertexGeminiConfig() + + # Simulate Gemini image generation model response metadata + # User sends text-only prompt, model generates image + text + usage_metadata_dict = { + "promptTokenCount": 101, + "candidatesTokenCount": 1290, + "totalTokenCount": 1391, + # Prompt is text-only (no image tokens in input) + "promptTokensDetails": [ + {"modality": "TEXT", "tokenCount": 101} + ], + # Response contains generated image + text + "candidatesTokensDetails": [ + {"modality": "IMAGE", "tokenCount": 1290} + ], + } + + completion_response = {"usageMetadata": usage_metadata_dict} + result = v._calculate_usage(completion_response=completion_response) + + # Verify basic token counts + assert result.prompt_tokens == 101 + assert result.completion_tokens == 1290 + assert result.total_tokens == 1391 + + # CRITICAL: Prompt tokens details should show NO image tokens (text-only input) + assert result.prompt_tokens_details.text_tokens == 101, \ + "Prompt text tokens should be 101" + assert result.prompt_tokens_details.image_tokens is None, \ + "Prompt image tokens should be None (text-only input, no images in prompt)" + assert result.prompt_tokens_details.audio_tokens is None, \ + "Prompt audio tokens should be None" + + # Completion tokens details should show the generated image tokens + assert result.completion_tokens_details.image_tokens == 1290, \ + "Completion image tokens should be 1290 (generated image)" + + # Verify text_tokens is auto-calculated for completion + # candidatesTokenCount (1290) - image_tokens (1290) = 0 + assert result.completion_tokens_details.text_tokens == 0, \ + "Completion text tokens should be 0 (image-only response)"