diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 2d8d82e6ad..63b835df9d 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/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", []) diff --git a/tests/llm_translation/test_gemini_image_usage.py b/tests/llm_translation/test_gemini_image_usage.py new file mode 100644 index 0000000000..4c72d544c5 --- /dev/null +++ b/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 +