Merge pull request #19070 from BerriAI/litellm_19066-bug-gemini-image-generation-returns-incorrect-prompt_tokens_details

Fix: [Bug]: Gemini Image Generation Returns Incorrect prompt_tokens_d…
2026-08-23 16:24:19 +00:00 · 2026-01-14 17:55:09 +05:30
parent 49d873e894 6991342dc4
commit 26d8ef734c
2 changed files with 85 additions and 18 deletions
@@ -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(
@@ -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)"