Add usage object in image generation of gemini

litellm/llms/gemini/image_generation/transformation.py

@@ -11,7 +11,12 @@ from litellm.types.llms.openai import (
     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
@@ -73,6 +78,33 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
             "896x1280": "3:4",
         }
         return aspect_ratio_map.get(size, "1:1")
+    
+    def _transform_image_usage(self, usage_metadata: dict) -> ImageUsage:
+        """
+        Transform Gemini usageMetadata to ImageUsage format
+        """
+        input_tokens_details = ImageUsageInputTokensDetails(
+            image_tokens=0,
+            text_tokens=0,
+        )
+        
+        # Extract detailed token counts from promptTokensDetails
+        tokens_details = usage_metadata.get("promptTokensDetails", [])
+        for details in tokens_details:
+            if isinstance(details, dict):
+                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_metadata.get("promptTokenCount", 0),
+            input_tokens_details=input_tokens_details,
+            output_tokens=usage_metadata.get("candidatesTokenCount", 0),
+            total_tokens=usage_metadata.get("totalTokenCount", 0),
+        )
 
     def get_complete_url(
         self,
@@ -227,6 +259,10 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
                                 b64_json=inline_data["data"],
                                 url=None,
                             ))
+            
+            # Extract usage metadata for Gemini models
+            if "usageMetadata" in response_data:
+                model_response.usage = self._transform_image_usage(response_data["usageMetadata"])
         else:
             # Original Imagen format - predictions with generated images
             predictions = response_data.get("predictions", [])

tests/llm_translation/test_gemini_image_usage.py

