Merge pull request #17070 from BerriAI/litellm_add_vertex_ai_image_support
Add vertex ai image gen support for both gemini and imagen models
This commit is contained in:
Sameer Kankute
GitHub
commit
3249f6dd2d
@ -1,18 +1,65 @@
|
||||
# Vertex AI Image Generation
|
||||
|
||||
Vertex AI Image Generation uses Google's Imagen models to generate high-quality images from text descriptions.
|
||||
Vertex AI supports two types of image generation:
|
||||
|
||||
1. **Gemini Image Generation Models** (Nano Banana 🍌) - Conversational image generation using `generateContent` API
|
||||
2. **Imagen Models** - Traditional image generation using `predict` API
|
||||
|
||||
| Property | Details |
|
||||
|----------|---------|
|
||||
| Description | Vertex AI Image Generation uses Google's Imagen models to generate high-quality images from text descriptions. |
|
||||
| Description | Vertex AI Image Generation supports both Gemini image generation models |
|
||||
| Provider Route on LiteLLM | `vertex_ai/` |
|
||||
| Provider Doc | [Google Cloud Vertex AI Image Generation ↗](https://cloud.google.com/vertex-ai/docs/generative-ai/image/generate-images) |
|
||||
| Gemini Image Generation Docs | [Gemini Image Generation ↗](https://ai.google.dev/gemini-api/docs/image-generation) |
|
||||
|
||||
## Quick Start
|
||||
|
||||
### LiteLLM Python SDK
|
||||
### Gemini Image Generation Models
|
||||
|
||||
```python showLineNumbers title="Basic Image Generation"
|
||||
Gemini image generation models support conversational image creation with features like:
|
||||
- Text-to-Image generation
|
||||
- Image editing (text + image → image)
|
||||
- Multi-turn image refinement
|
||||
- High-fidelity text rendering
|
||||
- Up to 4K resolution (Gemini 3 Pro)
|
||||
|
||||
```python showLineNumbers title="Gemini 2.5 Flash Image"
|
||||
import litellm
|
||||
|
||||
# Generate a single image
|
||||
response = await litellm.aimage_generation(
|
||||
prompt="A nano banana dish in a fancy restaurant with a Gemini theme",
|
||||
model="vertex_ai/gemini-2.5-flash-image",
|
||||
vertex_ai_project="your-project-id",
|
||||
vertex_ai_location="us-central1",
|
||||
n=1,
|
||||
size="1024x1024",
|
||||
)
|
||||
|
||||
print(response.data[0].b64_json) # Gemini returns base64 images
|
||||
```
|
||||
|
||||
```python showLineNumbers title="Gemini 3 Pro Image Preview (4K output)"
|
||||
import litellm
|
||||
|
||||
# Generate high-resolution image
|
||||
response = await litellm.aimage_generation(
|
||||
prompt="Da Vinci style anatomical sketch of a dissected Monarch butterfly",
|
||||
model="vertex_ai/gemini-3-pro-image-preview",
|
||||
vertex_ai_project="your-project-id",
|
||||
vertex_ai_location="us-central1",
|
||||
n=1,
|
||||
size="1024x1024",
|
||||
# Optional: specify image size for Gemini 3 Pro
|
||||
# imageSize="4K", # Options: "1K", "2K", "4K"
|
||||
)
|
||||
|
||||
print(response.data[0].b64_json)
|
||||
```
|
||||
|
||||
### Imagen Models
|
||||
|
||||
```python showLineNumbers title="Imagen Image Generation"
|
||||
import litellm
|
||||
|
||||
# Generate a single image
|
||||
@ -21,9 +68,11 @@ response = await litellm.aimage_generation(
|
||||
model="vertex_ai/imagen-4.0-generate-001",
|
||||
vertex_ai_project="your-project-id",
|
||||
vertex_ai_location="us-central1",
|
||||
n=1,
|
||||
size="1024x1024",
|
||||
)
|
||||
|
||||
print(response.data[0].url)
|
||||
print(response.data[0].b64_json) # Imagen also returns base64 images
|
||||
```
|
||||
|
||||
### LiteLLM Proxy
|
||||
@ -70,6 +119,18 @@ print(response.data[0].url)
|
||||
|
||||
## Supported Models
|
||||
|
||||
### Gemini Image Generation Models
|
||||
|
||||
- `vertex_ai/gemini-2.5-flash-image` - Fast, efficient image generation (1024px resolution)
|
||||
- `vertex_ai/gemini-3-pro-image-preview` - Advanced model with 4K output, Google Search grounding, and thinking mode
|
||||
- `vertex_ai/gemini-2.0-flash-preview-image` - Preview model
|
||||
- `vertex_ai/gemini-2.5-flash-image-preview` - Preview model
|
||||
|
||||
### Imagen Models
|
||||
|
||||
- `vertex_ai/imagegeneration@006` - Legacy Imagen model
|
||||
- `vertex_ai/imagen-4.0-generate-001` - Latest Imagen model
|
||||
- `vertex_ai/imagen-3.0-generate-001` - Imagen 3.0 model
|
||||
|
||||
:::tip
|
||||
|
||||
@ -77,7 +138,5 @@ print(response.data[0].url)
|
||||
|
||||
:::
|
||||
|
||||
LiteLLM supports all Vertex AI Imagen models available through Google Cloud.
|
||||
|
||||
For the complete and up-to-date list of supported models, visit: [https://models.litellm.ai/](https://models.litellm.ai/)
|
||||
|
||||
|
||||
@ -19,6 +19,8 @@ from litellm.llms.custom_llm import CustomLLM
|
||||
|
||||
#################### Initialize provider clients ####################
|
||||
llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler()
|
||||
from openai.types.audio.transcription_create_params import FileTypes # type: ignore
|
||||
|
||||
from litellm.main import (
|
||||
azure_chat_completions,
|
||||
base_llm_aiohttp_handler,
|
||||
@ -26,7 +28,6 @@ from litellm.main import (
|
||||
bedrock_image_generation,
|
||||
openai_chat_completions,
|
||||
openai_image_variations,
|
||||
vertex_image_generation,
|
||||
)
|
||||
|
||||
###########################################
|
||||
@ -36,7 +37,6 @@ from litellm.types.llms.openai import ImageGenerationRequestQuality
|
||||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import (
|
||||
LITELLM_IMAGE_VARIATION_PROVIDERS,
|
||||
FileTypes,
|
||||
LlmProviders,
|
||||
all_litellm_params,
|
||||
)
|
||||
@ -344,6 +344,7 @@ def image_generation( # noqa: PLR0915
|
||||
litellm.LlmProviders.GEMINI,
|
||||
litellm.LlmProviders.FAL_AI,
|
||||
litellm.LlmProviders.RUNWAYML,
|
||||
litellm.LlmProviders.VERTEX_AI,
|
||||
):
|
||||
if image_generation_config is None:
|
||||
raise ValueError(
|
||||
@ -430,46 +431,6 @@ def image_generation( # noqa: PLR0915
|
||||
api_base=api_base,
|
||||
api_key=api_key,
|
||||
)
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
vertex_ai_project = (
|
||||
optional_params.pop("vertex_project", None)
|
||||
or optional_params.pop("vertex_ai_project", None)
|
||||
or litellm.vertex_project
|
||||
or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
vertex_ai_location = (
|
||||
optional_params.pop("vertex_location", None)
|
||||
or optional_params.pop("vertex_ai_location", None)
|
||||
or litellm.vertex_location
|
||||
or get_secret_str("VERTEXAI_LOCATION")
|
||||
)
|
||||
vertex_credentials = (
|
||||
optional_params.pop("vertex_credentials", None)
|
||||
or optional_params.pop("vertex_ai_credentials", None)
|
||||
or get_secret_str("VERTEXAI_CREDENTIALS")
|
||||
)
|
||||
|
||||
api_base = (
|
||||
api_base
|
||||
or litellm.api_base
|
||||
or get_secret_str("VERTEXAI_API_BASE")
|
||||
or get_secret_str("VERTEX_API_BASE")
|
||||
)
|
||||
|
||||
model_response = vertex_image_generation.image_generation(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
timeout=timeout,
|
||||
logging_obj=litellm_logging_obj,
|
||||
optional_params=optional_params,
|
||||
model_response=model_response,
|
||||
vertex_project=vertex_ai_project,
|
||||
vertex_location=vertex_ai_location,
|
||||
vertex_credentials=vertex_credentials,
|
||||
aimg_generation=aimg_generation,
|
||||
api_base=api_base,
|
||||
client=client,
|
||||
)
|
||||
elif (
|
||||
custom_llm_provider in litellm._custom_providers
|
||||
): # Assume custom LLM provider
|
||||
|
||||
43
litellm/llms/vertex_ai/image_generation/__init__.py
Normal file
43
litellm/llms/vertex_ai/image_generation/__init__.py
Normal file
@ -0,0 +1,43 @@
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
VertexAIModelRoute,
|
||||
get_vertex_ai_model_route,
|
||||
)
|
||||
|
||||
from .vertex_gemini_transformation import VertexAIGeminiImageGenerationConfig
|
||||
from .vertex_imagen_transformation import VertexAIImagenImageGenerationConfig
|
||||
|
||||
__all__ = [
|
||||
"VertexAIGeminiImageGenerationConfig",
|
||||
"VertexAIImagenImageGenerationConfig",
|
||||
"get_vertex_ai_image_generation_config",
|
||||
]
|
||||
|
||||
|
||||
def get_vertex_ai_image_generation_config(model: str) -> BaseImageGenerationConfig:
|
||||
"""
|
||||
Get the appropriate image generation config for a Vertex AI model.
|
||||
|
||||
Routes to the correct transformation class based on the model type:
|
||||
- Gemini image generation models use generateContent API (VertexAIGeminiImageGenerationConfig)
|
||||
- Imagen models use predict API (VertexAIImagenImageGenerationConfig)
|
||||
|
||||
Args:
|
||||
model: The model name (e.g., "gemini-2.5-flash-image", "imagegeneration@006")
|
||||
|
||||
Returns:
|
||||
BaseImageGenerationConfig: The appropriate configuration class
|
||||
"""
|
||||
# Determine the model route
|
||||
model_route = get_vertex_ai_model_route(model)
|
||||
|
||||
if model_route == VertexAIModelRoute.GEMINI:
|
||||
# Gemini models use generateContent API
|
||||
return VertexAIGeminiImageGenerationConfig()
|
||||
else:
|
||||
# Default to Imagen for other models (imagegeneration, etc.)
|
||||
# This includes NON_GEMINI models like imagegeneration@006
|
||||
return VertexAIImagenImageGenerationConfig()
|
||||
|
||||
@ -0,0 +1,264 @@
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
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 VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
|
||||
"""
|
||||
Vertex AI Gemini Image Generation Configuration
|
||||
|
||||
Uses generateContent API for Gemini image generation models on Vertex AI
|
||||
Supports models like gemini-2.5-flash-image, gemini-3-pro-image-preview, etc.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
BaseImageGenerationConfig.__init__(self)
|
||||
VertexLLM.__init__(self)
|
||||
|
||||
def get_supported_openai_params(
|
||||
self, model: str
|
||||
) -> List[OpenAIImageGenerationOptionalParams]:
|
||||
"""
|
||||
Gemini image generation supported parameters
|
||||
"""
|
||||
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 Gemini format
|
||||
if k == "n":
|
||||
mapped_params["candidate_count"] = v
|
||||
elif k == "size":
|
||||
# Map OpenAI size format to Gemini 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:
|
||||
"""
|
||||
Map OpenAI size format to Gemini aspect ratio format
|
||||
"""
|
||||
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 _resolve_vertex_project(self) -> Optional[str]:
|
||||
return (
|
||||
getattr(self, "_vertex_project", None)
|
||||
or os.environ.get("VERTEXAI_PROJECT")
|
||||
or getattr(litellm, "vertex_project", None)
|
||||
or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
|
||||
def _resolve_vertex_location(self) -> Optional[str]:
|
||||
return (
|
||||
getattr(self, "_vertex_location", None)
|
||||
or os.environ.get("VERTEXAI_LOCATION")
|
||||
or os.environ.get("VERTEX_LOCATION")
|
||||
or getattr(litellm, "vertex_location", None)
|
||||
or get_secret_str("VERTEXAI_LOCATION")
|
||||
or get_secret_str("VERTEX_LOCATION")
|
||||
)
|
||||
|
||||
def _resolve_vertex_credentials(self) -> Optional[str]:
|
||||
return (
|
||||
getattr(self, "_vertex_credentials", None)
|
||||
or os.environ.get("VERTEXAI_CREDENTIALS")
|
||||
or getattr(litellm, "vertex_credentials", None)
|
||||
or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
|
||||
or get_secret_str("VERTEXAI_CREDENTIALS")
|
||||
)
|
||||
|
||||
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 Vertex AI Gemini generateContent API
|
||||
"""
|
||||
vertex_project = self._resolve_vertex_project()
|
||||
vertex_location = self._resolve_vertex_location()
|
||||
|
||||
if not vertex_project or not vertex_location:
|
||||
raise ValueError("vertex_project and vertex_location are required for Vertex AI")
|
||||
|
||||
# Use the model name as provided, handling vertex_ai prefix
|
||||
model_name = model
|
||||
if model.startswith("vertex_ai/"):
|
||||
model_name = model.replace("vertex_ai/", "")
|
||||
|
||||
if api_base:
|
||||
base_url = api_base.rstrip("/")
|
||||
else:
|
||||
base_url = f"https://{vertex_location}-aiplatform.googleapis.com"
|
||||
|
||||
return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}:generateContent"
|
||||
|
||||
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:
|
||||
headers = headers or {}
|
||||
vertex_project = self._resolve_vertex_project()
|
||||
vertex_credentials = self._resolve_vertex_credentials()
|
||||
access_token, _ = self._ensure_access_token(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
return self.set_headers(access_token, 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 Gemini format
|
||||
|
||||
Uses generateContent API with responseModalities: ["IMAGE"]
|
||||
"""
|
||||
# Prepare messages with the prompt
|
||||
contents = [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [{"text": prompt}]
|
||||
}
|
||||
]
|
||||
|
||||
# Prepare generation config
|
||||
generation_config: Dict[str, Any] = {
|
||||
"responseModalities": ["IMAGE"]
|
||||
}
|
||||
|
||||
# Handle image-specific config parameters
|
||||
image_config: Dict[str, Any] = {}
|
||||
|
||||
# Map aspectRatio
|
||||
if "aspectRatio" in optional_params:
|
||||
image_config["aspectRatio"] = optional_params["aspectRatio"]
|
||||
elif "aspect_ratio" in optional_params:
|
||||
image_config["aspectRatio"] = optional_params["aspect_ratio"]
|
||||
|
||||
# Map imageSize (for Gemini 3 Pro)
|
||||
if "imageSize" in optional_params:
|
||||
image_config["imageSize"] = optional_params["imageSize"]
|
||||
elif "image_size" in optional_params:
|
||||
image_config["imageSize"] = optional_params["image_size"]
|
||||
|
||||
if image_config:
|
||||
generation_config["imageConfig"] = image_config
|
||||
|
||||
# Handle candidate_count (n parameter)
|
||||
if "candidate_count" in optional_params:
|
||||
generation_config["candidateCount"] = optional_params["candidate_count"]
|
||||
elif "n" in optional_params:
|
||||
generation_config["candidateCount"] = optional_params["n"]
|
||||
|
||||
request_body: Dict[str, Any] = {
|
||||
"contents": contents,
|
||||
"generationConfig": generation_config
|
||||
}
|
||||
|
||||
return request_body
|
||||
|
||||
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 Gemini image generation 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 = []
|
||||
|
||||
# Gemini image generation models return in candidates format
|
||||
candidates = response_data.get("candidates", [])
|
||||
for candidate in candidates:
|
||||
content = candidate.get("content", {})
|
||||
parts = content.get("parts", [])
|
||||
for part in parts:
|
||||
# Look for inlineData with image
|
||||
if "inlineData" in part:
|
||||
inline_data = part["inlineData"]
|
||||
if "data" in inline_data:
|
||||
model_response.data.append(ImageObject(
|
||||
b64_json=inline_data["data"],
|
||||
url=None,
|
||||
))
|
||||
|
||||
return model_response
|
||||
|
||||
@ -0,0 +1,230 @@
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Any, List, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.llms.base_llm.image_generation.transformation import (
|
||||
BaseImageGenerationConfig,
|
||||
)
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
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 VertexAIImagenImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
|
||||
"""
|
||||
Vertex AI Imagen Image Generation Configuration
|
||||
|
||||
Uses predict API for Imagen models on Vertex AI
|
||||
Supports models like imagegeneration@006
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
BaseImageGenerationConfig.__init__(self)
|
||||
VertexLLM.__init__(self)
|
||||
|
||||
def get_supported_openai_params(
|
||||
self, model: str
|
||||
) -> List[OpenAIImageGenerationOptionalParams]:
|
||||
"""
|
||||
Imagen API supported parameters
|
||||
"""
|
||||
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 Imagen format
|
||||
if k == "n":
|
||||
mapped_params["sampleCount"] = v
|
||||
elif k == "size":
|
||||
# Map OpenAI size format to Imagen 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:
|
||||
"""
|
||||
Map OpenAI size format to Imagen aspect ratio format
|
||||
"""
|
||||
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 _resolve_vertex_project(self) -> Optional[str]:
|
||||
return (
|
||||
getattr(self, "_vertex_project", None)
|
||||
or os.environ.get("VERTEXAI_PROJECT")
|
||||
or getattr(litellm, "vertex_project", None)
|
||||
or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
|
||||
def _resolve_vertex_location(self) -> Optional[str]:
|
||||
return (
|
||||
getattr(self, "_vertex_location", None)
|
||||
or os.environ.get("VERTEXAI_LOCATION")
|
||||
or os.environ.get("VERTEX_LOCATION")
|
||||
or getattr(litellm, "vertex_location", None)
|
||||
or get_secret_str("VERTEXAI_LOCATION")
|
||||
or get_secret_str("VERTEX_LOCATION")
|
||||
)
|
||||
|
||||
def _resolve_vertex_credentials(self) -> Optional[str]:
|
||||
return (
|
||||
getattr(self, "_vertex_credentials", None)
|
||||
or os.environ.get("VERTEXAI_CREDENTIALS")
|
||||
or getattr(litellm, "vertex_credentials", None)
|
||||
or os.environ.get("GOOGLE_APPLICATION_CREDENTIALS")
|
||||
or get_secret_str("VERTEXAI_CREDENTIALS")
|
||||
)
|
||||
|
||||
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 Vertex AI Imagen predict API
|
||||
"""
|
||||
vertex_project = self._resolve_vertex_project()
|
||||
vertex_location = self._resolve_vertex_location()
|
||||
|
||||
if not vertex_project or not vertex_location:
|
||||
raise ValueError("vertex_project and vertex_location are required for Vertex AI")
|
||||
|
||||
# Use the model name as provided, handling vertex_ai prefix
|
||||
model_name = model
|
||||
if model.startswith("vertex_ai/"):
|
||||
model_name = model.replace("vertex_ai/", "")
|
||||
|
||||
if api_base:
|
||||
base_url = api_base.rstrip("/")
|
||||
else:
|
||||
base_url = f"https://{vertex_location}-aiplatform.googleapis.com"
|
||||
|
||||
return f"{base_url}/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model_name}:predict"
|
||||
|
||||
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:
|
||||
headers = headers or {}
|
||||
vertex_project = self._resolve_vertex_project()
|
||||
vertex_credentials = self._resolve_vertex_credentials()
|
||||
access_token, _ = self._ensure_access_token(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
return self.set_headers(access_token, 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 Imagen format
|
||||
|
||||
Uses predict API with instances and parameters
|
||||
"""
|
||||
# Default parameters
|
||||
default_params = {
|
||||
"sampleCount": 1,
|
||||
}
|
||||
|
||||
# Merge with optional params
|
||||
parameters = {**default_params, **optional_params}
|
||||
|
||||
request_body = {
|
||||
"instances": [{"prompt": prompt}],
|
||||
"parameters": parameters,
|
||||
}
|
||||
|
||||
return request_body
|
||||
|
||||
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 Imagen image generation 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 = []
|
||||
|
||||
# Imagen format - predictions with generated images
|
||||
predictions = response_data.get("predictions", [])
|
||||
for prediction in predictions:
|
||||
# Imagen returns images as bytesBase64Encoded
|
||||
if "bytesBase64Encoded" in prediction:
|
||||
model_response.data.append(ImageObject(
|
||||
b64_json=prediction["bytesBase64Encoded"],
|
||||
url=None,
|
||||
))
|
||||
|
||||
return model_response
|
||||
|
||||
@ -6858,7 +6858,7 @@ def convert_to_dict(message: Union[BaseModel, dict]) -> dict:
|
||||
dict: The converted message.
|
||||
"""
|
||||
if isinstance(message, BaseModel):
|
||||
return message.model_dump(exclude_none=True)
|
||||
return message.model_dump(exclude_none=True) # type: ignore
|
||||
elif isinstance(message, dict):
|
||||
return message
|
||||
else:
|
||||
@ -7688,6 +7688,12 @@ class ProviderConfigManager:
|
||||
)
|
||||
|
||||
return get_runwayml_image_generation_config(model)
|
||||
elif LlmProviders.VERTEX_AI == provider:
|
||||
from litellm.llms.vertex_ai.image_generation import (
|
||||
get_vertex_ai_image_generation_config,
|
||||
)
|
||||
|
||||
return get_vertex_ai_image_generation_config(model)
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
|
||||
@ -119,6 +119,38 @@ class TestVertexImageGeneration(BaseImageGenTest):
|
||||
}
|
||||
|
||||
|
||||
class TestVertexAIGeminiImageGeneration(BaseImageGenTest):
|
||||
"""Test Gemini image generation models (Nano Banana)"""
|
||||
def get_base_image_generation_call_args(self) -> dict:
|
||||
# comment this when running locally
|
||||
load_vertex_ai_credentials()
|
||||
|
||||
litellm.in_memory_llm_clients_cache = InMemoryCache()
|
||||
return {
|
||||
"model": "vertex_ai/gemini-2.5-flash-image",
|
||||
"vertex_ai_project": "pathrise-convert-1606954137718",
|
||||
"vertex_ai_location": "us-central1",
|
||||
"n": 1,
|
||||
"size": "1024x1024",
|
||||
}
|
||||
|
||||
|
||||
class TestVertexAIGemini3ProImageGeneration(BaseImageGenTest):
|
||||
"""Test Gemini 3 Pro image generation model"""
|
||||
def get_base_image_generation_call_args(self) -> dict:
|
||||
# comment this when running locally
|
||||
load_vertex_ai_credentials()
|
||||
|
||||
litellm.in_memory_llm_clients_cache = InMemoryCache()
|
||||
return {
|
||||
"model": "vertex_ai/gemini-3-pro-image-preview",
|
||||
"vertex_ai_project": "pathrise-convert-1606954137718",
|
||||
"vertex_ai_location": "us-central1",
|
||||
"n": 1,
|
||||
"size": "1024x1024",
|
||||
}
|
||||
|
||||
|
||||
class TestBedrockNovaCanvasTextToImage(BaseImageGenTest):
|
||||
def get_base_image_generation_call_args(self) -> dict:
|
||||
litellm.in_memory_llm_clients_cache = InMemoryCache()
|
||||
|
||||
@ -0,0 +1,457 @@
|
||||
import os
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../.."))
|
||||
|
||||
from litellm.llms.vertex_ai.image_generation import (
|
||||
get_vertex_ai_image_generation_config,
|
||||
)
|
||||
from litellm.llms.vertex_ai.image_generation.vertex_gemini_transformation import (
|
||||
VertexAIGeminiImageGenerationConfig,
|
||||
)
|
||||
from litellm.llms.vertex_ai.image_generation.vertex_imagen_transformation import (
|
||||
VertexAIImagenImageGenerationConfig,
|
||||
)
|
||||
|
||||
|
||||
class TestVertexAIGeminiImageGenerationConfig:
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures"""
|
||||
self.config = VertexAIGeminiImageGenerationConfig()
|
||||
|
||||
def test_get_supported_openai_params(self):
|
||||
"""Test get_supported_openai_params returns correct params"""
|
||||
supported = self.config.get_supported_openai_params("gemini-2.5-flash-image")
|
||||
assert "n" in supported
|
||||
assert "size" in supported
|
||||
|
||||
def test_map_openai_params_n(self):
|
||||
"""Test mapping n parameter to candidate_count"""
|
||||
non_default_params = {"n": 3}
|
||||
optional_params = {}
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params, optional_params, "gemini-2.5-flash-image", False
|
||||
)
|
||||
assert result.get("candidate_count") == 3
|
||||
|
||||
def test_map_openai_params_size(self):
|
||||
"""Test mapping size parameter to aspectRatio"""
|
||||
non_default_params = {"size": "1024x1024"}
|
||||
optional_params = {}
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params, optional_params, "gemini-2.5-flash-image", False
|
||||
)
|
||||
assert result.get("aspectRatio") == "1:1"
|
||||
|
||||
def test_map_openai_params_size_16_9(self):
|
||||
"""Test mapping 16:9 size"""
|
||||
non_default_params = {"size": "1792x1024"}
|
||||
optional_params = {}
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params, optional_params, "gemini-2.5-flash-image", False
|
||||
)
|
||||
assert result.get("aspectRatio") == "16:9"
|
||||
|
||||
def test_map_size_to_aspect_ratio(self):
|
||||
"""Test size to aspect ratio mapping"""
|
||||
assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1"
|
||||
assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9"
|
||||
assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16"
|
||||
assert self.config._map_size_to_aspect_ratio("1280x896") == "4:3"
|
||||
assert self.config._map_size_to_aspect_ratio("896x1280") == "3:4"
|
||||
assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default
|
||||
|
||||
def test_transform_image_generation_request_basic(self):
|
||||
"""Test basic request transformation"""
|
||||
request = self.config.transform_image_generation_request(
|
||||
model="gemini-2.5-flash-image",
|
||||
prompt="A nano banana",
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
assert "contents" in request
|
||||
assert "generationConfig" in request
|
||||
assert request["generationConfig"]["responseModalities"] == ["IMAGE"]
|
||||
assert request["contents"][0]["parts"][0]["text"] == "A nano banana"
|
||||
|
||||
def test_transform_image_generation_request_with_aspect_ratio(self):
|
||||
"""Test request transformation with aspectRatio"""
|
||||
request = self.config.transform_image_generation_request(
|
||||
model="gemini-2.5-flash-image",
|
||||
prompt="A nano banana",
|
||||
optional_params={"aspectRatio": "16:9"},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
assert request["generationConfig"]["imageConfig"]["aspectRatio"] == "16:9"
|
||||
|
||||
def test_transform_image_generation_request_with_image_size(self):
|
||||
"""Test request transformation with imageSize (Gemini 3 Pro)"""
|
||||
request = self.config.transform_image_generation_request(
|
||||
model="gemini-3-pro-image-preview",
|
||||
prompt="A nano banana",
|
||||
optional_params={"imageSize": "4K"},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
assert request["generationConfig"]["imageConfig"]["imageSize"] == "4K"
|
||||
|
||||
def test_transform_image_generation_request_with_candidate_count(self):
|
||||
"""Test request transformation with candidate_count"""
|
||||
request = self.config.transform_image_generation_request(
|
||||
model="gemini-2.5-flash-image",
|
||||
prompt="A nano banana",
|
||||
optional_params={"candidate_count": 2},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
assert request["generationConfig"]["candidateCount"] == 2
|
||||
|
||||
def test_transform_image_generation_request_with_n(self):
|
||||
"""Test request transformation with n parameter"""
|
||||
request = self.config.transform_image_generation_request(
|
||||
model="gemini-2.5-flash-image",
|
||||
prompt="A nano banana",
|
||||
optional_params={"n": 2},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
assert request["generationConfig"]["candidateCount"] == 2
|
||||
|
||||
def test_transform_image_generation_response(self):
|
||||
"""Test response transformation"""
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"inlineData": {
|
||||
"mimeType": "image/png",
|
||||
"data": "base64_encoded_image_data",
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
mock_response.headers = {}
|
||||
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
model_response = ImageResponse()
|
||||
result = self.config.transform_image_generation_response(
|
||||
model="gemini-2.5-flash-image",
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert len(result.data) == 1
|
||||
assert result.data[0].b64_json == "base64_encoded_image_data"
|
||||
assert result.data[0].url is None
|
||||
|
||||
def test_transform_image_generation_response_multiple_images(self):
|
||||
"""Test response transformation with multiple images"""
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [
|
||||
{
|
||||
"inlineData": {
|
||||
"mimeType": "image/png",
|
||||
"data": "image1",
|
||||
}
|
||||
},
|
||||
{
|
||||
"inlineData": {
|
||||
"mimeType": "image/png",
|
||||
"data": "image2",
|
||||
}
|
||||
},
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
mock_response.headers = {}
|
||||
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
model_response = ImageResponse()
|
||||
result = self.config.transform_image_generation_response(
|
||||
model="gemini-2.5-flash-image",
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert len(result.data) == 2
|
||||
assert result.data[0].b64_json == "image1"
|
||||
assert result.data[1].b64_json == "image2"
|
||||
|
||||
|
||||
class TestVertexAIImagenImageGenerationConfig:
|
||||
def setup_method(self):
|
||||
"""Set up test fixtures"""
|
||||
self.config = VertexAIImagenImageGenerationConfig()
|
||||
|
||||
def test_get_supported_openai_params(self):
|
||||
"""Test get_supported_openai_params returns correct params"""
|
||||
supported = self.config.get_supported_openai_params("imagegeneration@006")
|
||||
assert "n" in supported
|
||||
assert "size" in supported
|
||||
|
||||
def test_map_openai_params_n(self):
|
||||
"""Test mapping n parameter to sampleCount"""
|
||||
non_default_params = {"n": 3}
|
||||
optional_params = {}
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params, optional_params, "imagegeneration@006", False
|
||||
)
|
||||
assert result.get("sampleCount") == 3
|
||||
|
||||
def test_map_openai_params_size(self):
|
||||
"""Test mapping size parameter to aspectRatio"""
|
||||
non_default_params = {"size": "1024x1024"}
|
||||
optional_params = {}
|
||||
result = self.config.map_openai_params(
|
||||
non_default_params, optional_params, "imagegeneration@006", False
|
||||
)
|
||||
assert result.get("aspectRatio") == "1:1"
|
||||
|
||||
def test_map_size_to_aspect_ratio(self):
|
||||
"""Test size to aspect ratio mapping"""
|
||||
assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1"
|
||||
assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9"
|
||||
assert self.config._map_size_to_aspect_ratio("unknown") == "1:1" # default
|
||||
|
||||
def test_transform_image_generation_request_basic(self):
|
||||
"""Test basic request transformation"""
|
||||
request = self.config.transform_image_generation_request(
|
||||
model="imagegeneration@006",
|
||||
prompt="A cat",
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
assert "instances" in request
|
||||
assert "parameters" in request
|
||||
assert request["instances"][0]["prompt"] == "A cat"
|
||||
assert request["parameters"]["sampleCount"] == 1
|
||||
|
||||
def test_transform_image_generation_request_with_params(self):
|
||||
"""Test request transformation with parameters"""
|
||||
request = self.config.transform_image_generation_request(
|
||||
model="imagegeneration@006",
|
||||
prompt="A cat",
|
||||
optional_params={"sampleCount": 2, "aspectRatio": "16:9"},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
assert request["parameters"]["sampleCount"] == 2
|
||||
assert request["parameters"]["aspectRatio"] == "16:9"
|
||||
|
||||
def test_transform_image_generation_response(self):
|
||||
"""Test response transformation"""
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"predictions": [
|
||||
{"bytesBase64Encoded": "base64_encoded_image_data"}
|
||||
]
|
||||
}
|
||||
mock_response.headers = {}
|
||||
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
model_response = ImageResponse()
|
||||
result = self.config.transform_image_generation_response(
|
||||
model="imagegeneration@006",
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert len(result.data) == 1
|
||||
assert result.data[0].b64_json == "base64_encoded_image_data"
|
||||
assert result.data[0].url is None
|
||||
|
||||
def test_transform_image_generation_response_multiple_images(self):
|
||||
"""Test response transformation with multiple images"""
|
||||
mock_response = MagicMock(spec=httpx.Response)
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {
|
||||
"predictions": [
|
||||
{"bytesBase64Encoded": "image1"},
|
||||
{"bytesBase64Encoded": "image2"},
|
||||
]
|
||||
}
|
||||
mock_response.headers = {}
|
||||
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
model_response = ImageResponse()
|
||||
result = self.config.transform_image_generation_response(
|
||||
model="imagegeneration@006",
|
||||
raw_response=mock_response,
|
||||
model_response=model_response,
|
||||
logging_obj=MagicMock(),
|
||||
request_data={},
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=None,
|
||||
)
|
||||
|
||||
assert len(result.data) == 2
|
||||
assert result.data[0].b64_json == "image1"
|
||||
assert result.data[1].b64_json == "image2"
|
||||
|
||||
|
||||
class TestGetVertexAIImageGenerationConfig:
|
||||
"""Test the router function that selects the correct config"""
|
||||
|
||||
def test_get_gemini_model_config(self):
|
||||
"""Test that Gemini models return Gemini config"""
|
||||
config = get_vertex_ai_image_generation_config("gemini-2.5-flash-image")
|
||||
assert isinstance(config, VertexAIGeminiImageGenerationConfig)
|
||||
|
||||
config = get_vertex_ai_image_generation_config("gemini-3-pro-image-preview")
|
||||
assert isinstance(config, VertexAIGeminiImageGenerationConfig)
|
||||
|
||||
config = get_vertex_ai_image_generation_config(
|
||||
"vertex_ai/gemini-2.5-flash-image"
|
||||
)
|
||||
assert isinstance(config, VertexAIGeminiImageGenerationConfig)
|
||||
|
||||
def test_get_imagen_model_config(self):
|
||||
"""Test that Imagen models return Imagen config"""
|
||||
config = get_vertex_ai_image_generation_config("imagegeneration@006")
|
||||
assert isinstance(config, VertexAIImagenImageGenerationConfig)
|
||||
|
||||
config = get_vertex_ai_image_generation_config("imagen-4.0-generate-001")
|
||||
assert isinstance(config, VertexAIImagenImageGenerationConfig)
|
||||
|
||||
config = get_vertex_ai_image_generation_config(
|
||||
"vertex_ai/imagegeneration@006"
|
||||
)
|
||||
assert isinstance(config, VertexAIImagenImageGenerationConfig)
|
||||
|
||||
def test_get_non_gemini_model_config(self):
|
||||
"""Test that non-Gemini models default to Imagen config"""
|
||||
config = get_vertex_ai_image_generation_config("some-other-model")
|
||||
assert isinstance(config, VertexAIImagenImageGenerationConfig)
|
||||
|
||||
|
||||
class TestVertexAIImageGenerationIntegration:
|
||||
"""Integration tests for Vertex AI image generation"""
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not os.getenv("VERTEXAI_PROJECT"),
|
||||
reason="Vertex AI credentials not set",
|
||||
)
|
||||
def test_gemini_image_generation_config_validation(self):
|
||||
"""Test that Gemini config can validate environment"""
|
||||
config = VertexAIGeminiImageGenerationConfig()
|
||||
with patch.object(
|
||||
config, "_resolve_vertex_project", return_value="test-project"
|
||||
), patch.object(
|
||||
config, "_resolve_vertex_location", return_value="us-central1"
|
||||
), patch.object(
|
||||
config, "_ensure_access_token", return_value=("token", None)
|
||||
):
|
||||
headers = config.validate_environment(
|
||||
headers={},
|
||||
model="gemini-2.5-flash-image",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert "Authorization" in headers
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not os.getenv("VERTEXAI_PROJECT"),
|
||||
reason="Vertex AI credentials not set",
|
||||
)
|
||||
def test_imagen_image_generation_config_validation(self):
|
||||
"""Test that Imagen config can validate environment"""
|
||||
config = VertexAIImagenImageGenerationConfig()
|
||||
with patch.object(
|
||||
config, "_resolve_vertex_project", return_value="test-project"
|
||||
), patch.object(
|
||||
config, "_resolve_vertex_location", return_value="us-central1"
|
||||
), patch.object(
|
||||
config, "_ensure_access_token", return_value=("token", None)
|
||||
):
|
||||
headers = config.validate_environment(
|
||||
headers={},
|
||||
model="imagegeneration@006",
|
||||
messages=[],
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert "Authorization" in headers
|
||||
|
||||
def test_gemini_get_complete_url(self):
|
||||
"""Test Gemini config URL generation"""
|
||||
config = VertexAIGeminiImageGenerationConfig()
|
||||
with patch.object(
|
||||
config, "_resolve_vertex_project", return_value="test-project"
|
||||
), patch.object(
|
||||
config, "_resolve_vertex_location", return_value="us-central1"
|
||||
):
|
||||
url = config.get_complete_url(
|
||||
api_base=None,
|
||||
api_key=None,
|
||||
model="gemini-2.5-flash-image",
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert "test-project" in url
|
||||
assert "us-central1" in url
|
||||
assert "gemini-2.5-flash-image" in url
|
||||
assert "generateContent" in url
|
||||
|
||||
def test_imagen_get_complete_url(self):
|
||||
"""Test Imagen config URL generation"""
|
||||
config = VertexAIImagenImageGenerationConfig()
|
||||
with patch.object(
|
||||
config, "_resolve_vertex_project", return_value="test-project"
|
||||
), patch.object(
|
||||
config, "_resolve_vertex_location", return_value="us-central1"
|
||||
):
|
||||
url = config.get_complete_url(
|
||||
api_base=None,
|
||||
api_key=None,
|
||||
model="imagegeneration@006",
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
)
|
||||
assert "test-project" in url
|
||||
assert "us-central1" in url
|
||||
assert "imagegeneration@006" in url
|
||||
assert "predict" in url
|
||||
|
||||
Loading…
Reference in New Issue
Block a user