Propagate token usage when generating images with Gemini (#17987)

Author: komarovd95

litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py

@@ -13,7 +13,7 @@
     AllMessageValues,
     OpenAIImageGenerationOptionalParams,
 )
-from litellm.types.utils import ImageObject, ImageResponse
+from litellm.types.utils import ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
@@ -234,6 +234,27 @@ def transform_image_generation_request(
 
         return request_body
 
+    def _transform_image_usage(self, usage: dict) -> ImageUsage:
+        input_tokens_details = ImageUsageInputTokensDetails(
+            image_tokens=0,
+            text_tokens=0,
+        )
+        tokens_details = usage.get("promptTokensDetails", [])
+        for details in tokens_details:
+            if isinstance(details, dict) and (modality := details.get("modality")):
+                token_count = details.get("tokenCount", 0)
+                if modality == "TEXT":
+                    input_tokens_details.text_tokens += token_count
+                elif modality == "IMAGE":
+                    input_tokens_details.image_tokens += token_count
+
+        return ImageUsage(
+            input_tokens=usage.get("promptTokenCount", 0),
+            input_tokens_details=input_tokens_details,
+            output_tokens=usage.get("candidatesTokenCount", 0),
+            total_tokens=usage.get("totalTokenCount", 0),
+        )
+
     def transform_image_generation_response(
         self,
         model: str,
@@ -276,6 +297,9 @@ def transform_image_generation_response(
                             b64_json=inline_data["data"],
                             url=None,
                         ))
+
+        if usage_metadata := response_data.get("usageMetadata", None):
+            model_response.usage = self._transform_image_usage(usage_metadata)
 
         return model_response
 

tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py

@@ -141,7 +141,22 @@ def test_transform_image_generation_response(self):
                         ]
                     }
                 }
-            ]
+            ],
+            "usageMetadata": {
+                "promptTokenCount": 93,
+                "promptTokensDetails": [
+                    {
+                        "modality": "TEXT",
+                        "tokenCount": 54,
+                    },
+                    {
+                        "modality": "IMAGE",
+                        "tokenCount": 39,
+                    }
+                ],
+                "candidatesTokenCount": 17,
+                "totalTokenCount": 110,
+            }
         }
         mock_response.headers = {}
 
@@ -162,6 +177,12 @@ def test_transform_image_generation_response(self):
         assert len(result.data) == 1
         assert result.data[0].b64_json == "base64_encoded_image_data"
         assert result.data[0].url is None
+        assert result.usage.input_tokens == 93
+        assert result.usage.input_tokens_details.text_tokens == 54
+        assert result.usage.input_tokens_details.image_tokens == 39
+        assert result.usage.output_tokens == 17
+        assert result.usage.total_tokens == 110
+
 
     def test_transform_image_generation_response_multiple_images(self):
         """Test response transformation with multiple images"""