Remove hardcoded model name and fix breaking change

- Reverted GEMINI_2_5_FLASH_IMAGE_PREVIEW_MODEL constant usage
- Made endpoint selection conditional for gemini-2.5-flash-image-preview only
- Preserved existing Imagen models functionality with :predict endpoint
- Fixed potential breaking change that would affect 6 other Gemini image models
This commit is contained in:
Tim Elfrink 2025-09-19 09:29:14 +02:00
parent f5e6246143
commit 3a98fd6096
2 changed files with 58 additions and 4 deletions

View File

@ -86,8 +86,8 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
"""
Get the complete url for the request
Google AI API format: https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent
Note: Gemini image generation models use generateContent, not predict
Gemini 2.5 Flash Image Preview: :generateContent
Other Imagen models: :predict
"""
complete_url: str = (
api_base
@ -96,7 +96,14 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
)
complete_url = complete_url.rstrip("/")
complete_url = f"{complete_url}/models/{model}:generateContent"
# Gemini 2.5 Flash Image Preview uses generateContent endpoint
if "2.5-flash-image-preview" in model:
complete_url = f"{complete_url}/models/{model}:generateContent"
else:
# All other Imagen models use predict endpoint
complete_url = f"{complete_url}/models/{model}:predict"
return complete_url
def validate_environment(

View File

@ -323,7 +323,7 @@ def test_gemini_2_5_flash_image_preview():
call_args = mock_post.call_args
called_url = call_args[0][0] if call_args[0] else call_args.kwargs.get('url', '')
# Verify it uses generateContent endpoint (not predict)
# Verify it uses generateContent endpoint for gemini-2.5-flash-image-preview (not predict)
assert ":generateContent" in called_url
assert "gemini-2.5-flash-image-preview" in called_url
@ -333,6 +333,53 @@ def test_gemini_2_5_flash_image_preview():
assert "parts" in request_data["contents"][0]
def test_gemini_imagen_models_use_predict_endpoint():
"""
Test that Imagen models still use :predict endpoint (not broken by gemini-2.5-flash-image-preview fix)
"""
from unittest.mock import patch, MagicMock
from litellm.types.utils import ImageResponse, ImageObject
with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post:
# Mock successful HTTP response for Imagen
mock_http_response = MagicMock()
mock_http_response.json.return_value = {
"predictions": [
{
"bytesBase64Encoded": "test_base64_image_data"
}
]
}
mock_http_response.status_code = 200
mock_post.return_value = mock_http_response
# Test an Imagen model
response = litellm.image_generation(
model="gemini/imagen-3.0-generate-001",
prompt="Generate a simple test image",
api_key="test_api_key"
)
# Validate response structure
assert response is not None
assert hasattr(response, 'data')
# Validate the correct endpoint was called for Imagen models
mock_post.assert_called_once()
call_args = mock_post.call_args
called_url = call_args[0][0] if call_args[0] else call_args.kwargs.get('url', '')
# Verify Imagen models use predict endpoint (not generateContent)
assert ":predict" in called_url
assert "imagen-3.0-generate-001" in called_url
assert ":generateContent" not in called_url
# Verify request format is Imagen format (not Gemini)
request_data = call_args.kwargs.get('json', {})
assert "instances" in request_data
assert "parameters" in request_data
def test_gemini_thinking():
litellm._turn_on_debug()
from litellm.types.utils import Message, CallTypes