From a8371d2cb16a1c60e6aa25caa95a435a00abef14 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 28 Jul 2025 21:25:40 -0700 Subject: [PATCH] [Feat] Add Google AI Studio Imagen4 model family (#13065) * add gemini * add init files * add get_gemini_image_generation_config * refactor transform * TestGoogleImageGen * fix transform * fix transform * add gemini_image_cost_calculator * add cost tracking for gemini/imagen models * docs image gen * docs image gen * test_get_model_info_gemini --- docs/my-website/docs/image_generation.md | 13 ++ .../providers/google_ai_studio/image_gen.md | 214 ++++++++++++++++++ litellm/cost_calculator.py | 8 + litellm/images/main.py | 2 + .../llms/gemini/image_generation/__init__.py | 13 ++ .../image_generation/cost_calculator.py | 30 +++ .../gemini/image_generation/transformation.py | 200 ++++++++++++++++ litellm/llms/recraft/cost_calculator.py | 13 +- ...odel_prices_and_context_window_backup.json | 36 +++ litellm/types/llms/gemini.py | 71 ++++++ litellm/utils.py | 6 + model_prices_and_context_window.json | 36 +++ .../image_gen_tests/test_image_generation.py | 3 + tests/test_litellm/test_utils.py | 4 +- 14 files changed, 644 insertions(+), 5 deletions(-) create mode 100644 docs/my-website/docs/providers/google_ai_studio/image_gen.md create mode 100644 litellm/llms/gemini/image_generation/__init__.py create mode 100644 litellm/llms/gemini/image_generation/cost_calculator.py create mode 100644 litellm/llms/gemini/image_generation/transformation.py diff --git a/docs/my-website/docs/image_generation.md b/docs/my-website/docs/image_generation.md index 792a21fc1a6..60a6356f012 100644 --- a/docs/my-website/docs/image_generation.md +++ b/docs/my-website/docs/image_generation.md @@ -278,3 +278,16 @@ response = litellm.image_generation( ) print(f"response: {response}") ``` + +## Supported Providers + +| Provider | Documentation Link | +|----------|-------------------| +| OpenAI | [OpenAI Image Generation →](./providers/openai) | +| Azure OpenAI | [Azure OpenAI Image Generation →](./providers/azure/azure) | +| Google AI Studio | [Google AI Studio Image Generation →](./providers/google_ai_studio/image_gen) | +| Vertex AI | [Vertex AI Image Generation →](./providers/vertex_image) | +| AWS Bedrock | [Bedrock Image Generation →](./providers/bedrock) | +| Recraft | [Recraft Image Generation →](./providers/recraft#image-generation) | +| Xinference | [Xinference Image Generation →](./providers/xinference#image-generation) | +| Nscale | [Nscale Image Generation →](./providers/nscale#image-generation) | \ No newline at end of file diff --git a/docs/my-website/docs/providers/google_ai_studio/image_gen.md b/docs/my-website/docs/providers/google_ai_studio/image_gen.md new file mode 100644 index 00000000000..f4e96d5225a --- /dev/null +++ b/docs/my-website/docs/providers/google_ai_studio/image_gen.md @@ -0,0 +1,214 @@ +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + +# Google AI Studio Image Generation + +Google AI Studio provides powerful image generation capabilities using Google's Imagen models to create high-quality images from text descriptions. + +## Overview + +| Property | Details | +|----------|---------| +| Description | Google AI Studio Image Generation uses Google's Imagen models to generate high-quality images from text descriptions. | +| Provider Route on LiteLLM | `gemini/` | +| Provider Doc | [Google AI Studio Image Generation ↗](https://ai.google.dev/gemini-api/docs/imagen) | +| Supported Operations | [`/images/generations`](#image-generation) | + +## Setup + +### API Key + +```python showLineNumbers +# Set your Google AI Studio API key +import os +os.environ["GEMINI_API_KEY"] = "your-api-key-here" +``` + +Get your API key from [Google AI Studio](https://aistudio.google.com/app/apikey). + +## Image Generation + +### Usage - LiteLLM Python SDK + + + + +```python showLineNumbers title="Basic Image Generation" +import litellm +import os + +# Set your API key +os.environ["GEMINI_API_KEY"] = "your-api-key-here" + +# Generate a single image +response = litellm.image_generation( + model="gemini/imagen-4.0-generate-preview-06-06", + prompt="A cute baby sea otter swimming in crystal clear water" +) + +print(response.data[0].url) +``` + + + + + +```python showLineNumbers title="Async Image Generation" +import litellm +import asyncio +import os + +async def generate_image(): + # Set your API key + os.environ["GEMINI_API_KEY"] = "your-api-key-here" + + # Generate image asynchronously + response = await litellm.aimage_generation( + model="gemini/imagen-4.0-generate-preview-06-06", + prompt="A beautiful sunset over mountains with vibrant colors", + n=1, + ) + + print(response.data[0].url) + return response + +# Run the async function +asyncio.run(generate_image()) +``` + + + + + +```python showLineNumbers title="Advanced Image Generation with Parameters" +import litellm +import os + +# Set your API key +os.environ["GEMINI_API_KEY"] = "your-api-key-here" + +# Generate image with additional parameters +response = litellm.image_generation( + model="gemini/imagen-4.0-generate-preview-06-06", + prompt="A futuristic cityscape at night with neon lights", + n=1, + size="1024x1024", + quality="standard", + response_format="url" +) + +for image in response.data: + print(f"Generated image URL: {image.url}") +``` + + + + +### Usage - LiteLLM Proxy Server + +#### 1. Configure your config.yaml + +```yaml showLineNumbers title="Google AI Studio Image Generation Configuration" +model_list: + - model_name: google-imagen + litellm_params: + model: gemini/imagen-4.0-generate-preview-06-06 + api_key: os.environ/GEMINI_API_KEY + model_info: + mode: image_generation + +general_settings: + master_key: sk-1234 +``` + +#### 2. Start LiteLLM Proxy Server + +```bash showLineNumbers title="Start LiteLLM Proxy Server" +litellm --config /path/to/config.yaml + +# RUNNING on http://0.0.0.0:4000 +``` + +#### 3. Make requests with OpenAI Python SDK + + + + +```python showLineNumbers title="Google AI Studio Image Generation via Proxy - OpenAI SDK" +from openai import OpenAI + +# Initialize client with your proxy URL +client = OpenAI( + base_url="http://localhost:4000", # Your proxy URL + api_key="sk-1234" # Your proxy API key +) + +# Generate image +response = client.images.generate( + model="google-imagen", + prompt="A majestic eagle soaring over snow-capped mountains", + n=1, + size="1024x1024" +) + +print(response.data[0].url) +``` + + + + + +```python showLineNumbers title="Google AI Studio Image Generation via Proxy - LiteLLM SDK" +import litellm + +# Configure LiteLLM to use your proxy +response = litellm.image_generation( + model="litellm_proxy/google-imagen", + prompt="A serene Japanese garden with cherry blossoms", + api_base="http://localhost:4000", + api_key="sk-1234" +) + +print(response.data[0].url) +``` + + + + + +```bash showLineNumbers title="Google AI Studio Image Generation via Proxy - cURL" +curl --location 'http://localhost:4000/v1/images/generations' \ +--header 'Content-Type: application/json' \ +--header 'Authorization: Bearer sk-1234' \ +--data '{ + "model": "google-imagen", + "prompt": "A cozy coffee shop interior with warm lighting", + "n": 1, + "size": "1024x1024" +}' +``` + + + + +## Supported Parameters + +Google AI Studio Image Generation supports the following OpenAI-compatible parameters: + +| Parameter | Type | Description | Default | Example | +|-----------|------|-------------|---------|---------| +| `prompt` | string | Text description of the image to generate | Required | `"A sunset over the ocean"` | +| `model` | string | The model to use for generation | Required | `"gemini/imagen-4.0-generate-preview-06-06"` | +| `n` | integer | Number of images to generate (1-4) | `1` | `2` | +| `size` | string | Image dimensions | `"1024x1024"` | `"512x512"`, `"1024x1024"` | + +1. Create an account at [Google AI Studio](https://aistudio.google.com/) +2. Generate an API key from [API Keys section](https://aistudio.google.com/app/apikey) +3. Set your `GEMINI_API_KEY` environment variable +4. Start generating images using LiteLLM + +## Additional Resources + +- [Google AI Studio Documentation](https://ai.google.dev/gemini-api/docs) +- [Imagen Model Overview](https://ai.google.dev/gemini-api/docs/imagen) +- [LiteLLM Image Generation Guide](../../completion/image_generation) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index d83d15cb165..c8892cd26a5 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -793,6 +793,14 @@ def completion_cost( # noqa: PLR0915 model=model, image_response=completion_response, ) + elif custom_llm_provider == litellm.LlmProviders.GEMINI.value: + from litellm.llms.gemini.image_generation.cost_calculator import ( + cost_calculator as gemini_image_cost_calculator, + ) + return gemini_image_cost_calculator( + model=model, + image_response=completion_response, + ) else: return default_image_cost_calculator( model=model, diff --git a/litellm/images/main.py b/litellm/images/main.py index 3a675a8168e..b808388d83e 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -411,6 +411,8 @@ def image_generation( # noqa: PLR0915 ######################################################### elif custom_llm_provider in ( litellm.LlmProviders.RECRAFT, + litellm.LlmProviders.GEMINI, + ): if image_generation_config is None: raise ValueError(f"image generation config is not supported for {custom_llm_provider}") diff --git a/litellm/llms/gemini/image_generation/__init__.py b/litellm/llms/gemini/image_generation/__init__.py new file mode 100644 index 00000000000..f99ca1383a9 --- /dev/null +++ b/litellm/llms/gemini/image_generation/__init__.py @@ -0,0 +1,13 @@ +from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, +) + +from .transformation import GoogleImageGenConfig + +__all__ = [ + "GoogleImageGenConfig", +] + + +def get_gemini_image_generation_config(model: str) -> BaseImageGenerationConfig: + return GoogleImageGenConfig() diff --git a/litellm/llms/gemini/image_generation/cost_calculator.py b/litellm/llms/gemini/image_generation/cost_calculator.py new file mode 100644 index 00000000000..0a9ca2e5276 --- /dev/null +++ b/litellm/llms/gemini/image_generation/cost_calculator.py @@ -0,0 +1,30 @@ +""" +Google AI Image Generation Cost Calculator +""" + +from typing import Any + +import litellm +from litellm.types.utils import ImageResponse + + +def cost_calculator( + model: str, + image_response: Any, +) -> float: + """ + Vertex AI Image Generation Cost Calculator + """ + _model_info = litellm.get_model_info( + model=model, + custom_llm_provider="gemini", + ) + + output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 + num_images: int = 0 + if isinstance(image_response, ImageResponse): + if image_response.data: + num_images = len(image_response.data) + return output_cost_per_image * num_images + else: + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py new file mode 100644 index 00000000000..72ba5bcf1ee --- /dev/null +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -0,0 +1,200 @@ +from typing import TYPE_CHECKING, Any, List, Optional + +import httpx + +from litellm.llms.base_llm.image_generation.transformation import ( + BaseImageGenerationConfig, +) +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.gemini import GeminiImageGenerationRequest +from litellm.types.llms.openai import ( + AllMessageValues, + OpenAIImageGenerationOptionalParams, +) +from litellm.types.utils import ImageObject, ImageResponse + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj + + LiteLLMLoggingObj = _LiteLLMLoggingObj +else: + LiteLLMLoggingObj = Any + + +class GoogleImageGenConfig(BaseImageGenerationConfig): + DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta" + + def get_supported_openai_params( + self, model: str + ) -> List[OpenAIImageGenerationOptionalParams]: + """ + Google AI Imagen API supported parameters + https://ai.google.dev/gemini-api/docs/imagen + """ + return [ + "n", + "size" + ] + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + supported_params = self.get_supported_openai_params(model) + mapped_params = {} + + for k, v in non_default_params.items(): + if k not in optional_params.keys(): + if k in supported_params: + # Map OpenAI parameters to Google format + if k == "n": + mapped_params["sampleCount"] = v + elif k == "size": + # Map OpenAI size format to Google aspectRatio + mapped_params["aspectRatio"] = self._map_size_to_aspect_ratio(v) + else: + mapped_params[k] = v + return mapped_params + + + def _map_size_to_aspect_ratio(self, size: str) -> str: + """ + https://ai.google.dev/gemini-api/docs/image-generation + + """ + aspect_ratio_map = { + "1024x1024": "1:1", + "1792x1024": "16:9", + "1024x1792": "9:16", + "1280x896": "4:3", + "896x1280": "3:4" + } + return aspect_ratio_map.get(size, "1:1") + + def get_complete_url( + self, + api_base: Optional[str], + api_key: Optional[str], + model: str, + optional_params: dict, + litellm_params: dict, + stream: Optional[bool] = None, + ) -> str: + """ + Get the complete url for the request + + Google AI API format: https://generativelanguage.googleapis.com/v1beta/models/{model}:predict + """ + complete_url: str = ( + api_base + or get_secret_str("GEMINI_API_BASE") + or self.DEFAULT_BASE_URL + ) + + complete_url = complete_url.rstrip("/") + complete_url = f"{complete_url}/models/{model}:predict" + return complete_url + + def validate_environment( + self, + headers: dict, + model: str, + messages: List[AllMessageValues], + optional_params: dict, + litellm_params: dict, + api_key: Optional[str] = None, + api_base: Optional[str] = None, + ) -> dict: + final_api_key: Optional[str] = ( + api_key or + get_secret_str("GEMINI_API_KEY") + ) + if not final_api_key: + raise ValueError("GEMINI_API_KEY is not set") + + headers["x-goog-api-key"] = final_api_key + headers["Content-Type"] = "application/json" + return headers + + def transform_image_generation_request( + self, + model: str, + prompt: str, + optional_params: dict, + litellm_params: dict, + headers: dict, + ) -> dict: + """ + Transform the image generation request to Google AI Imagen format + + Google AI API format: + { + "instances": [ + { + "prompt": "Robot holding a red skateboard" + } + ], + "parameters": { + "sampleCount": 4, + "aspectRatio": "1:1", + "personGeneration": "allow_adult" + } + } + """ + from litellm.types.llms.gemini import ( + GeminiImageGenerationInstance, + GeminiImageGenerationParameters, + ) + request_body: GeminiImageGenerationRequest = GeminiImageGenerationRequest( + instances=[ + GeminiImageGenerationInstance( + prompt=prompt + ) + ], + parameters=GeminiImageGenerationParameters(**optional_params) + ) + return request_body.model_dump(exclude_none=True) + + def transform_image_generation_response( + self, + model: str, + raw_response: httpx.Response, + model_response: ImageResponse, + logging_obj: LiteLLMLoggingObj, + request_data: dict, + optional_params: dict, + litellm_params: dict, + encoding: Any, + api_key: Optional[str] = None, + json_mode: Optional[bool] = None, + ) -> ImageResponse: + """ + Transform Google AI Imagen response to litellm ImageResponse format + """ + try: + response_data = raw_response.json() + except Exception as e: + raise self.get_error_class( + error_message=f"Error transforming image generation response: {e}", + status_code=raw_response.status_code, + headers=raw_response.headers, + ) + + if not model_response.data: + model_response.data = [] + + # Google AI returns predictions with generated images + predictions = response_data.get("predictions", []) + for prediction in predictions: + # Google AI returns base64 encoded images in the prediction + generated_images = prediction.get("generatedImages", []) + for image_data in generated_images: + model_response.data.append(ImageObject( + b64_json=image_data.get("bytesBase64Encoded", None), + url=None, # Google AI returns base64, not URLs + )) + + return model_response \ No newline at end of file diff --git a/litellm/llms/recraft/cost_calculator.py b/litellm/llms/recraft/cost_calculator.py index 17642cc693d..5ab47e9395e 100644 --- a/litellm/llms/recraft/cost_calculator.py +++ b/litellm/llms/recraft/cost_calculator.py @@ -1,10 +1,12 @@ +from typing import Any + import litellm from litellm.types.utils import ImageResponse def cost_calculator( model: str, - image_response: ImageResponse, + image_response: Any, ) -> float: """ Recraft image generation cost calculator @@ -15,6 +17,9 @@ def cost_calculator( ) output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0 num_images: int = 0 - if image_response.data: - num_images = len(image_response.data) - return output_cost_per_image * num_images + if isinstance(image_response, ImageResponse): + if image_response.data: + num_images = len(image_response.data) + return output_cost_per_image * num_images + else: + raise ValueError(f"image_response must be of type ImageResponse got type={type(image_response)}") diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index f027d8ec3d1..40a07a74189 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -10094,6 +10094,42 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", "supports_tool_choice": true }, + "gemini/imagen-4.0-generate-preview-06-06": { + "output_cost_per_image": 0.04, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "gemini/imagen-4.0-ultra-generate-preview-06-06": { + "output_cost_per_image": 0.06, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "gemini/imagen-4.0-fast-generate-preview-06-06": { + "output_cost_per_image": 0.02, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "gemini/imagen-3.0-generate-002": { + "output_cost_per_image": 0.04, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "gemini/imagen-3.0-generate-001": { + "output_cost_per_image": 0.04, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "gemini/imagen-3.0-fast-generate-001": { + "output_cost_per_image": 0.02, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "command-a-03-2025": { "max_tokens": 8000, "max_input_tokens": 256000, diff --git a/litellm/types/llms/gemini.py b/litellm/types/llms/gemini.py index e39a2a8e820..cfc13cc44a8 100644 --- a/litellm/types/llms/gemini.py +++ b/litellm/types/llms/gemini.py @@ -150,3 +150,74 @@ class BidiGenerateContentSetup(TypedDict, total=False): outputAudioTranscription: dict """The output audio transcription to be used for the realtime session.""" + + +# Image Generation Types +from pydantic import BaseModel + + +class GeminiImageGenerationInstance(TypedDict): + """Instance data for Gemini image generation request""" + prompt: str + + +class GeminiImageGenerationParameters(BaseModel): + """Parameters for Gemini image generation request""" + sampleCount: Optional[int] = None + """Number of images to generate (maps to OpenAI 'n' parameter)""" + + aspectRatio: Optional[str] = None + """Aspect ratio for generated images (e.g., '1:1', '16:9', '9:16', '4:3', '3:4')""" + + personGeneration: Optional[str] = None + """Controls person generation in images""" + + # Additional parameters that might be passed through + background: Optional[str] = None + """Background specification""" + + input_fidelity: Optional[str] = None + """Input fidelity specification""" + + moderation: Optional[str] = None + """Moderation settings""" + + output_compression: Optional[str] = None + """Output compression settings""" + + output_format: Optional[str] = None + """Output format specification""" + + quality: Optional[str] = None + """Quality settings""" + + response_format: Optional[str] = None + """Response format specification""" + + style: Optional[str] = None + """Style specification""" + + user: Optional[str] = None + """User specification""" + + +class GeminiImageGenerationRequest(BaseModel): + """Complete request body for Gemini image generation""" + instances: List[GeminiImageGenerationInstance] + parameters: GeminiImageGenerationParameters + + +class GeminiGeneratedImage(TypedDict): + """Individual generated image data from Gemini response""" + bytesBase64Encoded: str + """Base64 encoded image data""" + + +class GeminiImageGenerationPrediction(TypedDict): + """Prediction object containing generated images""" + generatedImages: List[GeminiGeneratedImage] + + +class GeminiImageGenerationResponse(TypedDict): + """Complete response body from Gemini image generation API""" + predictions: List[GeminiImageGenerationPrediction] diff --git a/litellm/utils.py b/litellm/utils.py index 9a618a49a1e..420f85ac7f6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7187,6 +7187,12 @@ class ProviderConfigManager: ) return get_recraft_image_generation_config(model) + elif LlmProviders.GEMINI == provider: + from litellm.llms.gemini.image_generation import ( + get_gemini_image_generation_config, + ) + + return get_gemini_image_generation_config(model) return None @staticmethod diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index f027d8ec3d1..40a07a74189 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -10094,6 +10094,42 @@ "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models", "supports_tool_choice": true }, + "gemini/imagen-4.0-generate-preview-06-06": { + "output_cost_per_image": 0.04, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "gemini/imagen-4.0-ultra-generate-preview-06-06": { + "output_cost_per_image": 0.06, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "gemini/imagen-4.0-fast-generate-preview-06-06": { + "output_cost_per_image": 0.02, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "gemini/imagen-3.0-generate-002": { + "output_cost_per_image": 0.04, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "gemini/imagen-3.0-generate-001": { + "output_cost_per_image": 0.04, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, + "gemini/imagen-3.0-fast-generate-001": { + "output_cost_per_image": 0.02, + "litellm_provider": "gemini", + "mode": "image_generation", + "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" + }, "command-a-03-2025": { "max_tokens": 8000, "max_input_tokens": 256000, diff --git a/tests/image_gen_tests/test_image_generation.py b/tests/image_gen_tests/test_image_generation.py index 79f7f42e554..c34fd0b5e83 100644 --- a/tests/image_gen_tests/test_image_generation.py +++ b/tests/image_gen_tests/test_image_generation.py @@ -169,6 +169,9 @@ class TestRecraftImageGeneration(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: return {"model": "recraft/recraftv3"} +class TestGoogleImageGen(BaseImageGenTest): + def get_base_image_generation_call_args(self) -> dict: + return {"model": "gemini/imagen-4.0-generate-preview-06-06"} class TestAzureOpenAIDalle3(BaseImageGenTest): def get_base_image_generation_call_args(self) -> dict: diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 6b136f03c87..9ae18f8f0b4 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -11,6 +11,7 @@ sys.path.insert( ) # Adds the parent directory to the system path import litellm +from litellm.proxy.utils import is_valid_api_key from litellm.types.utils import ( Delta, LlmProviders, @@ -23,7 +24,6 @@ from litellm.utils import ( get_llm_provider, get_optional_params_image_gen, ) -from litellm.proxy.utils import is_valid_api_key # Adds the parent directory to the system path @@ -558,6 +558,7 @@ def test_get_model_info_gemini(): model.startswith("gemini/") and not "gemma" in model and not "learnlm" in model + and not "imagen" in model ): assert info.get("tpm") is not None, f"{model} does not have tpm" assert info.get("rpm") is not None, f"{model} does not have rpm" @@ -2158,6 +2159,7 @@ def test_image_response_utils(): def test_is_valid_api_key(): import hashlib + # Valid sk- keys assert is_valid_api_key("sk-abc123") assert is_valid_api_key("sk-ABC_123-xyz")