From 5713eae75cff5e453b8291835a16e7364581cf73 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Mon, 22 Dec 2025 11:45:49 +0530
Subject: [PATCH 1/2] Add image tokens in usage object

---
 .../vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py    | 4 ++++
 1 file changed, 4 insertions(+)

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 84a5958ee5e..b1810b40cf9 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
@@ -1476,6 +1476,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
         cached_tokens: Optional[int] = None
         audio_tokens: Optional[int] = None
         text_tokens: Optional[int] = None
+        image_tokens: Optional[int] = None
         prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
         reasoning_tokens: Optional[int] = None
         response_tokens: Optional[int] = None
@@ -1526,6 +1527,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
                     audio_tokens = detail.get("tokenCount", 0)
                 elif detail["modality"] == "TEXT":
                     text_tokens = detail.get("tokenCount", 0)
+                elif detail["modality"] == "IMAGE":
+                    image_tokens = detail.get("tokenCount", 0)
         if "thoughtsTokenCount" in usage_metadata:
             reasoning_tokens = usage_metadata["thoughtsTokenCount"]
             # Also add reasoning tokens to response_tokens_details
@@ -1546,6 +1549,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
             cached_tokens=cached_tokens,
             audio_tokens=audio_tokens,
             text_tokens=text_tokens,
+            image_tokens=image_tokens,
         )
 
         completion_tokens = response_tokens or completion_response["usageMetadata"].get(

From 60b71d48bd207bf7edf1bc5fdb14b2311ab4c98e Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Mon, 22 Dec 2025 11:46:14 +0530
Subject: [PATCH 2/2] Add test image tokens in output

---
 ...test_vertex_and_google_ai_studio_gemini.py | 53 +++++++++++++++++++
 1 file changed, 53 insertions(+)

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 c209b27bc99..783d85f471b 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
@@ -614,6 +614,59 @@ def test_vertex_ai_usage_metadata_with_image_tokens_auto_calculated_text():
     assert result.completion_tokens_details.reasoning_tokens == 158
 
 
+def test_vertex_ai_usage_metadata_with_image_tokens_in_prompt():
+    """Test promptTokensDetails with IMAGE modality for multimodal inputs
+    
+    This test verifies the fix for issue #18182 where image_tokens were missing
+    from prompt_tokens_details when calling Gemini models with image inputs.
+    
+    Example scenario: User sends a text prompt + image, and Gemini generates an image response.
+    The promptTokensDetails should include both TEXT and IMAGE token counts.
+    
+    In this test case, candidatesTokenCount is INCLUSIVE of thoughtsTokenCount because:
+    promptTokenCount (533) + candidatesTokenCount (1337) = totalTokenCount (1870)
+    """
+    v = VertexGeminiConfig()
+    usage_metadata = {
+        "promptTokenCount": 533,
+        "candidatesTokenCount": 1337,  # INCLUSIVE of thoughtsTokenCount
+        "totalTokenCount": 1870,
+        "promptTokensDetails": [
+            {"modality": "IMAGE", "tokenCount": 527},
+            {"modality": "TEXT", "tokenCount": 6}
+        ],
+        "candidatesTokensDetails": [
+            {"modality": "IMAGE", "tokenCount": 1120}
+        ],
+        "thoughtsTokenCount": 217
+    }
+    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 == 533
+    # candidatesTokenCount is INCLUSIVE, so completion_tokens = candidatesTokenCount
+    assert result.completion_tokens == 1337
+    assert result.total_tokens == 1870
+    
+    # Verify prompt_tokens_details includes both text and image tokens
+    assert result.prompt_tokens_details.text_tokens == 6
+    assert result.prompt_tokens_details.image_tokens == 527
+    
+    # Verify completion_tokens_details
+    assert result.completion_tokens_details.image_tokens == 1120
+    assert result.completion_tokens_details.reasoning_tokens == 217
+    
+    # Verify the math: prompt_tokens = text + image
+    # 533 = 6 (text) + 527 (image)
+    assert (
+        result.prompt_tokens_details.text_tokens
+        + result.prompt_tokens_details.image_tokens
+        == result.prompt_tokens
+    )
+
+
 def test_vertex_ai_map_thinking_param_with_budget_tokens_0():
     """
     If budget_tokens is 0, do not set includeThoughts to True