litellm/tests/llm_translation/test_gemini_image_usage.py

"""
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",
        "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