@@ -0,0 +1,214 @@
+"""
+Test for Gemini image generation usage metadata extraction.
+
+This test verifies the fix for issue #18323 where image_generation() 
+was returning usage=0 while completion() returned proper token usage.
+"""
+import pytest
+from unittest.mock import patch, MagicMock
+import litellm
+from litellm.types.utils import ImageResponse, ImageObject, ImageUsage
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    [
+        "gemini/gemini-2.5-flash-image-preview",
+        "gemini/gemini-2.0-flash-preview-image-generation",
+        "gemini/gemini-3-pro-image-preview",
+    ],
+)
+def test_gemini_image_generation_usage_metadata(model_name: str):
+    """
+    Test that image_generation() properly extracts and returns usage metadata
+    from Gemini API responses.
+    
+    This test verifies the fix for issue #18323.
+    """
+    
+    # Mock response data that includes usageMetadata (like real Gemini API)
+    mock_response_data = {
+        "candidates": [
+            {
+                "content": {
+                    "parts": [
+                        {
+                            "inlineData": {
+                                "mimeType": "image/png",
+                                "data": "test_base64_image_data"
+                            }
+                        }
+                    ]
+                }
+            }
+        ],
+        "usageMetadata": {
+            "promptTokenCount": 35,
+            "candidatesTokenCount": 1716,
+            "totalTokenCount": 1751,
+            "promptTokensDetails": [
+                {
+                    "modality": "TEXT",
+                    "tokenCount": 35
+                }
+            ],
+            "candidatesTokensDetails": [
+                {
+                    "modality": "TEXT",
+                    "tokenCount": 213
+                },
+                {
+                    "modality": "IMAGE",
+                    "tokenCount": 1120
+                }
+            ]
+        }
+    }
+    
+    with patch(
+        "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post"
+    ) as mock_post:
+        # Mock successful HTTP response
+        mock_http_response = MagicMock()
+        mock_http_response.json.return_value = mock_response_data
+        mock_http_response.status_code = 200
+        mock_http_response.headers = {}
+        mock_post.return_value = mock_http_response
+        
+        # Call image_generation
+        response = litellm.image_generation(
+            model=model_name,
+            prompt="A cute baby sea otter eating a cute baby spinach with cute starry cereals dressing",
+            api_key="test_api_key",
+        )
+        
+        # Validate response structure
+        assert response is not None
+        assert hasattr(response, "data")
+        assert response.data is not None
+        assert len(response.data) > 0
+        
+        # IMPORTANT: Validate usage metadata is properly extracted
+        assert response.usage is not None, "Usage should not be None"
+        
+        # Note: The usage object might be converted to Usage type by Pydantic/OpenAI SDK
+        # but it should still have the ImageUsage fields (input_tokens, output_tokens, etc.)
+        
+        # Validate token counts match the mock response
+        assert hasattr(response.usage, 'input_tokens'), "Usage should have input_tokens attribute"
+        assert hasattr(response.usage, 'output_tokens'), "Usage should have output_tokens attribute"
+        assert hasattr(response.usage, 'total_tokens'), "Usage should have total_tokens attribute"
+        
+        assert response.usage.input_tokens == 35, f"Expected input_tokens=35, got {response.usage.input_tokens}"
+        assert response.usage.output_tokens == 1716, f"Expected output_tokens=1716, got {response.usage.output_tokens}"
+        assert response.usage.total_tokens == 1751, f"Expected total_tokens=1751, got {response.usage.total_tokens}"
+        
+        # Validate input tokens details
+        assert hasattr(response.usage, 'input_tokens_details'), "Usage should have input_tokens_details attribute"
+        assert response.usage.input_tokens_details is not None, "Input tokens details should not be None"
+        
+        # input_tokens_details might be a dict or an object
+        if isinstance(response.usage.input_tokens_details, dict):
+            assert response.usage.input_tokens_details['text_tokens'] == 35, f"Expected text_tokens=35, got {response.usage.input_tokens_details['text_tokens']}"
+            assert response.usage.input_tokens_details['image_tokens'] == 0, f"Expected image_tokens=0, got {response.usage.input_tokens_details['image_tokens']}"
+        else:
+            assert response.usage.input_tokens_details.text_tokens == 35, f"Expected text_tokens=35, got {response.usage.input_tokens_details.text_tokens}"
+            assert response.usage.input_tokens_details.image_tokens == 0, f"Expected image_tokens=0, got {response.usage.input_tokens_details.image_tokens}"
+        
+        # Verify the usage is not all zeros (the bug we're fixing)
+        assert response.usage.total_tokens > 0, "Total tokens should be greater than 0"
+        assert response.usage.input_tokens > 0, "Input tokens should be greater than 0"
+        assert response.usage.output_tokens > 0, "Output tokens should be greater than 0"
+
+
+def test_gemini_image_generation_without_usage_metadata():
+    """
+    Test that image_generation() handles responses without usageMetadata gracefully.
+    """
+    
+    # Mock response data without usageMetadata
+    mock_response_data = {
+        "candidates": [
+            {
+                "content": {
+                    "parts": [
+                        {
+                            "inlineData": {
+                                "mimeType": "image/png",
+                                "data": "test_base64_image_data"
+                            }
+                        }
+                    ]
+                }
+            }
+        ]
+    }
+    
+    with patch(
+        "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post"
+    ) as mock_post:
+        # Mock successful HTTP response
+        mock_http_response = MagicMock()
+        mock_http_response.json.return_value = mock_response_data
+        mock_http_response.status_code = 200
+        mock_http_response.headers = {}
+        mock_post.return_value = mock_http_response
+        
+        # Call image_generation
+        response = litellm.image_generation(
+            model="gemini/gemini-3-pro-image-preview",
+            prompt="Test prompt",
+            api_key="test_api_key",
+        )
+        
+        # Validate response structure
+        assert response is not None
+        assert hasattr(response, "data")
+        assert response.data is not None
+        assert len(response.data) > 0
+        
+        # Usage should be None if not present in response
+        # (or have default values depending on implementation)
+        # This ensures we don't crash when usageMetadata is missing
+
+
+def test_gemini_imagen_models_no_usage_extraction():
+    """
+    Test that non-Gemini Imagen models don't attempt to extract usage metadata
+    from the different response format.
+    """
+    
+    # Mock response data for Imagen models (different format)
+    mock_response_data = {
+        "predictions": [
+            {
+                "bytesBase64Encoded": "test_base64_image_data"
+            }
+        ]
+    }
+    
+    with patch(
+        "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post"
+    ) as mock_post:
+        # Mock successful HTTP response
+        mock_http_response = MagicMock()
+        mock_http_response.json.return_value = mock_response_data
+        mock_http_response.status_code = 200
+        mock_http_response.headers = {}
+        mock_post.return_value = mock_http_response
+        
+        # Call image_generation with an Imagen model
+        response = litellm.image_generation(
+            model="gemini/imagen-3.0-generate-001",
+            prompt="Test prompt",
+            api_key="test_api_key",
+        )
+        
+        # Validate response structure
+        assert response is not None
+        assert hasattr(response, "data")
+        assert response.data is not None
+        
+        # For Imagen models, we don't extract usage from the predictions format
+        # This test just ensures we don't crash
+