diff --git a/py/packages/genkit-google-genai/pyproject.toml b/py/packages/genkit-google-genai/pyproject.toml index c2237af7c9..7315583c0a 100644 --- a/py/packages/genkit-google-genai/pyproject.toml +++ b/py/packages/genkit-google-genai/pyproject.toml @@ -50,11 +50,10 @@ dependencies = [ "structlog>=25.2.0", "strenum>=0.4.15; python_version < '3.11'", ] -description = "Gemini, Imagen, and Veo on Genkit" +description = "Gemini and Veo on Genkit" keywords = [ "genkit", "gemini", - "imagen", "veo", "vertex-ai", "google", diff --git a/py/packages/genkit-google-genai/src/genkit_google_genai/__init__.py b/py/packages/genkit-google-genai/src/genkit_google_genai/__init__.py index 3f3259c581..cdb4542a85 100644 --- a/py/packages/genkit-google-genai/src/genkit_google_genai/__init__.py +++ b/py/packages/genkit-google-genai/src/genkit_google_genai/__init__.py @@ -95,7 +95,6 @@ KnownGemma, VertexAIGeminiVersion, ) -from genkit_google_genai.models.imagen import ImagenConfigSchema, ImagenVersion, KnownImagen from genkit_google_genai.models.interactions_lyria import LyriaConfig from genkit_google_genai.models.interactions_registry import LyriaVersion from genkit_google_genai.models.veo import KnownVeo, VeoConfig, VeoVersion @@ -121,13 +120,10 @@ def package_name() -> str: 'DeepResearchConfig', 'GoogleAI', 'GoogleAIGeminiVersion', - 'ImagenConfigSchema', - 'ImagenVersion', 'KnownGemini', 'KnownGeminiImage', 'KnownGeminiTts', 'KnownGemma', - 'KnownImagen', 'KnownVeo', 'LyriaConfig', 'LyriaVersion', diff --git a/py/packages/genkit-google-genai/src/genkit_google_genai/google.py b/py/packages/genkit-google-genai/src/genkit_google_genai/google.py index a363832683..22e078b144 100644 --- a/py/packages/genkit-google-genai/src/genkit_google_genai/google.py +++ b/py/packages/genkit-google-genai/src/genkit_google_genai/google.py @@ -22,7 +22,7 @@ Google GenAI SDK to detect and register available models at runtime. Supported capabilities include text generation (Gemini/Gemma), text embeddings, -image generation (Imagen), and video generation (Veo). +image generation (Gemini native image), and video generation (Veo). Example: ```python @@ -92,7 +92,11 @@ family_embedder_ref, family_model_ref, ) -from genkit_google_genai.models._routing import is_unroutable_model_id, strip_ref_prefixes +from genkit_google_genai.models._routing import ( + is_unroutable_model_id, + is_unsupported_image_model_name, + strip_ref_prefixes, +) from genkit_google_genai.models.antigravity import AntigravityConfig, create_antigravity_action from genkit_google_genai.models.deep_research import ( DeepResearchConfig, @@ -122,15 +126,6 @@ is_tts_model, is_tuned_gemini_name, ) -from genkit_google_genai.models.imagen import ( - SUPPORTED_MODELS as IMAGE_SUPPORTED_MODELS, - ImagenConfigSchema, - ImagenModel, - KnownImagen, - is_imagen_model_name, - is_unsupported_image_model_name, - vertexai_image_model_info, -) from genkit_google_genai.models.interactions_lyria import ( LyriaConfig as InteractionsLyriaConfig, create_lyria_action, @@ -163,20 +158,17 @@ class GenaiModels: Attributes: gemini: List of Gemini/Gemma model names (generateContent action). - imagen: List of Imagen model names (predict action, Vertex AI only). embedders: List of embedding model names (embedContent action). veo: List of Veo video generation model names (generateVideos action). """ gemini: list[str] - imagen: list[str] embedders: list[str] veo: list[str] def __init__(self) -> None: """Initialize empty model lists.""" self.gemini = [] - self.imagen = [] self.embedders = [] self.veo = [] @@ -192,18 +184,16 @@ def _list_genai_models(client: genai.Client, is_vertex: bool) -> GenaiModels: - Google AI populates each model's ``supported_actions`` field, so models are categorized by action: - 'embedContent' action → embedders - - 'predict' + Imagen name (``imagen-``) → imagen - 'generateVideos' or Veo name (``veo-``) → veo - 'generateContent' + 'gemini'/'gemma' in name → gemini - Vertex AI returns ``supported_actions = None`` for every publisher model, so categorizing by action would skip them all. The Vertex path instead categorizes by model name: - - Imagen name (``imagen-``) → imagen - Veo name (``veo-``) → veo - 'gemini'/'gemma' in name (and not an embedding) → gemini - Ids with no working generate path here (``imagegeneration@*``, - ``imagetext@*``, ``virtual-try-on-*``) are not categorized at all, so - they are never advertised or registered. + Ids with no working generate path here (``imagen-*``, + ``imagegeneration@*``, ``imagetext@*``, ``virtual-try-on-*``) are not + categorized at all, so they are never advertised or registered. Embedders are intentionally NOT discovered here. The Vertex catalog over-lists embedders that are published but not callable, so they are advertised from a curated list (``VERTEX_KNOWN_EMBEDDERS``) instead. @@ -248,8 +238,6 @@ def _list_genai_models(client: genai.Client, is_vertex: bool) -> GenaiModels: continue elif is_unsupported_image_model_name(name): continue - elif is_imagen_model_name(name): - models.imagen.append(name) elif is_veo_model(name): models.veo.append(name) elif 'gemini' in lower_name or 'gemma' in lower_name: @@ -263,10 +251,6 @@ def _list_genai_models(client: genai.Client, is_vertex: bool) -> GenaiModels: if 'embedContent' in m.supported_actions: models.embedders.append(name) - # Imagen (imagen- prefix, not a bare "image" substring) - if 'predict' in m.supported_actions and is_imagen_model_name(name): - models.imagen.append(name) - # Veo if 'generateVideos' in m.supported_actions or is_veo_model(name): models.veo.append(name) @@ -500,21 +484,6 @@ def gemma_model( config=config, ) - @classmethod - def imagen_model( - cls, name: KnownImagen | str, *, config: ImagenConfigSchema | None = None - ) -> ModelRef[ImagenConfigSchema]: - """Typed ref for an Imagen model (``imagen-…``).""" - return family_model_ref( - name, - namespace=cls.name, - plugin_class=cls.__name__, - family='imagen', - method='imagen_model', - config_schema=ImagenConfigSchema, - config=config, - ) - @classmethod def veo_model(cls, name: KnownVeo | str, *, config: VeoConfig | None = None) -> ModelRef[VeoConfig]: """Typed ref for a Veo video model (``veo-…``).""" @@ -574,7 +543,6 @@ class GoogleAI(GoogleFamilyRefs, Plugin): | Type | Action Kind | Example | |---|---|---| | Gemini / Gemma | MODEL | ``googleai/gemini-flash-latest`` | - | Imagen | MODEL | ``googleai/imagen-3.0-generate-002`` | | Embedders | EMBEDDER | ``googleai/gemini-embedding-001`` | | Veo (Video) | BACKGROUND_MODEL | ``googleai/veo-3.1-generate-preview`` | @@ -740,11 +708,6 @@ async def init(self) -> list[Action]: if action := self._resolve_model(googleai_name(name)): actions.append(action) - # Imagen Models - for name in genai_models.imagen: - if action := self._resolve_model(googleai_name(name)): - actions.append(action) - # Veo Models (background models) for name in genai_models.veo: actions.extend(_background_actions(self._resolve_veo_model(googleai_name(name)))) @@ -785,15 +748,12 @@ async def init(self) -> list[Action]: return actions def _list_known_models(self) -> list[Action]: - """List known Gemini and Imagen models as Action objects.""" + """List known Gemini models as Action objects.""" genai_models = _list_genai_models(self._runtime_client(), is_vertex=False) actions = [] for name in genai_models.gemini: if action := self._resolve_model(googleai_name(name)): actions.append(action) - for name in genai_models.imagen: - if action := self._resolve_model(googleai_name(name)): - actions.append(action) return actions def _list_known_veo_models(self) -> list[Action]: @@ -901,15 +861,6 @@ def _resolve_model(self, name: str) -> Action | None: return None # One annotated closure per family. Action validates request.config # from the fn annotation; a single _run cannot switch schemas at runtime. - if is_imagen_model_name(clean_name): - model_info = vertexai_image_model_info(clean_name) - IMAGE_SUPPORTED_MODELS[clean_name] = model_info # pyright: ignore[reportArgumentType] - - async def _run_imagen(request: ModelRequest[ImagenConfigSchema], ctx: ActionRunContext) -> ModelResponse: - return await ImagenModel(clean_name, self._runtime_client()).generate(request, ctx) - - return _model_action(name, _run_imagen, model_info, ImagenConfigSchema) - model_info = google_model_info(clean_name) SUPPORTED_MODELS[clean_name] = model_info @@ -976,15 +927,6 @@ async def list_actions(self) -> list[ActionMetadata]: ) ) - for name in genai_models.imagen: - actions_list.append( - model_action_metadata( - name=googleai_name(name), - info=vertexai_image_model_info(name).model_dump(by_alias=True), - config_schema=ImagenConfigSchema, - ) - ) - for name in genai_models.veo: actions_list.append(_veo_background_action_metadata(googleai_name(name))) @@ -1040,7 +982,7 @@ class VertexAI(GoogleFamilyRefs, Plugin): """Vertex AI plugin for Genkit with dynamic model discovery. This plugin provides access to Google Cloud Vertex AI models including - Gemini, Imagen, Veo, and embedders. Models are discovered dynamically, + Gemini, Veo, and embedders. Models are discovered dynamically, ensuring new models are available without SDK updates. Vertex AI vs Google AI: @@ -1049,14 +991,12 @@ class VertexAI(GoogleFamilyRefs, Plugin): - Customer-managed encryption keys (CMEK) - Data residency controls - IAM-based access control - - Imagen image generation models Model Types: | Type | Action Kind | Example | |---|---|---| | Gemini / Gemma | MODEL | ``vertexai/gemini-flash-latest`` | - | Imagen | MODEL | ``vertexai/imagen-3.0-generate-002`` | - | Veo (Video) | BACKGROUND_MODEL | ``vertexai/veo-3.1-generate-preview`` | + | Veo (Video) | BACKGROUND_MODEL | ``vertexai/veo-3.1-generate-001`` | | Embedders | EMBEDDER | ``vertexai/text-embedding-005`` | Example: @@ -1185,10 +1125,6 @@ async def init(self) -> list[Action]: if action := self._resolve_model(vertexai_name(name)): actions.append(action) - for name in genai_models.imagen: - if action := self._resolve_model(vertexai_name(name)): - actions.append(action) - # Veo Models (background models) for name in genai_models.veo: bg_action = self._resolve_veo_model(vertexai_name(name)) @@ -1226,9 +1162,6 @@ def _list_known_models(self) -> list[Action]: for name in genai_models.gemini: if action := self._resolve_model(vertexai_name(name)): actions.append(action) - for name in genai_models.imagen: - if action := self._resolve_model(vertexai_name(name)): - actions.append(action) return actions def _list_known_veo_models(self) -> list[Action]: @@ -1368,15 +1301,6 @@ async def _run_tuned(request: ModelRequest[GeminiConfigSchema], ctx: ActionRunCo return _model_action(name, _run_tuned, model_info, GeminiConfigSchema) - if is_imagen_model_name(clean_name): - model_info = vertexai_image_model_info(clean_name) - IMAGE_SUPPORTED_MODELS[clean_name] = model_info # pyright: ignore[reportArgumentType] - - async def _run_imagen(request: ModelRequest[ImagenConfigSchema], ctx: ActionRunContext) -> ModelResponse: - return await ImagenModel(clean_name, self._runtime_client()).generate(request, ctx) - - return _model_action(name, _run_imagen, model_info, ImagenConfigSchema) - model_info = google_model_info(clean_name) SUPPORTED_MODELS[clean_name] = model_info @@ -1443,15 +1367,6 @@ async def list_actions(self) -> list[ActionMetadata]: ) ) - for name in genai_models.imagen: - actions_list.append( - model_action_metadata( - name=vertexai_name(name), - info=vertexai_image_model_info(name).model_dump(by_alias=True), - config_schema=ImagenConfigSchema, - ) - ) - for name in genai_models.veo: actions_list.append(_veo_background_action_metadata(vertexai_name(name))) diff --git a/py/packages/genkit-google-genai/src/genkit_google_genai/models/_model_refs.py b/py/packages/genkit-google-genai/src/genkit_google_genai/models/_model_refs.py index 5679f6e7c8..12126fc933 100644 --- a/py/packages/genkit-google-genai/src/genkit_google_genai/models/_model_refs.py +++ b/py/packages/genkit-google-genai/src/genkit_google_genai/models/_model_refs.py @@ -17,7 +17,7 @@ """Shared machinery for the typed family ref constructors. Each plugin class exposes one constructor per config family -(``gemini_model``, ``imagen_model``, ...). The return type is the contract: +(``gemini_model``, ``gemini_image_model``, ...). The return type is the contract: ``ModelRef[GeminiConfigSchema]`` tells generate-time code exactly which config is legal, so a constructor must refuse ids whose runtime action would validate a different schema — otherwise the ref lies and the wrong @@ -41,7 +41,6 @@ 'tts': 'gemini_tts_model', 'image': 'gemini_image_model', 'gemma': 'gemma_model', - 'imagen': 'imagen_model', 'veo': 'veo_model', 'embedder': 'embedding', } @@ -65,7 +64,10 @@ def wrong_family_error(*, plugin_class: str, method: str, family: str, local: st else f"'{local}' has no ref constructor in this plugin." ) elif actual == 'unsupported': - hint = f"'{local}' is not a supported model." + if local.lower().startswith('imagen-'): + hint = f"'{local}' is not a supported model; for image generation use {plugin_class}.gemini_image_model()." + else: + hint = f"'{local}' is not a supported model." elif actual in FAMILY_METHOD: hint = f"'{local}' is not a {family} model; use {plugin_class}.{FAMILY_METHOD[actual]}()." else: diff --git a/py/packages/genkit-google-genai/src/genkit_google_genai/models/_routing.py b/py/packages/genkit-google-genai/src/genkit_google_genai/models/_routing.py index de61e8e432..26c0d6c646 100644 --- a/py/packages/genkit-google-genai/src/genkit_google_genai/models/_routing.py +++ b/py/packages/genkit-google-genai/src/genkit_google_genai/models/_routing.py @@ -28,10 +28,6 @@ is_image_model, is_tts_model, ) -from genkit_google_genai.models.imagen import ( - is_imagen_model_name, - is_unsupported_image_model_name, -) from genkit_google_genai.models.lyria import is_lyria_model from genkit_google_genai.models.veo import is_veo_model @@ -49,8 +45,8 @@ ) # Families with no MODEL generate path on this plugin. Constructor -# refusal is a different table — TTS, native image, Gemma, and Imagen -# still resolve as MODEL. +# refusal is a different table — TTS, native image, and Gemma still +# resolve as MODEL. UNROUTABLE_FAMILIES = frozenset({ 'embedder', 'unsupported', @@ -74,6 +70,17 @@ def strip_ref_prefixes(name: str) -> str: return local +def is_unsupported_image_model_name(name: str) -> bool: + """True for image ids with no generate path here. + + Matches the ``imagen-``, ``imagegeneration@``, ``imagetext@`` and + ``virtual-try-on-`` prefixes on the last path segment. Gemini native + image (``gemini-…-image``) is a different family and routes normally. + """ + local = name.split('/')[-1].lower() + return local.startswith(('imagen-', 'imagegeneration@', 'imagetext@', 'virtual-try-on-')) + + def classify_family(name: str) -> str: """Bucket a model id by the last path segment. @@ -93,8 +100,6 @@ def classify_family(name: str) -> str: return 'deep-research' if leaf.startswith('antigravity-'): return 'antigravity' - if is_imagen_model_name(leaf): - return 'imagen' if is_tts_model(leaf): return 'tts' if is_image_model(leaf): diff --git a/py/packages/genkit-google-genai/src/genkit_google_genai/models/gemini.py b/py/packages/genkit-google-genai/src/genkit_google_genai/models/gemini.py index 0f17e8767d..fa6db95589 100644 --- a/py/packages/genkit-google-genai/src/genkit_google_genai/models/gemini.py +++ b/py/packages/genkit-google-genai/src/genkit_google_genai/models/gemini.py @@ -960,9 +960,9 @@ def is_tts_model(name: str) -> bool: def is_image_model(name: str) -> bool: """Check if the model is a Gemini native image generation model. - Native image is a ``gemini-`` name that contains ``-image``. Imagen - (``imagen-…``) is a different family — a bare ``image`` substring - would catch both. + Native image is a ``gemini-`` name that contains ``-image``. The + ``gemini-`` prefix is required: a bare ``image`` substring would also + match ``imagen-`` ids, which have no generate path here. Args: name: The model name to check. diff --git a/py/packages/genkit-google-genai/src/genkit_google_genai/models/imagen.py b/py/packages/genkit-google-genai/src/genkit_google_genai/models/imagen.py deleted file mode 100644 index 0b74ca0ac0..0000000000 --- a/py/packages/genkit-google-genai/src/genkit_google_genai/models/imagen.py +++ /dev/null @@ -1,291 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - -"""Imagen model implementation for Google GenAI plugin.""" - -import base64 -import sys - -if sys.version_info < (3, 11): - from strenum import StrEnum -else: - from enum import StrEnum - -import json -from functools import cached_property -from typing import Any, Literal, TypeAlias - -from google import genai -from google.genai import types as genai_types -from google.genai.errors import APIError -from pydantic import BaseModel, ConfigDict, ValidationError - -from genkit import ( - GenkitError, - Media, - MediaPart, - Message, - ModelInfo, - ModelRequest, - ModelResponse, - Part, - Role, - Supports, - TextPart, -) -from genkit.plugin_api import ActionRunContext, tracer, wrap_http_error -from genkit_google_genai.models._sdk_config import ( - attach_leftovers, - dump_family_config, - sdk_config_error, - split_sdk_fields, -) - - -def _to_dict(obj: Any) -> Any: # noqa: ANN401 - """Convert object to dict if it's a Pydantic model, otherwise return as-is.""" - return obj.model_dump() if isinstance(obj, BaseModel) else obj - - -class ImagenVersion(StrEnum): - """Supported text-to-image models.""" - - IMAGEN3 = 'imagen-3.0-generate-002' - IMAGEN3_FAST = 'imagen-3.0-fast-generate-001' - - -# Quote autocomplete needs a Literal. The enum above is the catalog; a test -# requires these members and the enum values to be the same set. -KnownImagen: TypeAlias = Literal[ - 'imagen-3.0-generate-002', - 'imagen-3.0-fast-generate-001', -] - - -SUPPORTED_MODELS = { - ImagenVersion.IMAGEN3: ModelInfo( - label='Vertex AI - Imagen3', - supports=Supports( - media=True, - multiturn=False, - tools=False, - system_role=True, - output=['media'], - ), - ), - ImagenVersion.IMAGEN3_FAST: ModelInfo( - label='Vertex AI - Imagen3 Fast', - supports=Supports( - media=False, - multiturn=False, - tools=False, - system_role=True, - output=['media'], - ), - ), -} - -DEFAULT_IMAGE_SUPPORT = Supports( - media=True, - multiturn=False, - tools=False, - system_role=True, - output=['media'], -) - - -def is_imagen_model_name(name: str) -> bool: - """Return True if ``name`` is an Imagen model. - - Imagen ids start with ``imagen-`` on the local name after stripping the - plugin / ``models/`` prefix. Gemini native image (``gemini-…-image``) is - not Imagen. - """ - return name.split('/')[-1].lower().startswith('imagen-') - - -def is_unsupported_image_model_name(name: str) -> bool: - """Return True for image ids that must not route anywhere. - - ``imagegeneration@*`` and ``imagetext@*`` were shut down by Google in - June 2026, and ``virtual-try-on-*`` needs a person+product image request - shape this plugin does not implement. Letting any of those fall through - to the Gemini default would answer with the wrong model, so callers - treat these ids as not-a-model instead. - """ - local = name.split('/')[-1].lower() - return local.startswith('imagegeneration@') or local.startswith('imagetext@') or local.startswith('virtual-try-on-') - - -def vertexai_image_model_info( - version: str, -) -> ModelInfo: - """Generates a ModelInfo object. - - This function tries to get the best ModelInfo Supports - for the given version. - - Args: - version: Version of the model. - - Returns: - ModelInfo object. - """ - return ModelInfo( - label=f'Vertex AI - {version}', - supports=DEFAULT_IMAGE_SUPPORT, - ) - - -class ImagenConfigSchema(BaseModel): - """Imagen Config Schema.""" - - model_config = ConfigDict(extra='allow') - - -class ImagenModel: - """Imagen text-to-image model.""" - - def __init__(self, version: str | ImagenVersion, client: genai.Client) -> None: - """Initialize Imagen model. - - Args: - version: Imagen version - client: Google AI client - """ - self._version = version - self._client = client - - def _build_prompt(self, request: ModelRequest) -> str: - """Build prompt request from Genkit request. - - Args: - request: Genkit request. - - Returns: - prompt for Imagen - """ - prompt = [] - for message in request.messages: - for part in message.content: - if isinstance(part.root, TextPart): - prompt.append(part.root.text) - else: - raise GenkitError(status='INVALID_ARGUMENT', message='Non-text messages are not supported') - return ' '.join(prompt) - - async def generate(self, request: ModelRequest, _: ActionRunContext) -> ModelResponse: - """Handle a generation request. - - Args: - request: The generation request containing messages and parameters. - _: action context - - Returns: - The model's response to the generation request. - """ - prompt = self._build_prompt(request) - config = self._get_config(request) - if request.tools: - raise GenkitError(status='UNIMPLEMENTED', message='Tools are not supported for this model.') - - with tracer.start_as_current_span('generate_images') as span: - span.set_attribute( - 'genkit:input', - json.dumps({ - 'config': _to_dict(config), - 'contents': prompt, - 'model': self._version, - }), - ) - try: - response = await self._client.aio.models.generate_images( - model=self._version, prompt=prompt, config=config - ) - except APIError as e: - raise wrap_http_error(e, status_code=e.code, message=e.message or str(e)) from e - span.set_attribute('genkit:output', json.dumps(_to_dict(response), default=str)) - - content = self._contents_from_response(response) - - return ModelResponse( - message=Message( - content=content, - role=Role.MODEL, - ) - ) - - def _get_config(self, request: ModelRequest) -> genai_types.GenerateImagesConfig | None: - dumped = dump_family_config( - config=request.config, - expected_type=ImagenConfigSchema, - action_name=self._version, - ) - if not dumped: - return None - - known, leftovers = split_sdk_fields(dumped, genai_types.GenerateImagesConfig) - try: - cfg = genai_types.GenerateImagesConfig(**known) if known else genai_types.GenerateImagesConfig() - except ValidationError as e: - raise sdk_config_error(action_name=self._version, error=e) from e - return attach_leftovers(cfg, leftovers, nest='parameters') - - def _contents_from_response(self, response: genai_types.GenerateImagesResponse) -> list: - """Retrieve contents from google-genai response. - - Args: - response: google-genai response. - - Returns: - list of generated contents. - """ - content = [] - if response.generated_images: - for image in response.generated_images: - if image.image and image.image.image_bytes: - b64_data = base64.b64encode(image.image.image_bytes).decode('utf-8') - content.append( - Part( - root=MediaPart( - media=Media( - url=f'data:{image.image.mime_type};base64,{b64_data}', - content_type=image.image.mime_type, - ) - ) - ) - ) - - return content - - @cached_property - def metadata(self) -> dict: - """Model metadata. - - Returns: - model metadata. - """ - supports = {} - if self._version in SUPPORTED_MODELS: - model_supports = SUPPORTED_MODELS[self._version].supports # pyright: ignore[reportArgumentType] - if model_supports: - supports = model_supports.model_dump(by_alias=True) - else: - model_supports = vertexai_image_model_info(self._version).supports - if model_supports: - supports = model_supports.model_dump(by_alias=True) - - return {'model': {'supports': supports}} diff --git a/py/packages/genkit-google-genai/src/genkit_google_genai/models/veo.py b/py/packages/genkit-google-genai/src/genkit_google_genai/models/veo.py index 0862ddfa6e..9329b8f04a 100644 --- a/py/packages/genkit-google-genai/src/genkit_google_genai/models/veo.py +++ b/py/packages/genkit-google-genai/src/genkit_google_genai/models/veo.py @@ -62,30 +62,27 @@ class VeoVersion(StrEnum): """Supported Veo video generation models. Note: Models are discovered dynamically. This enum provides convenience - constants for commonly used Veo models. + constants for commonly used Veo models. Google AI serves Veo 3.1 as the + ``-preview`` ids; Vertex AI serves it as the ``-001`` ids. """ - VEO_2_0 = 'veo-2.0-generate-001' - VEO_2_0_EXP = 'veo-2.0-generate-exp' - VEO_3_0 = 'veo-3.0-generate-001' - VEO_3_0_FAST = 'veo-3.0-fast-generate-001' VEO_3_1_PREVIEW = 'veo-3.1-generate-preview' VEO_3_1_FAST_PREVIEW = 'veo-3.1-fast-generate-preview' + VEO_3_1_LITE_PREVIEW = 'veo-3.1-lite-generate-preview' VEO_3_1 = 'veo-3.1-generate-001' VEO_3_1_FAST = 'veo-3.1-fast-generate-001' + VEO_3_1_LITE = 'veo-3.1-lite-generate-001' # Quote autocomplete needs a Literal. The enum above is the catalog; a test # requires these members and the enum values to be the same set. KnownVeo: TypeAlias = Literal[ - 'veo-2.0-generate-001', - 'veo-2.0-generate-exp', - 'veo-3.0-generate-001', - 'veo-3.0-fast-generate-001', 'veo-3.1-generate-preview', 'veo-3.1-fast-generate-preview', + 'veo-3.1-lite-generate-preview', 'veo-3.1-generate-001', 'veo-3.1-fast-generate-001', + 'veo-3.1-lite-generate-001', ] diff --git a/py/packages/genkit-google-genai/test/google_plugin_test.py b/py/packages/genkit-google-genai/test/google_plugin_test.py index 78bf78fd51..1c0f2ca17a 100644 --- a/py/packages/genkit-google-genai/test/google_plugin_test.py +++ b/py/packages/genkit-google-genai/test/google_plugin_test.py @@ -34,10 +34,6 @@ GeminiConfigSchema, GeminiModel, ) -from genkit_google_genai.models.imagen import ( - DEFAULT_IMAGE_SUPPORT, - SUPPORTED_MODELS as IMAGE_SUPPORTED_MODELS, -) from google import genai from google.auth.credentials import Credentials from google.genai.types import HttpOptions @@ -414,7 +410,7 @@ class MockModel: description: str = '' models_return_value = [ - MockModel(supported_actions=['generateVideos'], name='models/veo-2.0-generate-001'), + MockModel(supported_actions=['generateVideos'], name='models/veo-3.1-generate-preview'), ] mock_client = MagicMock() @@ -424,7 +420,7 @@ class MockModel: result = googleai_plugin_instance._list_known_veo_models() # Check Veo - action1 = next(a for a in result if a.name == googleai_name('veo-2.0-generate-001')) + action1 = next(a for a in result if a.name == googleai_name('veo-3.1-generate-preview')) assert action1 is not None @@ -712,53 +708,42 @@ async def test_vertexai_resolve_action_embedder( 'genkit_google_genai.models.gemini.google_model_info', new_callable=MagicMock, ) -@patch( - 'genkit_google_genai.models.imagen.vertexai_image_model_info', - new_callable=MagicMock, -) @pytest.mark.parametrize( - 'model_name, expected_model_name, key, image', + 'model_name, expected_model_name, key', [ ( 'gemini-pro-deluxe-max', 'vertexai/gemini-pro-deluxe-max', 'gemini-pro-deluxe-max', - False, ), ( 'vertexai/gemini-pro-deluxe-max', 'vertexai/gemini-pro-deluxe-max', 'gemini-pro-deluxe-max', - False, ), - # A bare "image" prefix is not Imagen; only imagen- ids route there. + # A bare "image" substring still resolves as a Gemini text model. ( 'vertexai/image-gemini-pro-deluxe-max', 'vertexai/image-gemini-pro-deluxe-max', 'image-gemini-pro-deluxe-max', - False, ), ( 'image-gemini-pro-deluxe-max', 'vertexai/image-gemini-pro-deluxe-max', 'image-gemini-pro-deluxe-max', - False, ), ( 'gemini-pro-deluxe-max-image', 'vertexai/gemini-pro-deluxe-max-image', 'gemini-pro-deluxe-max-image', - False, ), ], ) def test_vertexai__resolve_model( mock_google_model_info: MagicMock, - mock_vertexai_image_model_info: MagicMock, model_name: str, expected_model_name: str, key: str, - image: bool, vertexai_plugin_instance: VertexAI, ) -> None: """Tests for VertexAI._resolve_model method.""" @@ -770,21 +755,12 @@ def test_vertexai__resolve_model( supports=DEFAULT_SUPPORTS_MODEL, ) - mock_vertexai_image_model_info.return_value = ModelInfo( - label=f'Vertex AI - {model_name}', - supports=DEFAULT_IMAGE_SUPPORT, - ) - action = plugin._resolve_model(name=expected_model_name) assert action is not None assert action.kind == ActionKind.MODEL assert action.name == expected_model_name - - if image: - assert key in IMAGE_SUPPORTED_MODELS - else: - assert key in SUPPORTED_MODELS + assert key in SUPPORTED_MODELS @pytest.mark.parametrize( @@ -821,19 +797,6 @@ def test_vertexai__resolve_embedder( @pytest.mark.asyncio async def test_vertexai_list_actions(vertexai_plugin_instance: VertexAI) -> None: """Unit test for list actions.""" - - @dataclass - class MockModel: - name: str - description: str = '' - - [ - MockModel(name='publishers/google/models/gemini-1.5-flash'), - MockModel(name='publishers/google/models/gemini-embedding-001'), - MockModel(name='publishers/google/models/imagen-3.0-generate-001'), - MockModel(name='publishers/google/models/veo-2.0-generate-001'), - ] - mock_client = MagicMock() # Create sophisticated mocks that have supported_actions m1 = MagicMock() @@ -848,11 +811,11 @@ class MockModel: m3 = MagicMock() m3.name = 'publishers/google/models/imagen-3.0-generate-001' - m3.supported_actions = ['predict'] # Imagen uses predict + m3.supported_actions = ['predict'] m3.description = 'Imagen' m4 = MagicMock() - m4.name = 'publishers/google/models/veo-2.0-generate-001' + m4.name = 'publishers/google/models/veo-3.1-generate-001' m4.supported_actions = ['generateVideos'] # Veo uses generateVideos m4.description = 'Veo' @@ -869,13 +832,11 @@ class MockModel: action2 = next(a for a in result if a.name == vertexai_name('gemini-embedding-001')) assert action2 is not None - # Verify Imagen - action3 = next(a for a in result if a.name == vertexai_name('imagen-3.0-generate-001')) - assert action3 is not None - assert action3.action_type == ActionKind.MODEL + # Imagen has no generate path here, so it is never advertised. + assert not any(a.name == vertexai_name('imagen-3.0-generate-001') for a in result) # Verify Veo - action4 = next(a for a in result if a.name == vertexai_name('veo-2.0-generate-001')) + action4 = next(a for a in result if a.name == vertexai_name('veo-3.1-generate-001')) assert action4 is not None # from genkit_google_genai.models.veo import VeoConfigSchema # assert action4.config_schema == VeoConfigSchema @@ -903,7 +864,7 @@ def mock_model(name: str) -> MagicMock: mock_model('publishers/google/models/gemini-embedding-001'), mock_model('publishers/google/models/gemini-embedding-2'), mock_model('publishers/google/models/imagen-3.0-generate-002'), - mock_model('publishers/google/models/veo-2.0-generate-001'), + mock_model('publishers/google/models/veo-3.1-generate-001'), ] vertexai_plugin_instance._runtime_client = lambda: mock_client @@ -912,9 +873,9 @@ def mock_model(name: str) -> MagicMock: # Gemini text model discovered despite supported_actions=None. assert vertexai_name('gemini-2.5-pro') in names - # Imagen and Veo discovered. - assert vertexai_name('imagen-3.0-generate-002') in names - assert vertexai_name('veo-2.0-generate-001') in names + # Veo discovered; Imagen is not advertised. + assert vertexai_name('veo-3.1-generate-001') in names + assert vertexai_name('imagen-3.0-generate-002') not in names # gemini-embedding-001 is registered as an embedder, not a gemini model. embedder = next(a for a in result if a.name == vertexai_name('gemini-embedding-001')) @@ -928,10 +889,12 @@ async def test_googleai_resolve_background_model(googleai_plugin_instance: Googl """Test resolve action for background model.""" plugin = googleai_plugin_instance - action = await plugin.resolve(action_type=ActionKind.BACKGROUND_MODEL, name=googleai_name('veo-2.0-generate-001')) + action = await plugin.resolve( + action_type=ActionKind.BACKGROUND_MODEL, name=googleai_name('veo-3.1-generate-preview') + ) assert action is not None assert action.kind == ActionKind.BACKGROUND_MODEL - assert action.name == googleai_name('veo-2.0-generate-001') + assert action.name == googleai_name('veo-3.1-generate-preview') @pytest.mark.asyncio @@ -940,29 +903,16 @@ async def test_googleai_resolve_check_operation(googleai_plugin_instance: Google plugin = googleai_plugin_instance action = await plugin.resolve( - action_type=ActionKind.CHECK_OPERATION, name=googleai_name('veo-2.0-generate-001/check') + action_type=ActionKind.CHECK_OPERATION, name=googleai_name('veo-3.1-generate-preview/check') ) assert action is not None assert action.kind == ActionKind.CHECK_OPERATION - assert action.name == googleai_name('veo-2.0-generate-001/check') + assert action.name == googleai_name('veo-3.1-generate-preview/check') @pytest.mark.asyncio async def test_vertexai_list_known_models(vertexai_plugin_instance: VertexAI) -> None: """Unit test for list known models.""" - - @dataclass - class MockModel: - name: str - description: str = '' - - [ - MockModel(name='publishers/google/models/gemini-1.5-flash'), - MockModel(name='publishers/google/models/gemini-embedding-001'), - MockModel(name='publishers/google/models/imagen-3.0-generate-001'), - MockModel(name='publishers/google/models/veo-2.0-generate-001'), - ] - mock_client = MagicMock() # Create sophisticated mocks that have supported_actions m1 = MagicMock() @@ -981,7 +931,7 @@ class MockModel: m3.description = 'Imagen' m4 = MagicMock() - m4.name = 'publishers/google/models/veo-2.0-generate-001' + m4.name = 'publishers/google/models/veo-3.1-generate-001' m4.supported_actions = ['generateVideos'] m4.description = 'Veo' @@ -994,12 +944,11 @@ class MockModel: action1 = next(a for a in result if a.name == vertexai_name('gemini-1.5-flash')) assert action1 is not None - # Verify Imagen - action3 = next(a for a in result if a.name == vertexai_name('imagen-3.0-generate-001')) - assert action3 is not None + # Imagen has no generate path here, so it is not a known MODEL. + assert not any(a.name == vertexai_name('imagen-3.0-generate-001') for a in result) # Veo is background-only, so it is not a known generate MODEL. - assert not any(a.name == vertexai_name('veo-2.0-generate-001') for a in result) + assert not any(a.name == vertexai_name('veo-3.1-generate-001') for a in result) veo_actions = vertexai_plugin_instance._list_known_veo_models() assert {a.kind for a in veo_actions} == {ActionKind.BACKGROUND_MODEL, ActionKind.CHECK_OPERATION} diff --git a/py/packages/genkit-google-genai/test/models/family_name_test.py b/py/packages/genkit-google-genai/test/models/family_name_test.py index aa6f615722..1225e132b2 100644 --- a/py/packages/genkit-google-genai/test/models/family_name_test.py +++ b/py/packages/genkit-google-genai/test/models/family_name_test.py @@ -17,41 +17,21 @@ """Family name checks for Google model routing.""" import pytest -from genkit_google_genai.models._routing import classify_family, is_unroutable_model_id +from genkit_google_genai.models._routing import ( + classify_family, + is_unroutable_model_id, + is_unsupported_image_model_name, +) from genkit_google_genai.models.gemini import ( is_gemini_model, is_gemma_model, is_image_model, is_tts_model, ) -from genkit_google_genai.models.imagen import ( - is_imagen_model_name, - is_unsupported_image_model_name, -) from genkit_google_genai.models.lyria import is_lyria_model from genkit_google_genai.models.veo import is_veo_model -@pytest.mark.parametrize( - ('name', 'expected'), - [ - ('imagen-3.0-generate-002', True), - ('IMAGEN-3.0-generate-002', True), - ('googleai/imagen-3.0-generate-002', True), - ('vertexai/imagen-3.0-generate-002', True), - ('models/imagen-3.0-generate-002', True), - ('gemini-2.5-flash-image', False), - ('gemini-2.5-flash-image-preview', False), - ('imagegeneration@006', False), - ('virtual-try-on-001', False), - ('veo-3.0-generate-001', False), - ], -) -def test_is_imagen_model_name(name: str, expected: bool) -> None: - """Imagen is the ``imagen-`` prefix only, on both plugins.""" - assert is_imagen_model_name(name) is expected - - @pytest.mark.parametrize( ('name', 'expected'), [ @@ -62,13 +42,18 @@ def test_is_imagen_model_name(name: str, expected: bool) -> None: ('vertexai/imagetext@001', True), ('virtual-try-on-001', True), ('vertexai/virtual-try-on-001', True), - ('imagen-3.0-generate-002', False), + ('imagen-3.0-generate-002', True), + ('IMAGEN-4.0-generate-001', True), + ('googleai/imagen-4.0-generate-001', True), + ('models/imagen-4.0-ultra-generate-001', True), ('gemini-2.5-flash-image', False), + ('gemini-2.5-flash-image-preview', False), ('gemini-2.5-flash', False), + ('veo-3.1-generate-001', False), ], ) def test_is_unsupported_image_model_name(name: str, expected: bool) -> None: - """Retired / unimplemented image ids fail closed instead of routing to Gemini.""" + """Image ids with no generate path fail closed instead of routing to Gemini.""" assert is_unsupported_image_model_name(name) is expected @@ -137,9 +122,9 @@ def test_is_gemini_model(name: str, expected: bool) -> None: @pytest.mark.parametrize( ('name', 'expected'), [ - ('veo-3.0-generate-001', True), + ('veo-3.1-generate-001', True), ('googleai/veo-3.1-generate-preview', True), - ('VEO-2.0-generate-001', True), + ('VEO-3.1-generate-001', True), ('gemini-2.0-flash', False), ('devotional-hymn', False), ], @@ -174,7 +159,9 @@ def test_is_lyria_model(name: str, expected: bool) -> None: 'imagegeneration@006', 'imagetext@001', 'virtual-try-on-001', - 'veo-3.0-generate-001', + 'imagen-3.0-generate-002', + 'vertexai/imagen-4.0-generate-001', + 'veo-3.1-generate-001', 'googleai/lyria-002', 'models/deep-research-pro-preview', 'googleai/deep-research-pro-preview', @@ -199,6 +186,7 @@ def test_unroutable_ids_fail_closed(name: str) -> None: ('publishers/google/models/antigravity-preview-05-2026', 'antigravity'), ('lyria-002', 'lyria'), ('imagetext@001', 'unsupported'), + ('imagen-4.0-generate-001', 'unsupported'), ], ) def test_classify_family_uses_leaf_name(name: str, family: str) -> None: @@ -211,7 +199,7 @@ def test_classify_family_uses_leaf_name(name: str, family: str) -> None: [ 'gemini-2.5-flash', 'GEMINI-2.5-FLASH-PREVIEW-TTS', - 'imagen-3.0-generate-002', + 'gemini-2.5-flash-image', 'gemma-3-12b-it', ], ) diff --git a/py/packages/genkit-google-genai/test/models/googlegenai_imagen_test.py b/py/packages/genkit-google-genai/test/models/googlegenai_imagen_test.py deleted file mode 100644 index c6c7d096e0..0000000000 --- a/py/packages/genkit-google-genai/test/models/googlegenai_imagen_test.py +++ /dev/null @@ -1,136 +0,0 @@ -# Copyright 2025 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -# -# SPDX-License-Identifier: Apache-2.0 - - -"""Tests for the Imagen model implementation.""" - -import base64 -from unittest.mock import MagicMock - -import pytest -from genkit_google_genai.models.imagen import ImagenConfigSchema, ImagenModel, ImagenVersion -from google import genai -from pytest_mock import MockerFixture - -from genkit import ( - ActionRunContext, - GenkitError, - MediaPart, - Message, - ModelRequest, - ModelResponse, - Part, - Role, - TextPart, -) - - -@pytest.mark.asyncio -@pytest.mark.parametrize('version', [x for x in ImagenVersion]) -async def test_generate_media_response(mocker: MockerFixture, version: ImagenVersion) -> None: - """Test generate method for media responses.""" - request_text = 'response question' - response_byte_string = b'\x89PNG\r\n\x1a\n' - response_mimetype = 'image/png' - - request = ModelRequest( - messages=[ - Message( - role=Role.USER, - content=[ - Part(root=TextPart(text=request_text)), - ], - ), - ], - ) - - response_images = genai.types.GenerateImagesResponse( - generated_images=[ - genai.types.GeneratedImage( - image=genai.types.Image(image_bytes=response_byte_string, mime_type=response_mimetype) - ) - ] - ) - - googleai_client_mock = mocker.AsyncMock() - googleai_client_mock.aio.models.generate_images.return_value = response_images - - imagen = ImagenModel(version, googleai_client_mock) - - ctx = ActionRunContext() - response = await imagen.generate(request, ctx) - - googleai_client_mock.assert_has_calls([ - mocker.call.aio.models.generate_images(model=version, prompt=request_text, config=None) - ]) - assert isinstance(response, ModelResponse) - assert response.message is not None - content = response.message.content[0] - assert isinstance(content.root, MediaPart) - - assert content.root.media.content_type == response_mimetype - - # Verify the data URL contains the correct base64-encoded content - # Data URLs have format: data:;base64, - data_url = content.root.media.url - assert data_url.startswith(f'data:{response_mimetype};base64,') - encoded_data = data_url.split(',', 1)[1] - assert base64.b64decode(encoded_data) == response_byte_string - - -def test_imagen_unknown_extra_rides_on_extra_body() -> None: - """Leftover keys ride on extra_body so a newly supported field still reaches the API.""" - imagen = ImagenModel(ImagenVersion.IMAGEN3, MagicMock()) - request = ModelRequest( - messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='a cat'))])], - config=ImagenConfigSchema.model_validate({'fooBar': 1}), - ) - - cfg = imagen._get_config(request) - - assert cfg is not None - assert cfg.http_options is not None - assert cfg.http_options.extra_body == {'parameters': {'fooBar': 1}} - - -def test_imagen_rejects_raw_dicts() -> None: - """A dict at the dump leaf means Action never produced the family instance.""" - imagen = ImagenModel(ImagenVersion.IMAGEN3, MagicMock()) - request = ModelRequest( - messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='a cat'))])], - config={'number_of_images': 1}, # type: ignore[arg-type] - ) - - with pytest.raises(GenkitError) as exc_info: - imagen._get_config(request) - - assert exc_info.value.status == 'INVALID_ARGUMENT' - assert imagen._version in str(exc_info.value) - - -def test_imagen_invalid_sdk_field_is_invalid_argument() -> None: - """SDK type errors become a named INVALID_ARGUMENT.""" - imagen = ImagenModel(ImagenVersion.IMAGEN3, MagicMock()) - request = ModelRequest( - messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='a cat'))])], - config=ImagenConfigSchema.model_validate({'number_of_images': 'nope'}), - ) - - with pytest.raises(GenkitError) as exc_info: - imagen._get_config(request) - - assert exc_info.value.status == 'INVALID_ARGUMENT' - assert 'number_of_images' in str(exc_info.value) diff --git a/py/packages/genkit-google-genai/test/models/model_refs_test.py b/py/packages/genkit-google-genai/test/models/model_refs_test.py index 02d9e37557..6263b736d1 100644 --- a/py/packages/genkit-google-genai/test/models/model_refs_test.py +++ b/py/packages/genkit-google-genai/test/models/model_refs_test.py @@ -16,10 +16,12 @@ """Tests for the typed family ref constructors on GoogleAI / VertexAI.""" +import importlib from collections.abc import Callable from enum import Enum from typing import get_args +import genkit_google_genai import pytest from genkit_google_genai import ( GoogleAI, @@ -27,7 +29,6 @@ KnownGeminiImage, KnownGeminiTts, KnownGemma, - KnownImagen, KnownVeo, VertexAI, ) @@ -44,11 +45,6 @@ is_image_model, is_tts_model, ) -from genkit_google_genai.models.imagen import ( - ImagenConfigSchema, - ImagenVersion, - is_imagen_model_name, -) from genkit_google_genai.models.veo import VeoConfig, VeoVersion, is_veo_model from genkit import GenkitError @@ -72,23 +68,22 @@ def test_gemini_model_both_plugins(self) -> None: def test_enum_names_still_work(self) -> None: """The existing version enums remain valid constructor input.""" assert GoogleAI.gemini_model(GoogleAIGeminiVersion.GEMINI_2_5_FLASH).name == 'googleai/gemini-2.5-flash' - assert GoogleAI.imagen_model(ImagenVersion.IMAGEN3).name == 'googleai/imagen-3.0-generate-002' assert GoogleAI.veo_model(VeoVersion.VEO_3_1_FAST_PREVIEW).name == 'googleai/veo-3.1-fast-generate-preview' assert VertexAI.veo_model(VeoVersion.VEO_3_1).name == 'vertexai/veo-3.1-generate-001' + assert GoogleAI.veo_model(VeoVersion.VEO_3_1_LITE_PREVIEW).name == 'googleai/veo-3.1-lite-generate-preview' + assert VertexAI.veo_model(VeoVersion.VEO_3_1_LITE).name == 'vertexai/veo-3.1-lite-generate-001' def test_family_constructors_type_their_config(self) -> None: """Each family constructor carries its own config schema.""" tts = GoogleAI.gemini_tts_model('gemini-2.5-flash-preview-tts') image = VertexAI.gemini_image_model('gemini-2.5-flash-image') gemma = GoogleAI.gemma_model('gemma-3-12b-it') - imagen = VertexAI.imagen_model('imagen-3.0-generate-002') veo = GoogleAI.veo_model('veo-3.1-fast-generate-preview') assert tts.config_schema is GeminiTtsConfigSchema assert image.config_schema is GeminiImageConfigSchema + assert image.name == 'vertexai/gemini-2.5-flash-image' assert gemma.config_schema is GemmaConfigSchema - assert imagen.config_schema is ImagenConfigSchema - assert imagen.name == 'vertexai/imagen-3.0-generate-002' assert veo.config_schema is VeoConfig assert veo.name == 'googleai/veo-3.1-fast-generate-preview' @@ -104,7 +99,7 @@ def test_unknown_id_allowed_on_gemini_model_and_embedding(self) -> None: assert GoogleAI.embedding('totally-new-embedder').name == 'googleai/totally-new-embedder' with pytest.raises(GenkitError): - GoogleAI.imagen_model('totally-new-model') + GoogleAI.gemma_model('totally-new-model') class TestStripThenPrefix: @@ -161,7 +156,7 @@ class TestClosedRejectSet: @pytest.mark.parametrize( 'bad_id', [ - 'veo-3.0-generate-001', # wrong family: has its own constructor + 'veo-3.1-generate-preview', # wrong family: has its own constructor 'lyria-002', # no constructor in this plugin 'googleai/lyria-002', # prefix must not defeat the gate 'deep-research-pro-preview', # Interactions API family @@ -169,7 +164,8 @@ class TestClosedRejectSet: 'imagegeneration@006', # retired June 2026 'virtual-try-on-001', # predict shape not implemented 'gemini-embedding-001', # embedder, not a generate model - 'imagen-3.0-generate-002', # wrong family: has its own constructor + 'imagen-3.0-generate-002', # not a supported model + 'imagen-4.0-generate-001', # not a supported model 'gemini-2.5-flash-preview-tts', # wrong family: TTS 'gemini-2.5-flash-image', # wrong family: native image 'gemma-3-12b-it', # wrong family: Gemma @@ -186,11 +182,11 @@ def test_error_points_at_the_right_constructor(self) -> None: with pytest.raises(GenkitError, match=r'gemini_tts_model'): GoogleAI.gemini_model('gemini-2.5-flash-preview-tts') with pytest.raises(GenkitError, match=r'VertexAI\.gemini_model'): - VertexAI.imagen_model('gemini-2.5-flash') + VertexAI.gemma_model('gemini-2.5-flash') with pytest.raises(GenkitError, match=r'embedding'): GoogleAI.gemini_model('gemini-embedding-001') with pytest.raises(GenkitError, match=r'veo_model'): - GoogleAI.gemini_model('veo-3.0-generate-001') + GoogleAI.gemini_model('veo-3.1-generate-preview') with pytest.raises(GenkitError, match=r'lyria_model'): GoogleAI.gemini_model('lyria-002') with pytest.raises(GenkitError, match=r'deep_research_model'): @@ -205,13 +201,15 @@ def test_error_points_at_the_right_constructor(self) -> None: GoogleAI.gemini_model('imagegeneration@006') with pytest.raises(GenkitError, match=r'is not a supported model'): GoogleAI.gemini_model('virtual-try-on-001') - with pytest.raises(GenkitError, match=r'is not a imagen model'): - GoogleAI.imagen_model('totally-new-model') + with pytest.raises(GenkitError, match=r'GoogleAI\.gemini_image_model'): + GoogleAI.gemini_model('imagen-4.0-generate-001') + with pytest.raises(GenkitError, match=r'VertexAI\.gemini_image_model'): + VertexAI.gemini_image_model('imagen-3.0-generate-002') + with pytest.raises(GenkitError, match=r'is not a gemma model'): + GoogleAI.gemma_model('totally-new-model') def test_family_constructors_reject_other_families(self) -> None: """Non-gemini constructors take only their own family ids.""" - with pytest.raises(GenkitError): - GoogleAI.imagen_model('gemini-2.5-flash') with pytest.raises(GenkitError): GoogleAI.gemini_tts_model('gemini-2.5-flash') with pytest.raises(GenkitError): @@ -274,7 +272,48 @@ def test_family_literals_match_their_catalog(self) -> None: assert set(get_args(KnownGeminiTts)) == _family_catalog(is_tts_model) assert set(get_args(KnownGeminiImage)) == _family_catalog(is_image_model) assert set(get_args(KnownGemma)) == _family_catalog(is_gemma_model) - assert set(get_args(KnownImagen)) == {str(member.value) for member in ImagenVersion} - assert all(is_imagen_model_name(value) for value in get_args(KnownImagen)) assert set(get_args(KnownVeo)) == {str(member.value) for member in VeoVersion} assert all(is_veo_model(value) for value in get_args(KnownVeo)) + + def test_veo_catalog_is_the_3_1_family(self) -> None: + """The Veo catalog is the 3.1 family only: preview ids on Google AI, 001 ids on Vertex AI.""" + assert {str(member.value) for member in VeoVersion} == { + 'veo-3.1-generate-preview', + 'veo-3.1-fast-generate-preview', + 'veo-3.1-lite-generate-preview', + 'veo-3.1-generate-001', + 'veo-3.1-fast-generate-001', + 'veo-3.1-lite-generate-001', + } + + +class TestNoImagenSurface: + """Imagen ids have no constructor, no export, and no module.""" + + def test_no_imagen_exports(self) -> None: + """Neither the package nor the plugin classes expose an Imagen name.""" + assert not {name for name in genkit_google_genai.__all__ if 'Imagen' in name} + assert not hasattr(GoogleAI, 'imagen_model') + assert not hasattr(VertexAI, 'imagen_model') + + def test_no_imagen_module(self) -> None: + """The Imagen model module is gone.""" + with pytest.raises(ModuleNotFoundError): + importlib.import_module('genkit_google_genai.models.imagen') + + @pytest.mark.parametrize('plugin', [GoogleAI, VertexAI]) + def test_every_constructor_rejects_imagen_ids(self, plugin: type[GoogleAI] | type[VertexAI]) -> None: + """An imagen- id is refused by every family constructor with the image hint.""" + for method in ('gemini_model', 'gemini_tts_model', 'gemini_image_model', 'gemma_model'): + with pytest.raises(GenkitError, match=r'for image generation use \w+\.gemini_image_model\(\)'): + getattr(plugin, method)('imagen-4.0-generate-001') + with pytest.raises(GenkitError, match=r'for image generation use \w+\.gemini_image_model\(\)'): + plugin.embedding('imagen-4.0-generate-001') + + @pytest.mark.parametrize('bad_id', ['imagegeneration@006', 'imagetext@001', 'virtual-try-on-001']) + def test_other_unsupported_ids_omit_the_image_hint(self, bad_id: str) -> None: + """Only imagen- ids are redirected to image generation.""" + with pytest.raises(GenkitError) as exc_info: + GoogleAI.gemini_model(bad_id) + assert 'is not a supported model' in str(exc_info.value) + assert 'gemini_image_model' not in str(exc_info.value) diff --git a/py/packages/genkit-google-genai/tests/google_genai_plugin_test.py b/py/packages/genkit-google-genai/tests/google_genai_plugin_test.py index ed43d937c9..a2aa5abcb3 100644 --- a/py/packages/genkit-google-genai/tests/google_genai_plugin_test.py +++ b/py/packages/genkit-google-genai/tests/google_genai_plugin_test.py @@ -47,7 +47,6 @@ GeminiTtsConfigSchema, GemmaConfigSchema, ) -from genkit_google_genai.models.imagen import ImagenConfigSchema from genkit_google_genai.models.veo import VeoConfig, VeoModel from genkit import ActionKind, Genkit, GenkitError, Message, ModelRequest, Part, Role, TextPart @@ -83,7 +82,7 @@ def test_googleai_name() -> None: def test_vertexai_name() -> None: """Test vertexai_name helper function.""" assert vertexai_name('gemini-2.0-flash') == 'vertexai/gemini-2.0-flash' - assert vertexai_name('imagen-3.0-generate-001') == 'vertexai/imagen-3.0-generate-001' + assert vertexai_name('gemini-2.5-flash-image') == 'vertexai/gemini-2.5-flash-image' def test_plugin_names() -> None: @@ -179,7 +178,6 @@ def test_genai_models_container() -> None: """Test GenaiModels container initialization.""" models = GenaiModels() assert models.gemini == [] - assert models.imagen == [] assert models.embedders == [] assert models.veo == [] @@ -202,24 +200,10 @@ async def test_googleai_resolve_model(mock_list_models: MagicMock, mock_client: @patch('genkit_google_genai.google.genai.client.Client') @patch('genkit_google_genai.google._list_genai_models') @pytest.mark.asyncio -async def test_googleai_resolve_imagen_model(mock_list_models: MagicMock, mock_client: MagicMock) -> None: - """Test GoogleAI plugin resolves Imagen image generation models.""" - mock_list_models.return_value = GenaiModels() - - plugin = GoogleAI(api_key='test-key') - action = await plugin.resolve(ActionKind.MODEL, 'googleai/imagen-3.0-generate-002') - - assert action is not None - assert action.kind == ActionKind.MODEL - assert action.name == 'googleai/imagen-3.0-generate-002' - assert _custom_options(action) == to_json_schema(ImagenConfigSchema) - - -@patch('genkit_google_genai.google.genai.client.Client') -@patch('genkit_google_genai.google._list_genai_models') -@pytest.mark.asyncio -async def test_googleai_resolve_gemini_image_is_not_imagen(mock_list_models: MagicMock, mock_client: MagicMock) -> None: - """Native Gemini image models must not route through Imagen.""" +async def test_googleai_resolve_gemini_image_uses_image_config( + mock_list_models: MagicMock, mock_client: MagicMock +) -> None: + """Native Gemini image models validate the image config schema.""" mock_list_models.return_value = GenaiModels() plugin = GoogleAI(api_key='test-key') @@ -241,7 +225,6 @@ async def test_googleai_resolve_gemini_image_is_not_imagen(mock_list_models: Mag ('googleai/gemini-2.5-flash-preview-tts', GeminiTtsConfigSchema), ('googleai/gemini-2.5-flash-image', GeminiImageConfigSchema), ('googleai/gemma-3-12b-it', GemmaConfigSchema), - ('googleai/imagen-3.0-generate-002', ImagenConfigSchema), ], ) async def test_googleai_resolve_types_family_config( @@ -271,7 +254,6 @@ async def test_googleai_resolve_types_family_config( ('vertexai/gemini-2.5-flash-preview-tts', GeminiTtsConfigSchema), ('vertexai/gemini-2.5-flash-image', GeminiImageConfigSchema), ('vertexai/gemma-3-12b-it', GemmaConfigSchema), - ('vertexai/imagen-3.0-generate-002', ImagenConfigSchema), ], ) async def test_vertexai_resolve_types_family_config( @@ -321,8 +303,8 @@ async def test_veo_start_types_family_config(mock_list_models: MagicMock, mock_c mock_list_models.return_value = GenaiModels() for plugin, name in ( - (GoogleAI(api_key='test-key'), 'googleai/veo-3.0-generate-001'), - (VertexAI(project='test-project'), 'vertexai/veo-3.0-generate-001'), + (GoogleAI(api_key='test-key'), 'googleai/veo-3.1-generate-preview'), + (VertexAI(project='test-project'), 'vertexai/veo-3.1-generate-001'), ): action = await plugin.resolve(ActionKind.BACKGROUND_MODEL, name) assert action is not None @@ -363,7 +345,7 @@ async def test_veo_action_run_dumps_leftover_and_stamps(mock_list_models: MagicM mock_client.return_value.aio.models.generate_videos = AsyncMock(return_value=op) plugin = VertexAI(project='test-project') - action = await plugin.resolve(ActionKind.BACKGROUND_MODEL, 'vertexai/veo-3.0-generate-001') + action = await plugin.resolve(ActionKind.BACKGROUND_MODEL, 'vertexai/veo-3.1-generate-001') assert action is not None started = await action.run({ @@ -378,7 +360,7 @@ async def test_veo_action_run_dumps_leftover_and_stamps(mock_list_models: MagicM assert cfg.duration_seconds == 5 assert cfg.http_options is not None assert cfg.http_options.extra_body == {'parameters': {'fooBar': 1}} - assert started.response.action == '/background-model/vertexai/veo-3.0-generate-001' + assert started.response.action == '/background-model/vertexai/veo-3.1-generate-001' @patch('genkit_google_genai.google.genai.client.Client') @@ -389,7 +371,7 @@ async def test_veo_action_run_rejects_bad_duration(mock_list_models: MagicMock, mock_list_models.return_value = GenaiModels() plugin = VertexAI(project='test-project') - action = await plugin.resolve(ActionKind.BACKGROUND_MODEL, 'vertexai/veo-3.0-generate-001') + action = await plugin.resolve(ActionKind.BACKGROUND_MODEL, 'vertexai/veo-3.1-generate-001') assert action is not None with pytest.raises(GenkitError) as exc_info: @@ -410,46 +392,69 @@ async def test_veo_check_is_typed(mock_list_models: MagicMock, mock_client: Magi mock_list_models.return_value = GenaiModels() plugin = VertexAI(project='test-project') - action = await plugin.resolve(ActionKind.CHECK_OPERATION, 'vertexai/veo-3.0-generate-001/check') + action = await plugin.resolve(ActionKind.CHECK_OPERATION, 'vertexai/veo-3.1-generate-001/check') assert action is not None hints = get_type_hints(action._fn) # noqa: SLF001 assert hints['op'] is Operation assert hints['return'] is Operation -@patch('genkit_google_genai.google.genai.client.Client') -@patch('genkit_google_genai.google._list_genai_models') -@pytest.mark.asyncio -async def test_googleai_init_registers_imagen_models(mock_list_models: MagicMock, mock_client: MagicMock) -> None: - """Test GoogleAI init registers Imagen models from dynamic discovery.""" - models = GenaiModels() - models.imagen = ['imagen-3.0-generate-002'] - mock_list_models.return_value = models +def test_list_genai_models_googleai_skips_imagen() -> None: + """A predict-only ``imagen-`` entry lands in no bucket.""" - plugin = GoogleAI(api_key='test-key') - actions = await plugin.init() + def _model(name: str, actions: list[str]) -> MagicMock: + item = MagicMock() + item.name = name + item.supported_actions = actions + item.description = '' + return item - imagen_actions = [a for a in actions if 'imagen' in a.name] - assert len(imagen_actions) == 1 - assert imagen_actions[0].name == 'googleai/imagen-3.0-generate-002' - assert imagen_actions[0].kind == ActionKind.MODEL + client = MagicMock() + client.models.list.return_value = [ + _model('models/gemini-2.5-flash', ['generateContent']), + _model('models/imagen-4.0-generate-001', ['predict']), + _model('models/imagen-4.0-ultra-generate-001', ['predict', 'generateContent']), + ] + catalog = _list_genai_models(client, is_vertex=False) + assert vars(catalog) == {'gemini': ['gemini-2.5-flash'], 'embedders': [], 'veo': []} @patch('genkit_google_genai.google.genai.client.Client') -@patch('genkit_google_genai.google._list_genai_models') @pytest.mark.asyncio -async def test_googleai_list_actions_includes_imagen(mock_list_models: MagicMock, mock_client: MagicMock) -> None: - """Test GoogleAI list_actions includes Imagen models.""" - models = GenaiModels() - models.imagen = ['imagen-3.0-generate-002'] - mock_list_models.return_value = models +@pytest.mark.parametrize('backend', ['googleai', 'vertexai']) +async def test_list_actions_never_advertise_imagen(mock_client: MagicMock, backend: str) -> None: + """An ``imagen-`` id served by the API reaches neither list_actions nor init.""" - plugin = GoogleAI(api_key='test-key') - actions_list = await plugin.list_actions() + def _model(name: str, actions: list[str] | None) -> MagicMock: + item = MagicMock() + item.name = name + item.supported_actions = actions + item.description = '' + return item - imagen_actions = [a for a in actions_list if 'imagen' in a.name] - assert len(imagen_actions) == 1 - assert imagen_actions[0].name == 'googleai/imagen-3.0-generate-002' + if backend == 'googleai': + plugin: GoogleAI | VertexAI = GoogleAI(api_key='test-key') + listing = [ + _model('models/gemini-2.5-flash', ['generateContent']), + _model('models/imagen-4.0-generate-001', ['predict']), + ] + else: + plugin = VertexAI(project='test-project') + listing = [ + _model('publishers/google/models/gemini-2.5-flash', None), + _model('publishers/google/models/imagen-4.0-generate-001', None), + ] + mock_client.return_value.models.list.return_value = listing + + listed = await plugin.list_actions() + registered = await plugin.init() + + # The Gemini id proves the listing reached discovery, so the imagen + # assertions below are not passing on an empty catalog. + assert any('gemini-2.5-flash' in a.name for a in listed) + assert any('gemini-2.5-flash' in a.name for a in registered) + assert not any('imagen' in a.name for a in listed) + assert not any('imagen' in a.name for a in registered) @patch('genkit_google_genai.google.genai.client.Client') @@ -503,6 +508,8 @@ async def test_vertexai_resolve_model(mock_list_models: MagicMock, mock_client: 'virtual-try-on-001', 'imagegeneration@006', 'imagetext@001', + 'imagen-3.0-generate-002', + 'imagen-4.0-generate-001', 'lyria-002', 'deep-research-pro-preview', 'gemini-embedding-001', @@ -531,6 +538,8 @@ async def test_vertexai_unroutable_ids_fail_closed( 'virtual-try-on-001', 'imagegeneration@006', 'imagetext@001', + 'imagen-3.0-generate-002', + 'imagen-4.0-generate-001', 'deep-research-pro-preview', 'gemini-embedding-001', 'models/deep-research-pro-preview', @@ -557,7 +566,7 @@ async def test_googleai_resolve_veo_as_model_returns_none(mock_list_models: Magi mock_list_models.return_value = GenaiModels() plugin = GoogleAI(api_key='test-key') - action = await plugin.resolve(ActionKind.MODEL, 'googleai/veo-3.0-generate-001') + action = await plugin.resolve(ActionKind.MODEL, 'googleai/veo-3.1-generate-preview') assert action is None @@ -570,7 +579,7 @@ async def test_vertexai_resolve_veo_as_model_returns_none(mock_list_models: Magi mock_list_models.return_value = GenaiModels() plugin = VertexAI(project='test-project') - action = await plugin.resolve(ActionKind.MODEL, 'vertexai/veo-3.0-generate-001') + action = await plugin.resolve(ActionKind.MODEL, 'vertexai/veo-3.1-generate-001') assert action is None @@ -583,11 +592,11 @@ async def test_resolve_model_finds_veo_as_background(mock_list_models: MagicMock mock_list_models.return_value = GenaiModels() ai = Genkit(plugins=[GoogleAI(api_key='test-key')]) - action = await ai.registry.resolve_model('googleai/veo-3.0-generate-001') + action = await ai.registry.resolve_model('googleai/veo-3.1-generate-preview') assert action is not None assert action.kind == ActionKind.BACKGROUND_MODEL - assert action.name == 'googleai/veo-3.0-generate-001' + assert action.name == 'googleai/veo-3.1-generate-preview' @patch('genkit_google_genai.models.veo.genai.Client') @@ -613,7 +622,7 @@ async def test_generate_and_check_operation_apply_veo_context_and_config( context = {'secrets': {'api_key': 'tenant-key'}} operation = await ai.generate_operation( - model='googleai/veo-3.0-generate-001', + model='googleai/veo-3.1-generate-preview', prompt='a cat walking', config={'aspectRatio': '16:9', 'baseUrl': 'https://request.example', 'apiVersion': 'v1'}, context=context, @@ -650,18 +659,18 @@ async def test_vertexai_resolve_veo_background_model(mock_list_models: MagicMock mock_list_models.return_value = GenaiModels() plugin = VertexAI(project='test-project') - start = await plugin.resolve(ActionKind.BACKGROUND_MODEL, 'vertexai/veo-3.0-generate-001') - check = await plugin.resolve(ActionKind.CHECK_OPERATION, 'vertexai/veo-3.0-generate-001/check') + start = await plugin.resolve(ActionKind.BACKGROUND_MODEL, 'vertexai/veo-3.1-generate-001') + check = await plugin.resolve(ActionKind.CHECK_OPERATION, 'vertexai/veo-3.1-generate-001/check') assert start is not None assert start.kind == ActionKind.BACKGROUND_MODEL - assert start.name == 'vertexai/veo-3.0-generate-001' + assert start.name == 'vertexai/veo-3.1-generate-001' model_meta = cast('dict[str, object]', start.metadata['model']) supports = cast('dict[str, object]', model_meta['supports']) assert supports['longRunning'] is True assert check is not None assert check.kind == ActionKind.CHECK_OPERATION - assert check.name == 'vertexai/veo-3.0-generate-001/check' + assert check.name == 'vertexai/veo-3.1-generate-001/check' @patch('genkit_google_genai.google.create_vertex_evaluators') @@ -673,7 +682,7 @@ async def test_vertexai_init_registers_veo_as_background( ) -> None: """Vertex init registers Veo start/check, never a blocking MODEL action.""" models = GenaiModels() - models.veo = ['veo-3.0-generate-001'] + models.veo = ['veo-3.1-generate-001'] mock_list_models.return_value = models mock_evaluators.return_value = [] @@ -691,7 +700,7 @@ async def test_vertexai_init_registers_veo_as_background( async def test_list_actions_advertises_veo_as_background(mock_list_models: MagicMock, mock_client: MagicMock) -> None: """Both plugins list Veo with the kind that resolve() actually serves.""" models = GenaiModels() - models.veo = ['veo-3.0-generate-001'] + models.veo = ['veo-3.1-generate-001'] mock_list_models.return_value = models googleai_actions = await GoogleAI(api_key='test-key').list_actions() @@ -700,12 +709,12 @@ async def test_list_actions_advertises_veo_as_background(mock_list_models: Magic for actions, plugin_name in ((googleai_actions, 'googleai'), (vertexai_actions, 'vertexai')): veo_entries = [a for a in actions if 'veo' in a.name] assert len(veo_entries) == 1 - assert veo_entries[0].name == f'{plugin_name}/veo-3.0-generate-001' + assert veo_entries[0].name == f'{plugin_name}/veo-3.1-generate-001' assert veo_entries[0].action_type == ActionKind.BACKGROUND_MODEL -def test_list_genai_models_vertex_skips_substring_veo_and_retired_image() -> None: - """Discovery buckets on the ``veo-`` prefix, not a ``veo`` substring.""" +def test_list_genai_models_vertex_skips_substring_veo_and_unsupported_image() -> None: + """Discovery buckets on the ``veo-`` prefix and drops unsupported image ids.""" def _model(name: str) -> MagicMock: item = MagicMock() @@ -717,17 +726,20 @@ def _model(name: str) -> MagicMock: client = MagicMock() client.models.list.return_value = [ _model('publishers/google/models/gemini-2.5-flash'), - _model('publishers/google/models/veo-3.0-generate-001'), + _model('publishers/google/models/veo-3.1-generate-001'), _model('publishers/google/models/braveo-lab'), _model('publishers/google/models/imagegeneration@006'), _model('publishers/google/models/virtual-try-on-001'), _model('publishers/google/models/imagetext@001'), + _model('publishers/google/models/imagen-3.0-generate-002'), + _model('publishers/google/models/imagen-4.0-generate-001'), ] catalog = _list_genai_models(client, is_vertex=True) - assert catalog.veo == ['veo-3.0-generate-001'] - assert catalog.imagen == [] - assert 'imagetext@001' not in catalog.gemini - assert 'braveo-lab' not in catalog.gemini + assert vars(catalog) == { + 'gemini': ['gemini-2.5-flash'], + 'embedders': [], + 'veo': ['veo-3.1-generate-001'], + } @patch('genkit_google_genai.google.genai.client.Client') @@ -737,19 +749,19 @@ async def test_veo_start_stamps_background_action_key(mock_list_models: MagicMoc """Start and check stamp ``/background-model/{name}`` so a later check can resolve.""" mock_list_models.return_value = GenaiModels() plugin = VertexAI(project='test-project') - start = await plugin.resolve(ActionKind.BACKGROUND_MODEL, 'vertexai/veo-3.0-generate-001') - check = await plugin.resolve(ActionKind.CHECK_OPERATION, 'vertexai/veo-3.0-generate-001/check') + start = await plugin.resolve(ActionKind.BACKGROUND_MODEL, 'vertexai/veo-3.1-generate-001') + check = await plugin.resolve(ActionKind.CHECK_OPERATION, 'vertexai/veo-3.1-generate-001/check') assert start is not None assert check is not None request = ModelRequest(messages=[Message(role=Role.USER, content=[Part(root=TextPart(text='a clip'))])]) with patch.object(VeoModel, 'start', new=AsyncMock(return_value=Operation(id='ops/1'))): started = await start.run(request) - assert started.response.action == '/background-model/vertexai/veo-3.0-generate-001' + assert started.response.action == '/background-model/vertexai/veo-3.1-generate-001' with patch.object(VeoModel, 'check', new=AsyncMock(return_value=Operation(id='ops/1'))): checked = await check.run(Operation(id='ops/1')) - assert checked.response.action == '/background-model/vertexai/veo-3.0-generate-001' + assert checked.response.action == '/background-model/vertexai/veo-3.1-generate-001' @patch('genkit_google_genai.google.genai.client.Client') diff --git a/py/packages/genkit-google-genai/tests/veo_test.py b/py/packages/genkit-google-genai/tests/veo_test.py index 0d44063b45..5fc9cf2e61 100644 --- a/py/packages/genkit-google-genai/tests/veo_test.py +++ b/py/packages/genkit-google-genai/tests/veo_test.py @@ -24,7 +24,6 @@ from genkit_google_genai.models.veo import ( VeoConfig, VeoModel, - VeoVersion, _from_veo_operation, is_veo_model, ) @@ -77,11 +76,11 @@ class TestIsVeoModel: def test_veo_model_name(self) -> None: """Veo model names are recognized.""" - assert is_veo_model('veo-2.0-generate-001') is True + assert is_veo_model('veo-3.1-generate-001') is True def test_veo_uppercase(self) -> None: """Case-insensitive matching works.""" - assert is_veo_model('VEO-2.0-generate-001') is True + assert is_veo_model('VEO-3.1-generate-001') is True def test_non_veo_model(self) -> None: """Non-Veo model names are rejected.""" @@ -89,30 +88,13 @@ def test_non_veo_model(self) -> None: def test_namespaced_veo_model(self) -> None: """Plugin prefixes are stripped before the ``veo-`` check.""" - assert is_veo_model('googleai/veo-3.0-generate-001') is True + assert is_veo_model('googleai/veo-3.1-generate-preview') is True def test_substring_veo_is_rejected(self) -> None: """A bare ``veo`` substring is not enough; the id has to start with ``veo-``.""" assert is_veo_model('devotional-hymn') is False -class TestVeoVersion: - """Tests for VeoVersion enum convenience constants.""" - - @pytest.mark.parametrize( - 'version', - [ - VeoVersion.VEO_3_1_PREVIEW, - VeoVersion.VEO_3_1_FAST_PREVIEW, - VeoVersion.VEO_3_0, - VeoVersion.VEO_3_0_FAST, - ], - ) - def test_new_googleai_models_are_recognized(self, version: VeoVersion) -> None: - """New Veo 3.0/3.1 model constants map to valid Veo names.""" - assert is_veo_model(version.value) is True - - class TestFromVeoOperation: """``_from_veo_operation`` reads the SDK operation and resolves playable media.""" @@ -341,7 +323,7 @@ async def test_start_passes_generate_videos_config_and_returns_ticket(self) -> N """A typed Veo config dumps aspectRatio / durationSeconds onto generate_videos.""" client = MagicMock() client.aio.models.generate_videos = AsyncMock(return_value=_sdk_op(name='operations/1', done=False)) - veo = VeoModel('veo-3.0-generate-001', client) + veo = VeoModel('veo-3.1-generate-001', client) request = _text_request( config=VeoConfig.model_validate({'aspectRatio': '16:9', 'durationSeconds': 5, 'fooBar': 1}), ) @@ -361,7 +343,7 @@ async def test_start_no_config_sends_none(self) -> None: """No config is a valid start; generate_videos gets no knobs.""" client = MagicMock() client.aio.models.generate_videos = AsyncMock(return_value=_sdk_op(name='operations/1', done=False)) - veo = VeoModel('veo-3.0-generate-001', client) + veo = VeoModel('veo-3.1-generate-001', client) await veo.start(_text_request(), ActionRunContext()) @@ -384,7 +366,7 @@ async def test_check_polls_operation_by_sdk_name_and_returns_updated_operation(s ), ), ) - model = VeoModel('veo-3.0-generate-001', client) + model = VeoModel('veo-3.1-generate-001', client) updated = await model.check(Operation(id='operations/123'), ActionRunContext()) assert updated.done is True @@ -395,7 +377,7 @@ async def test_check_wraps_api_error_into_genkit_error(self) -> None: """A 503 on the poll must stay retryable UNAVAILABLE, not collapse to INTERNAL.""" client = MagicMock() client.aio.operations.get = AsyncMock(side_effect=APIError(503, {'error': {'message': 'overloaded'}})) - model = VeoModel('veo-3.0-generate-001', client) + model = VeoModel('veo-3.1-generate-001', client) with pytest.raises(GenkitError) as raised: await model.check(Operation(id='operations/abc'), ActionRunContext()) @@ -403,7 +385,7 @@ async def test_check_wraps_api_error_into_genkit_error(self) -> None: def test_invalid_sdk_field_raises_invalid_argument(self) -> None: """SDK type errors become a named INVALID_ARGUMENT.""" - veo = VeoModel('veo-3.0-generate-001', MagicMock()) + veo = VeoModel('veo-3.1-generate-001', MagicMock()) request = _text_request(config=VeoConfig.model_construct(duration_seconds='nope')) with pytest.raises(GenkitError) as exc_info: @@ -447,7 +429,7 @@ async def test_empty_context_uses_plugin_client(self) -> None: plugin = MagicMock() plugin.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) plugin.aio.operations.get = AsyncMock(return_value=_pending_sdk_op()) - veo = VeoModel('veo-3.0-generate-001', plugin) + veo = VeoModel('veo-3.1-generate-001', plugin) with patch('genkit_google_genai.models.veo.genai.Client') as ctor: started = await veo.start(_text_request(), ActionRunContext()) @@ -463,7 +445,7 @@ async def test_secrets_api_key_builds_request_client(self) -> None: plugin.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) override = MagicMock() override.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op(name='operations/tenant')) - veo = VeoModel('veo-3.0-generate-001', plugin) + veo = VeoModel('veo-3.1-generate-001', plugin) with patch('genkit_google_genai.models.veo.genai.Client', return_value=override) as ctor: op = await veo.start( @@ -486,7 +468,7 @@ async def test_start_config_routes_client_and_stays_out_of_model_parameters(self plugin.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) override = MagicMock() override.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) - veo = VeoModel('veo-3.0-generate-001', plugin, client_kwargs={'api_key': 'plugin-key'}) + veo = VeoModel('veo-3.1-generate-001', plugin, client_kwargs={'api_key': 'plugin-key'}) request = _text_request( config=VeoConfig.model_validate({ 'aspectRatio': '16:9', @@ -517,7 +499,7 @@ async def test_start_config_routes_vertex_location(self) -> None: override = MagicMock() override.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) veo = VeoModel( - 'veo-3.0-generate-001', + 'veo-3.1-generate-001', plugin, client_kwargs={'vertexai': True, 'project': 'p', 'location': 'us-central1'}, ) @@ -536,7 +518,7 @@ async def test_check_without_tenant_context_returns_to_plugin_client(self) -> No plugin.aio.operations.get = AsyncMock(return_value=_pending_sdk_op()) override = MagicMock() override.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) - veo = VeoModel('veo-3.0-generate-001', plugin) + veo = VeoModel('veo-3.1-generate-001', plugin) with patch('genkit_google_genai.models.veo.genai.Client', return_value=override): ticket = await veo.start( @@ -553,7 +535,7 @@ async def test_secrets_apikey_alias(self) -> None: plugin = MagicMock() override = MagicMock() override.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) - veo = VeoModel('veo-3.0-generate-001', plugin) + veo = VeoModel('veo-3.1-generate-001', plugin) with patch('genkit_google_genai.models.veo.genai.Client', return_value=override) as ctor: await veo.start( @@ -570,7 +552,7 @@ async def test_check_uses_secrets_and_does_not_write_key_on_op(self) -> None: plugin.aio.operations.get = AsyncMock(return_value=_pending_sdk_op()) override = MagicMock() override.aio.operations.get = AsyncMock(return_value=_pending_sdk_op()) - veo = VeoModel('veo-3.0-generate-001', plugin) + veo = VeoModel('veo-3.1-generate-001', plugin) ticket = Operation(id='operations/1', done=False) with patch('genkit_google_genai.models.veo.genai.Client', return_value=override) as ctor: @@ -591,7 +573,7 @@ async def test_secrets_and_base_url_together(self) -> None: override = MagicMock() override.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) override.aio.operations.get = AsyncMock(return_value=_pending_sdk_op()) - veo = VeoModel('veo-3.0-generate-001', plugin) + veo = VeoModel('veo-3.1-generate-001', plugin) ctx = ActionRunContext( context={ 'secrets': {'api_key': 'sk-tenant'}, @@ -617,7 +599,7 @@ async def test_vertex_secrets_drop_project_and_location(self) -> None: override = MagicMock() override.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) veo = VeoModel( - 'veo-3.0-generate-001', + 'veo-3.1-generate-001', plugin, client_kwargs={ 'vertexai': True, @@ -645,7 +627,7 @@ async def test_googleai_location_is_ignored(self) -> None: plugin = MagicMock() plugin.vertexai = False plugin.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) - veo = VeoModel('veo-3.0-generate-001', plugin, client_kwargs={'api_key': 'plugin-key'}) + veo = VeoModel('veo-3.1-generate-001', plugin, client_kwargs={'api_key': 'plugin-key'}) with patch('genkit_google_genai.models.veo.genai.Client') as ctor: await veo.start( @@ -663,7 +645,7 @@ async def test_vertex_location_rewrites_base_url(self) -> None: override = MagicMock() override.aio.operations.get = AsyncMock(return_value=_pending_sdk_op()) veo = VeoModel( - 'veo-3.0-generate-001', + 'veo-3.1-generate-001', plugin, client_kwargs={ 'vertexai': True, @@ -690,7 +672,7 @@ async def test_vertex_regional_location_clears_rep_url(self) -> None: override = MagicMock() override.aio.operations.get = AsyncMock(return_value=_pending_sdk_op()) veo = VeoModel( - 'veo-3.0-generate-001', + 'veo-3.1-generate-001', plugin, client_kwargs={ 'vertexai': True, @@ -713,7 +695,7 @@ async def test_vertex_regional_location_clears_rep_url(self) -> None: @pytest.mark.asyncio @pytest.mark.parametrize('bag', ({'api_key': 'sk-wrong'}, {'apiKey': 'sk-wrong'})) async def test_config_api_key_is_invalid_argument(self, bag: dict[str, str]) -> None: - veo = VeoModel('veo-3.0-generate-001', MagicMock()) + veo = VeoModel('veo-3.1-generate-001', MagicMock()) with pytest.raises(GenkitError) as raised: await veo.start( @@ -726,7 +708,7 @@ async def test_config_api_key_is_invalid_argument(self, bag: dict[str, str]) -> @pytest.mark.asyncio async def test_secrets_must_be_a_dict(self) -> None: - veo = VeoModel('veo-3.0-generate-001', MagicMock()) + veo = VeoModel('veo-3.1-generate-001', MagicMock()) with pytest.raises(GenkitError) as raised: await veo.start( @@ -738,7 +720,7 @@ async def test_secrets_must_be_a_dict(self) -> None: @pytest.mark.asyncio async def test_secret_api_key_must_be_a_string(self) -> None: - veo = VeoModel('veo-3.0-generate-001', MagicMock()) + veo = VeoModel('veo-3.1-generate-001', MagicMock()) with pytest.raises(GenkitError) as raised: await veo.start( @@ -755,7 +737,7 @@ async def test_api_version_overlay(self) -> None: plugin.vertexai = False override = MagicMock() override.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) - veo = VeoModel('veo-3.0-generate-001', plugin, client_kwargs={'api_key': 'plugin-key'}) + veo = VeoModel('veo-3.1-generate-001', plugin, client_kwargs={'api_key': 'plugin-key'}) with patch('genkit_google_genai.models.veo.genai.Client', return_value=override) as ctor: await veo.start( @@ -769,7 +751,7 @@ async def test_api_version_overlay(self) -> None: @pytest.mark.asyncio async def test_empty_secrets_pocket_is_invalid_argument(self) -> None: - veo = VeoModel('veo-3.0-generate-001', MagicMock()) + veo = VeoModel('veo-3.1-generate-001', MagicMock()) for pocket in ({}, {'api_key': None}, {'api_key': ''}): with pytest.raises(GenkitError) as raised: await veo.start( @@ -780,7 +762,7 @@ async def test_empty_secrets_pocket_is_invalid_argument(self) -> None: @pytest.mark.asyncio async def test_top_level_api_key_is_invalid_argument(self) -> None: - veo = VeoModel('veo-3.0-generate-001', MagicMock()) + veo = VeoModel('veo-3.1-generate-001', MagicMock()) for bag in ({'api_key': 'sk-wrong'}, {'apiKey': 'sk-wrong'}): with pytest.raises(GenkitError) as raised: await veo.start(_text_request(), ActionRunContext(context=bag)) @@ -789,7 +771,7 @@ async def test_top_level_api_key_is_invalid_argument(self) -> None: @pytest.mark.asyncio async def test_request_config_api_key_is_invalid_argument(self) -> None: - veo = VeoModel('veo-3.0-generate-001', MagicMock()) + veo = VeoModel('veo-3.1-generate-001', MagicMock()) cfg = VeoConfig.model_validate({'api_key': 'sk-gemini-habit'}) with pytest.raises(GenkitError) as raised: @@ -802,7 +784,7 @@ async def test_request_config_api_key_is_invalid_argument(self) -> None: async def test_client_ctor_failure_is_invalid_argument(self) -> None: plugin = MagicMock() plugin.vertexai = False - veo = VeoModel('veo-3.0-generate-001', plugin, client_kwargs={'api_key': 'plugin-key'}) + veo = VeoModel('veo-3.1-generate-001', plugin, client_kwargs={'api_key': 'plugin-key'}) with ( patch( @@ -826,7 +808,7 @@ async def test_vertex_tenant_key_clears_plugin_base_url(self) -> None: override = MagicMock() override.aio.models.generate_videos = AsyncMock(return_value=_pending_sdk_op()) veo = VeoModel( - 'veo-3.0-generate-001', + 'veo-3.1-generate-001', plugin, client_kwargs={ 'vertexai': True, diff --git a/py/packages/genkit/tests/genkit/core/registry_test.py b/py/packages/genkit/tests/genkit/core/registry_test.py index 7118122f63..9a83b9324c 100644 --- a/py/packages/genkit/tests/genkit/core/registry_test.py +++ b/py/packages/genkit/tests/genkit/core/registry_test.py @@ -538,8 +538,8 @@ async def resolve(self, action_type: ActionKind, name: str) -> Action | None: return Action(name=name, kind=ActionKind.BACKGROUND_MODEL, fn=_bg_start) ai = Genkit(plugins=[VeoPlugin()]) - got = await ai.registry.resolve_model('plug/veo-2.0-generate-001') + got = await ai.registry.resolve_model('plug/veo-3.1-generate-preview') assert got is not None assert got.kind == ActionKind.BACKGROUND_MODEL - assert got.name == 'plug/veo-2.0-generate-001' + assert got.name == 'plug/veo-3.1-generate-preview'