From 4fb413f138a60ff947a60cdd4713502ee21669b4 Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Thu, 30 Jul 2026 11:03:50 +0100 Subject: [PATCH] refactor(py/plugins/google-genai): move Veo to background model actions Veo slice extracted from #5806: rewrites veo.py around create_veo_background_action on the shared background-model plumbing, wires Veo through BACKGROUND_MODEL/CHECK_OPERATION resolution in both GoogleAI and VertexAI (dropping the plain MODEL registration), and updates the google-genai-media sample to poll generate_operation directly. Co-authored-by: Jeff Huang --- .../src/genkit_google_genai/google.py | 123 +++--- .../src/genkit_google_genai/models/veo.py | 389 +++++------------- .../test/google_plugin_test.py | 19 +- .../genkit-google-genai/tests/veo_test.py | 268 +++++++----- py/samples/google-genai-media/README.md | 2 + py/samples/google-genai-media/src/main.py | 101 +---- 6 files changed, 338 insertions(+), 564 deletions(-) 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 1d8dd4cc17..5903a15d57 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 @@ -97,8 +97,8 @@ vertexai_image_model_info, ) from genkit_google_genai.models.veo import ( - VeoConfigSchema, - VeoModel, + VeoConfig, + create_veo_background_action, is_veo_model, veo_model_info, ) @@ -284,7 +284,7 @@ def _create_embedder_action( ), ) - async def _run(request: Any) -> Any: # noqa: ANN401 + async def run(request: Any) -> Any: # noqa: ANN401 embedder = Embedder( version=clean_name, client=client_getter(), @@ -295,7 +295,7 @@ async def _run(request: Any) -> Any: # noqa: ANN401 action = Action( kind=ActionKind.EMBEDDER, name=full_name, - fn=_run, + fn=run, metadata=action_metadata.metadata, ) @@ -338,6 +338,16 @@ class GoogleAI(Plugin): # 3. Inspect output text directly print(res.text) # => Quantum computing utilizes quantum bits to solve complex problems faster... + + # 4. Video generation runs as a background model + operation = await ai.generate_operation( + model='googleai/veo-2.0-generate-001', + prompt='A sunset over mountains', + ) + while not operation.done: + await asyncio.sleep(5) + operation = await ai.check_operation(operation) + print(operation.output.message.content[0].root.media.url) ``` Attributes: @@ -483,9 +493,13 @@ async def resolve(self, action_type: ActionKind, name: str) -> Action | None: Action object if found, None otherwise. """ if action_type == ActionKind.MODEL: + prefix = GOOGLEAI_PLUGIN_NAME + '/' + clean_name = name.replace(prefix, '') if name.startswith(prefix) else name + # Background-only families: leave MODEL empty so resolve_model falls back. + if is_veo_model(clean_name): + return None return self._resolve_model(name) elif action_type == ActionKind.BACKGROUND_MODEL: - # For Veo models, return the start action prefix = GOOGLEAI_PLUGIN_NAME + '/' clean_name = name.replace(prefix, '') if name.startswith(prefix) else name if is_veo_model(clean_name): @@ -493,10 +507,8 @@ async def resolve(self, action_type: ActionKind, name: str) -> Action | None: return bg_action.start_action return None elif action_type == ActionKind.CHECK_OPERATION: - # Check action names are in format {model_name}/check - # Extract the model name and resolve if it's a Veo model if name.endswith('/check'): - model_name = name[:-6] # Remove '/check' suffix + model_name = name[:-6] prefix = GOOGLEAI_PLUGIN_NAME + '/' clean_name = model_name.replace(prefix, '') if model_name.startswith(prefix) else model_name if is_veo_model(clean_name): @@ -508,52 +520,8 @@ async def resolve(self, action_type: ActionKind, name: str) -> Action | None: return None def _resolve_veo_model(self, name: str) -> BackgroundAction: - """Create a BackgroundAction for a Veo video generation model. - - Args: - name: The namespaced name of the model. - - Returns: - BackgroundAction for the Veo model. - """ - clean_name = name.replace(GOOGLEAI_PLUGIN_NAME + '/', '') if name.startswith(GOOGLEAI_PLUGIN_NAME) else name - - # Create actions manually since we don't have registry access here - - async def _start(request: Any, ctx: Any) -> Any: # noqa: ANN401 - veo = VeoModel(clean_name, self._runtime_client()) - return await veo.start(request, ctx) - - async def _check(op: Any, _ctx: Any) -> Any: # noqa: ANN401 - veo = VeoModel(clean_name, self._runtime_client()) - return await veo.check(op) - - # Prepare metadata matching model_action_metadata structure - info = veo_model_info(clean_name).model_dump(by_alias=True) - config_schema = VeoConfigSchema - - start_action = Action( - kind=ActionKind.BACKGROUND_MODEL, - name=name, - fn=_start, - metadata={ - 'model': {**info, 'customOptions': to_json_schema(config_schema)}, - 'type': 'background-model', - }, - ) - - check_action = Action( - kind=ActionKind.CHECK_OPERATION, - name=f'{name}/check', - fn=_check, - metadata={'type': 'check-operation'}, - ) - - return BackgroundAction( - start_action=start_action, - check_action=check_action, - cancel_action=None, - ) + """Create a BackgroundAction for a Veo video generation model.""" + return create_veo_background_action(name, self._runtime_client()) def _resolve_model(self, name: str) -> Action: """Create an Action object for a Google AI model. @@ -577,7 +545,7 @@ def _resolve_model(self, name: str) -> Action: SUPPORTED_MODELS[clean_name] = model_ref config_schema = get_model_config_schema(clean_name) - async def _run(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: + async def run(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: if clean_name.lower().startswith('image'): model = ImagenModel(clean_name, self._runtime_client()) else: @@ -592,7 +560,7 @@ async def _run(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: return Action( kind=ActionKind.MODEL, name=name, - fn=_run, + fn=run, metadata=model_action_metadata( name=name, info=model_ref.model_dump(by_alias=True), @@ -649,7 +617,7 @@ async def list_actions(self) -> list[ActionMetadata]: model_action_metadata( name=googleai_name(name), info=veo_model_info(name).model_dump(by_alias=True), - config_schema=VeoConfigSchema, + config_schema=VeoConfig, ) ) @@ -688,7 +656,7 @@ class VertexAI(Plugin): |---|---|---| | Gemini / Gemma | MODEL | ``vertexai/gemini-flash-latest`` | | Imagen | MODEL | ``vertexai/imagen-3.0-generate-002`` | - | Veo (Video) | MODEL | ``vertexai/veo-2.0-generate-001`` | + | Veo (Video) | BACKGROUND_MODEL | ``vertexai/veo-2.0-generate-001`` | | Embedders | EMBEDDER | ``vertexai/text-embedding-005`` | Example: @@ -820,7 +788,9 @@ async def init(self) -> list[Action]: actions.append(self._resolve_model(vertexai_name(name))) for name in genai_models.veo: - actions.append(self._resolve_model(vertexai_name(name))) + bg_action = self._resolve_veo_model(vertexai_name(name)) + actions.append(bg_action.start_action) + actions.append(bg_action.check_action) for name in VERTEX_KNOWN_EMBEDDERS: actions.append(self._resolve_embedder(vertexai_name(name))) @@ -854,10 +824,12 @@ def _list_known_models(self) -> list[Action]: actions.append(self._resolve_model(vertexai_name(name))) for name in genai_models.imagen: actions.append(self._resolve_model(vertexai_name(name))) - for name in genai_models.veo: - actions.append(self._resolve_model(vertexai_name(name))) return actions + def _resolve_veo_model(self, name: str) -> BackgroundAction: + """Create a BackgroundAction for a Veo video generation model.""" + return create_veo_background_action(name, self._runtime_client()) + def _list_known_embedders(self) -> list[Action]: """List known embedders as Action objects. @@ -881,7 +853,25 @@ async def resolve(self, action_type: ActionKind, name: str) -> Action | None: Action object if found, None otherwise. """ if action_type == ActionKind.MODEL: + prefix = VERTEXAI_PLUGIN_NAME + '/' + clean_name = name.replace(prefix, '') if name.startswith(prefix) else name + if is_veo_model(clean_name): + return None return self._resolve_model(name) + elif action_type == ActionKind.BACKGROUND_MODEL: + prefix = VERTEXAI_PLUGIN_NAME + '/' + clean_name = name.replace(prefix, '') if name.startswith(prefix) else name + if is_veo_model(clean_name): + return self._resolve_veo_model(name).start_action + return None + elif action_type == ActionKind.CHECK_OPERATION: + if name.endswith('/check'): + model_name = name[:-6] + prefix = VERTEXAI_PLUGIN_NAME + '/' + clean_name = model_name.replace(prefix, '') if model_name.startswith(prefix) else model_name + if is_veo_model(clean_name): + return self._resolve_veo_model(model_name).check_action + return None elif action_type == ActionKind.EMBEDDER: return self._resolve_embedder(name) elif action_type == ActionKind.EVALUATOR: @@ -947,15 +937,12 @@ def _resolve_model(self, name: str) -> Action: model_ref = vertexai_image_model_info(clean_name) IMAGE_SUPPORTED_MODELS[clean_name] = model_ref # pyright: ignore[reportArgumentType] config_schema = ImagenConfigSchema - elif is_veo_model(clean_name): - model_ref = veo_model_info(clean_name) - config_schema = VeoConfigSchema else: model_ref = google_model_info(clean_name) SUPPORTED_MODELS[clean_name] = model_ref config_schema = get_model_config_schema(clean_name) - async def _run(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: + async def run(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: if is_tuned_gemini_name(clean_name): model = GeminiModel( clean_name, @@ -965,8 +952,6 @@ async def _run(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: ) elif clean_name.lower().startswith('image'): model = ImagenModel(clean_name, self._runtime_client()) - elif is_veo_model(clean_name): - model = VeoModel(clean_name, self._runtime_client()) else: model = GeminiModel( clean_name, @@ -979,7 +964,7 @@ async def _run(request: ModelRequest, ctx: ActionRunContext) -> ModelResponse: return Action( kind=ActionKind.MODEL, name=name, - fn=_run, + fn=run, metadata=model_action_metadata( name=name, info=model_ref.model_dump(by_alias=True), @@ -1036,7 +1021,7 @@ async def list_actions(self) -> list[ActionMetadata]: model_action_metadata( name=vertexai_name(name), info=veo_model_info(name).model_dump(by_alias=True), - config_schema=VeoConfigSchema, + config_schema=VeoConfig, ) ) 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 660d8e3c1e..0167a84981 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 @@ -14,25 +14,16 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Veo video generation model for Google GenAI plugin. +"""Veo video generation via the GenAI generate_videos long-running API.""" -Veo is Google's video generation model that creates videos from text prompts. -""" - -import asyncio -import sys -from typing import Any, cast - -if sys.version_info < (3, 11): - from strenum import StrEnum -else: - from enum import StrEnum +from __future__ import annotations from google import genai from google.genai import types as genai_types from pydantic import BaseModel, ConfigDict, Field from genkit import ( + FinishReason, Media, MediaPart, Message, @@ -44,16 +35,16 @@ Supports, TextPart, ) -from genkit.model import Error, Operation -from genkit.plugin_api import ActionRunContext, tracer +from genkit._core._background import define_background_model +from genkit._core._compat import StrEnum +from genkit._core._registry import Registry +from genkit.model import BackgroundAction, Error, Operation +from genkit.plugin_api import ActionRunContext +from genkit_google_genai.models.interactions_utils import extract_version class VeoVersion(StrEnum): - """Supported Veo video generation models. - - Note: Models are discovered dynamically. This enum provides convenience - constants for commonly used Veo models. - """ + """Commonly used Veo model version identifiers.""" VEO_2_0 = 'veo-2.0-generate-001' VEO_2_0_EXP = 'veo-2.0-generate-exp' @@ -65,39 +56,17 @@ class VeoVersion(StrEnum): VEO_3_1_FAST = 'veo-3.1-fast-generate-001' -def is_veo_model(name: str) -> bool: - """Check if a model name is a Veo model. - - Args: - name: The model name to check. - - Returns: - True if this is a Veo model name. - """ - return name.lower().startswith('veo') - - -class VeoConfigSchema(BaseModel): - """Veo Config Schema.""" +class VeoConfig(BaseModel): + """Veo video generation configuration.""" model_config = ConfigDict(extra='allow', populate_by_name=True) - negative_prompt: str | None = Field( - default=None, alias='negativePrompt', description='Negative prompt for video generation.' - ) - aspect_ratio: str | None = Field( - default=None, alias='aspectRatio', description='Desired aspect ratio of the output video (e.g. "16:9").' - ) - person_generation: str | None = Field(default=None, alias='personGeneration', description='Person generation mode.') - duration_seconds: int | None = Field( - default=None, alias='durationSeconds', description='Length of video in seconds.' - ) - resolution: str | None = Field(default=None, description='Desired output resolution (e.g. "720p").') - seed: int | None = Field(default=None, description='Random seed for deterministic generation.') - enhance_prompt: bool | None = Field(default=None, alias='enhancePrompt', description='Enable prompt enhancement.') - - -# Alias for backwards compatibility with __init__.py exports -VeoConfig = VeoConfigSchema + negative_prompt: str | None = Field(default=None, alias='negativePrompt') + aspect_ratio: str | None = Field(default=None, alias='aspectRatio') + person_generation: str | None = Field(default=None, alias='personGeneration') + duration_seconds: int | None = Field(default=None, alias='durationSeconds') + resolution: str | None = None + seed: int | None = None + enhance_prompt: bool | None = Field(default=None, alias='enhancePrompt') DEFAULT_VEO_SUPPORT = Supports( @@ -106,280 +75,118 @@ class VeoConfigSchema(BaseModel): tools=False, system_role=True, output=['media'], + long_running=True, ) -def veo_model_info(version: str) -> ModelInfo: - """Get model info for a Veo model. +def is_veo_model(name: str) -> bool: + """Return True when the model name belongs to the Veo family.""" + return name.lower().startswith('veo') - Args: - version: The Veo model version. - Returns: - ModelInfo describing the model's capabilities. - """ +def veo_model_info(version: str) -> ModelInfo: + """Return capability metadata for a Veo model.""" + clean = extract_version(version) return ModelInfo( - label=f'Google AI - {version}', + label=f'Google AI - {clean}', supports=DEFAULT_VEO_SUPPORT, ) -def _extract_text(request: ModelRequest) -> str: - """Extract text prompt from a ModelRequest. +def to_veo_parameters(config: VeoConfig) -> genai_types.GenerateVideosConfig: + """Convert VeoConfig into the SDK GenerateVideosConfig.""" + return genai_types.GenerateVideosConfig.model_validate(config.model_dump(exclude_none=True)) - Args: - request: The generation request. - Returns: - The text prompt string. - """ - prompt_parts = [ - str(part.root.text) +def extract_text_prompt(request: ModelRequest[VeoConfig]) -> str: + """Join text parts from the request into a single prompt string.""" + parts = [ + part.root.text for message in request.messages or [] for part in message.content - if hasattr(part.root, 'text') and part.root.text + if isinstance(part.root, TextPart) and part.root.text ] - return ' '.join(prompt_parts) - - -def _to_veo_parameters(config: Any) -> dict[str, Any]: # noqa: ANN401 - """Convert config to Veo API parameters. - - Args: - config: The model configuration (VeoConfigSchema or dict). + return ' '.join(parts) - Returns: - Dictionary of Veo API parameters. - """ - if config is None: - return {} - if isinstance(config, VeoConfigSchema): - params = config.model_dump(by_alias=True, exclude_none=True) - elif isinstance(config, dict): - params = {k: v for k, v in config.items() if v is not None} - else: - return {} +def video_parts_from_uris(uris: list[str]) -> list[Part]: + """Build model message parts for generated video URIs.""" + return [Part(root=MediaPart(media=Media(url=uri, content_type='video/mp4'))) for uri in uris] - return params +def extract_video_uris(response: genai_types.GenerateVideosResponse) -> list[str]: + """Extract video URIs from a GenerateVideosResponse.""" + uris: list[str] = [] + for item in response.generated_videos or []: + if item.video and item.video.uri: + uris.append(item.video.uri) + return uris + + +def model_response_from_veo( + response: genai_types.GenerateVideosResponse, +) -> ModelResponse[genai_types.GenerateVideosResponse]: + """Build a ModelResponse from a completed GenerateVideosResponse.""" + return ModelResponse[genai_types.GenerateVideosResponse]( + finish_reason=FinishReason.STOP, + message=Message( + role=Role.MODEL, + content=video_parts_from_uris(extract_video_uris(response)), + ), + raw=response.model_dump(exclude_none=True), + ) -def _from_veo_operation(api_op: dict[str, Any]) -> Operation: - """Convert Veo API operation to Genkit Operation. - - The ``response`` value in ``api_op`` may be either: - - * A plain dict (from the ``start`` method, or legacy REST responses). - * A ``GenerateVideosResponse`` Pydantic model (from the ``check`` method, - which stores the SDK object directly). - - This function handles both cases when extracting video URIs. - - Args: - api_op: The raw API operation response dict. - Returns: - A Genkit Operation object. - """ +def from_veo_operation(operation: genai_types.GenerateVideosOperation) -> Operation: + """Convert a GenerateVideosOperation into a Genkit Operation.""" + # LRO can omit or null `done` while still running — treat that as pending. op = Operation( - id=api_op.get('name', ''), - done=api_op.get('done', False), + id=operation.name or '', + done=bool(operation.done), ) - - # Handle error - if api_op.get('error'): - op.error = Error(message=api_op['error'].get('message', 'Unknown error')) - return op - - # Handle response with generated videos. - response = api_op.get('response') - if response is None: + if operation.error: + op.error = Error(message=str(operation.error.get('message', 'Unknown error'))) return op - # Extract video URIs — response may be a Pydantic model or a dict. - uris: list[str] = [] - if hasattr(response, 'generated_videos'): - # Pydantic GenerateVideosResponse from the SDK (check path). - for gv in response.generated_videos or []: - if gv.video and gv.video.uri: - uris.append(gv.video.uri) - elif isinstance(response, dict): - # Plain dict (start path or legacy REST). - video_response = response.get('generateVideoResponse', {}) - for sample in video_response.get('generatedSamples', []): - video = sample.get('video', {}) - uri = video.get('uri') - if uri: - uris.append(uri) - - if uris: - content = [{'media': {'url': uri}} for uri in uris] - op.output = { - 'finishReason': 'stop', - 'message': { - 'role': 'model', - 'content': content, - }, - } - + response = operation.response or operation.result + if response is not None and extract_video_uris(response): + output: ModelResponse[genai_types.GenerateVideosResponse] = model_response_from_veo(response) + op.output = output return op -class VeoModel: - """Veo video generation model. - - This class implements both the standard model interface (for Vertex AI) - and the background model pattern (for GoogleAI) for Veo video generation. - """ - - def __init__(self, version: str, client: genai.Client) -> None: - """Initialize Veo model. - - Args: - version: The Veo model version. - client: The Google GenAI client. - """ - self._version = version - self._client = client - - def _build_prompt(self, request: ModelRequest) -> str: - """Build prompt request from Genkit request.""" - prompt = [] - for message in request.messages: - for part in message.content: - if isinstance(part.root, TextPart): - prompt.append(part.root.text) - else: - # TODO(#4363): Support image input if Veo supports it (e.g. for image-to-video) - # For now, strict text text-to-video - pass - return ' '.join(prompt) - - async def generate(self, request: ModelRequest, _: ActionRunContext) -> ModelResponse: - """Handle a generation request (synchronous/blocking mode for Vertex AI). - - Args: - request: The generation request. - _: action context - - Returns: - The model's response. - """ - if request.tools: - raise ValueError('Tools are not supported for this model.') - - prompt = self._build_prompt(request) - config = self._get_config(request) - - with tracer.start_as_current_span('generate_videos'): - operation = await self._client.aio.models.generate_videos(model=self._version, prompt=prompt, config=config) - - # Handling LRO. Using cast(Any) to avoid strict type definition issues for operation.result() - op = cast(Any, operation) - if hasattr(op, 'result'): - # Check if result is a coroutine (awaitable) or direct value - res = op.result() - if asyncio.iscoroutine(res): - response = await res - else: - response = res - else: - response = op - - content = self._contents_from_response(cast(genai_types.GenerateVideosResponse, response)) - - return ModelResponse( - message=Message( - content=content, - role=Role.MODEL, - ) - ) - - async def start(self, request: ModelRequest, ctx: ActionRunContext) -> Operation: - """Start a video generation operation (background model pattern for GoogleAI). +def create_veo_background_action(name: str, client: genai.Client) -> BackgroundAction: + """Build a Veo background model: start returns an Operation; check refreshes it once.""" + version = extract_version(name) + info = veo_model_info(version) - Args: - request: The generation request. - ctx: The action run context. - - Returns: - An Operation with the job ID. - """ + async def start(request: ModelRequest[VeoConfig], _: ActionRunContext) -> Operation: if request.tools: raise ValueError('Tools are not supported for this model.') - - prompt = _extract_text(request) + prompt = extract_text_prompt(request) if not prompt: raise ValueError('Veo requires a text prompt') - - # Call the generateVideos API - response = await self._client.aio.models.generate_videos( - model=self._version, + config = request.config or VeoConfig() + sdk_op = await client.aio.models.generate_videos( + model=version, prompt=prompt, - # pyrefly: ignore[bad-argument-type] - config dict matches GenerateVideosConfigDict - config=request.config if isinstance(request.config, dict) else None, # pyright: ignore[reportArgumentType] + config=to_veo_parameters(config), ) + return from_veo_operation(sdk_op) - # Convert to Operation - return _from_veo_operation({ - 'name': response.name if hasattr(response, 'name') else str(response), - 'done': getattr(response, 'done', False), - }) - - async def check(self, operation: Operation) -> Operation: - """Check the status of a video generation operation. - - Args: - operation: The operation to check. - - Returns: - Updated Operation with current status. - """ - # Get the operation status using the public operations.get() API - # See: https://ai.google.dev/gemini-api/docs/video - # Create a GenerateVideosOperation object from the operation ID - op_request = genai_types.GenerateVideosOperation.model_validate({'name': operation.id}) - response = await self._client.aio.operations.get(operation=op_request) - - # Convert response to dict for processing - op_dict = { - 'name': getattr(response, 'name', operation.id), - 'done': getattr(response, 'done', False), - } - - if hasattr(response, 'error') and response.error: - op_dict['error'] = {'message': str(response.error)} - - if hasattr(response, 'response') and response.response: - op_dict['response'] = response.response - - return _from_veo_operation(op_dict) - - def _get_config(self, request: ModelRequest) -> genai_types.GenerateVideosConfigOrDict | None: - if not request.config: - return None - return cast(genai_types.GenerateVideosConfigOrDict, request.config) - - def _contents_from_response(self, response: genai_types.GenerateVideosResponse) -> list[Part]: - content = [] - if response.generated_videos: - for video in response.generated_videos: - # Video URI is typically in video.video.uri - if video.video and video.video.uri: - uri = video.video.uri - content.append( - Part( - root=MediaPart( - media=Media( - url=uri, - content_type='video/mp4', - ) - ) - ) - ) - return content - - @property - def metadata(self) -> dict: - """Model metadata.""" - return {'model': {'supports': DEFAULT_VEO_SUPPORT.model_dump(by_alias=True)}} + async def check(operation: Operation) -> Operation: + sdk_op = await client.aio.operations.get( + operation=genai_types.GenerateVideosOperation.model_validate({'name': operation.id}), + ) + return from_veo_operation(sdk_op) + + return define_background_model( + registry=Registry(), + name=name, + start=start, + check=check, + cancel=None, + label=info.label or name, + info=info, + config_schema=VeoConfig, + ) 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 f4a1e8e1d6..4891494d9c 100644 --- a/py/packages/genkit-google-genai/test/google_plugin_test.py +++ b/py/packages/genkit-google-genai/test/google_plugin_test.py @@ -876,8 +876,6 @@ class MockModel: # Verify Veo action4 = next(a for a in result if a.name == vertexai_name('veo-2.0-generate-001')) assert action4 is not None - # from genkit_google_genai.models.veo import VeoConfigSchema - # assert action4.config_schema == VeoConfigSchema @pytest.mark.asyncio @@ -988,18 +986,13 @@ class MockModel: vertexai_plugin_instance._runtime_client = lambda: mock_client result = vertexai_plugin_instance._list_known_models() + names = {a.name for a in result} - # Verify Gemini - 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 - - # Verify Veo - action4 = next(a for a in result if a.name == vertexai_name('veo-2.0-generate-001')) - assert action4 is not None + assert vertexai_name('gemini-1.5-flash') in names + assert vertexai_name('imagen-3.0-generate-001') in names + # Veo only produces video through a long-running operation, so it is offered + # as a background model and never as a plain MODEL action. + assert vertexai_name('veo-2.0-generate-001') not in names @pytest.mark.asyncio diff --git a/py/packages/genkit-google-genai/tests/veo_test.py b/py/packages/genkit-google-genai/tests/veo_test.py index f6729dbf39..2e9babfa34 100644 --- a/py/packages/genkit-google-genai/tests/veo_test.py +++ b/py/packages/genkit-google-genai/tests/veo_test.py @@ -14,23 +14,30 @@ # # SPDX-License-Identifier: Apache-2.0 -"""Tests for Veo video generation model helpers. +"""Tests for Veo video generation model helpers and background actions.""" -Verifies _from_veo_operation handles both dict-based responses (from the -start path) and Pydantic GenerateVideosResponse objects (from the check -path where the SDK returns a model instance). -""" +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from genkit_google_genai import GoogleAI +from genkit_google_genai.google import GenaiModels from genkit_google_genai.models.veo import ( - VeoConfigSchema, + VeoConfig, VeoVersion, - _from_veo_operation, - _to_veo_parameters, + from_veo_operation, is_veo_model, + to_veo_parameters, ) from google.genai import types as genai_types +from genkit import FinishReason, Genkit, ModelResponse + + +def veo_operation(**fields: Any) -> genai_types.GenerateVideosOperation: + """Build an SDK video operation the way the API delivers one.""" + return genai_types.GenerateVideosOperation.model_validate(fields) + class TestIsVeoModel: """Tests for is_veo_model.""" @@ -43,9 +50,10 @@ def test_veo_uppercase(self) -> None: """Case-insensitive matching works.""" assert is_veo_model('VEO-2.0-generate-001') is True - def test_non_veo_model(self) -> None: - """Non-Veo model names are rejected.""" + def test_non_veo_model_name(self) -> None: + """Non-Veo models return False.""" assert is_veo_model('gemini-2.0-flash') is False + assert is_veo_model('imagen-3.0-generate-001') is False class TestVeoVersion: @@ -66,98 +74,67 @@ def test_new_googleai_models_are_recognized(self, version: VeoVersion) -> None: class TestToVeoParameters: - """Tests for _to_veo_parameters.""" + """Tests for to_veo_parameters.""" - def test_none_config(self) -> None: - """None config returns empty dict.""" - assert _to_veo_parameters(None) == {} - - def test_dict_config(self) -> None: - """Dict config filters out None values.""" - config = {'aspect_ratio': '16:9', 'duration_seconds': 5, 'empty': None} - result = _to_veo_parameters(config) - assert result == {'aspect_ratio': '16:9', 'duration_seconds': 5} + def test_empty_config(self) -> None: + """Empty VeoConfig becomes an empty GenerateVideosConfig.""" + result = to_veo_parameters(VeoConfig()) + assert isinstance(result, genai_types.GenerateVideosConfig) + assert result.aspect_ratio is None def test_schema_config(self) -> None: - """VeoConfigSchema is converted with camelCase keys.""" - config = VeoConfigSchema(aspect_ratio='16:9', duration_seconds=5) - result = _to_veo_parameters(config) - assert result['aspectRatio'] == '16:9' - assert result['durationSeconds'] == 5 + """VeoConfig maps onto GenerateVideosConfig fields.""" + config = VeoConfig(aspect_ratio='16:9', duration_seconds=5) + result = to_veo_parameters(config) + assert result.aspect_ratio == '16:9' + assert result.duration_seconds == 5 def test_schema_config_includes_new_fields(self) -> None: - """VeoConfigSchema includes newer Veo parameters.""" - config = VeoConfigSchema(resolution='1080p', seed=7) - result = _to_veo_parameters(config) - assert result['resolution'] == '1080p' - assert result['seed'] == 7 + """VeoConfig includes newer Veo parameters.""" + config = VeoConfig(resolution='1080p', seed=7) + result = to_veo_parameters(config) + assert result.resolution == '1080p' + assert result.seed == 7 class TestFromVeoOperation: - """Tests for _from_veo_operation. - - This function must handle two shapes for the 'response' value: - - 1. A plain dict — returned by the start() path or legacy REST. - 2. A GenerateVideosResponse Pydantic model — returned by the check() - path where the SDK object is stored directly. - - Regression: before the fix, case 2 raised - ``AttributeError: 'GenerateVideosResponse' object has no attribute 'get'`` - because the code unconditionally called ``.get()`` on the response. - """ + """Tests for from_veo_operation with typed GenerateVideosOperation.""" def test_pending_operation(self) -> None: """An in-progress operation has no response — output stays None.""" - op = _from_veo_operation({ - 'name': 'operations/123', - 'done': False, - }) + op = from_veo_operation(veo_operation(name='operations/123', done=False)) assert op.id == 'operations/123' assert op.done is False assert op.output is None assert op.error is None + def test_pending_operation_null_done_normalized(self) -> None: + """API null/omitted done is pending, not a missing flag.""" + for sdk_op in ( + veo_operation(name='operations/null-done', done=None), + veo_operation(name='operations/omitted-done'), + ): + op = from_veo_operation(sdk_op) + assert op.done is False + assert op.output is None + def test_error_operation(self) -> None: """An operation with an error populates op.error.""" - op = _from_veo_operation({ - 'name': 'operations/456', - 'done': True, - 'error': {'message': 'Quota exceeded'}, - }) + op = from_veo_operation( + veo_operation( + name='operations/456', + done=True, + error={'message': 'Quota exceeded'}, + ) + ) assert op.id == 'operations/456' assert op.done is True assert op.error is not None assert op.error.message == 'Quota exceeded' assert op.output is None - def test_dict_response_with_videos(self) -> None: - """Dict-shaped response extracts video URIs (start path).""" - op = _from_veo_operation({ - 'name': 'operations/789', - 'done': True, - 'response': { - 'generateVideoResponse': { - 'generatedSamples': [ - {'video': {'uri': 'https://example.com/v1.mp4'}}, - {'video': {'uri': 'https://example.com/v2.mp4'}}, - ] - } - }, - }) - assert op.done is True - assert op.output is not None - assert op.output['finishReason'] == 'stop' - content = op.output['message']['content'] - assert len(content) == 2 - assert content[0]['media']['url'] == 'https://example.com/v1.mp4' - assert content[1]['media']['url'] == 'https://example.com/v2.mp4' - def test_pydantic_response_with_videos(self) -> None: - """Pydantic GenerateVideosResponse extracts video URIs (check path). - - This is the regression case — previously this raised AttributeError. - """ + """GenerateVideosResponse extracts video URIs (check path).""" pydantic_response = genai_types.GenerateVideosResponse( generated_videos=[ genai_types.GeneratedVideo( @@ -172,47 +149,120 @@ def test_pydantic_response_with_videos(self) -> None: ), ], ) - op = _from_veo_operation({ - 'name': 'models/veo-2.0-generate-001/operations/abc', - 'done': True, - 'response': pydantic_response, - }) + op = from_veo_operation( + veo_operation( + name='models/veo-2.0-generate-001/operations/abc', + done=True, + response=pydantic_response, + ) + ) assert op.done is True - assert op.output is not None - assert op.output['finishReason'] == 'stop' - content = op.output['message']['content'] + assert isinstance(op.output, ModelResponse) + assert op.output.finish_reason == FinishReason.STOP + content = op.output.message.content if op.output.message else [] assert len(content) == 2 - assert content[0]['media']['url'] == 'https://example.com/video_a.mp4' - assert content[1]['media']['url'] == 'https://example.com/video_b.mp4' + media0 = content[0].root.media + media1 = content[1].root.media + assert media0 is not None and media0.url == 'https://example.com/video_a.mp4' + assert media1 is not None and media1.url == 'https://example.com/video_b.mp4' def test_pydantic_response_empty_videos(self) -> None: - """Pydantic response with no generated_videos produces no output.""" - pydantic_response = genai_types.GenerateVideosResponse( - generated_videos=[], + """Response with no generated_videos produces no output.""" + op = from_veo_operation( + veo_operation( + name='operations/empty', + done=True, + response=genai_types.GenerateVideosResponse(generated_videos=[]), + ) ) - op = _from_veo_operation({ - 'name': 'operations/empty', - 'done': True, - 'response': pydantic_response, - }) assert op.done is True assert op.output is None def test_response_none_explicit(self) -> None: """Explicit None response is handled (no crash).""" - op = _from_veo_operation({ - 'name': 'operations/null', - 'done': False, - 'response': None, - }) + op = from_veo_operation( + veo_operation( + name='operations/null', + done=False, + response=None, + ) + ) assert op.output is None - def test_dict_response_no_videos(self) -> None: - """Dict response with empty generatedSamples produces no output.""" - op = _from_veo_operation({ - 'name': 'operations/empty-dict', - 'done': True, - 'response': {'generateVideoResponse': {'generatedSamples': []}}, - }) - assert op.done is True - assert op.output is None + +def mock_veo_client(start_done: bool = False) -> MagicMock: + """Build a mocked GenAI client for Veo background-model tests.""" + client = MagicMock() + start_op = veo_operation( + name='operations/veo-start', + done=start_done, + ) + completed_response = genai_types.GenerateVideosResponse( + generated_videos=[ + genai_types.GeneratedVideo( + video=genai_types.Video(uri='https://example.com/generated.mp4'), + ), + ], + ) + check_op = veo_operation( + name='operations/veo-start', + done=True, + response=completed_response, + ) + + client.aio.models.generate_videos = AsyncMock(return_value=start_op) + client.aio.operations.get = AsyncMock(return_value=check_op) + return client + + +@patch('genkit_google_genai.google.genai.client.Client') +@patch('genkit_google_genai.google._list_genai_models') +@pytest.mark.asyncio +async def test_veo_generate_returns_operation(mock_list_models: MagicMock, mock_client_ctor: MagicMock) -> None: + """generate() on a Veo model returns an Operation to poll.""" + models = GenaiModels() + models.veo = ['veo-2.0-generate-001'] + mock_list_models.return_value = models + mock_client_ctor.return_value = mock_veo_client() + + ai = Genkit(plugins=[GoogleAI(api_key='test-key')]) + response = await ai.generate( + model='googleai/veo-2.0-generate-001', + prompt='a cat surfing', + ) + + assert response.operation is not None + assert response.operation.id == 'operations/veo-start' + assert response.operation.done is False + assert response.operation.action == '/background-model/googleai/veo-2.0-generate-001' + assert response.message is None + + +@patch('genkit_google_genai.google.genai.client.Client') +@patch('genkit_google_genai.google._list_genai_models') +@pytest.mark.asyncio +async def test_veo_generate_operation_poll_loop(mock_list_models: MagicMock, mock_client_ctor: MagicMock) -> None: + """generate_operation + check_operation poll Veo to a ModelResponse output.""" + models = GenaiModels() + models.veo = ['veo-2.0-generate-001'] + mock_list_models.return_value = models + mock_client_ctor.return_value = mock_veo_client() + + ai = Genkit(plugins=[GoogleAI(api_key='test-key')]) + operation = await ai.generate_operation( + model='googleai/veo-2.0-generate-001', + prompt='a cat surfing', + ) + + assert operation.id == 'operations/veo-start' + assert operation.done is False + + operation = await ai.check_operation(operation) + + assert operation.done is True + assert isinstance(operation.output, ModelResponse) + assert operation.output.finish_reason == FinishReason.STOP + content = operation.output.message.content if operation.output.message else [] + assert len(content) == 1 + media = content[0].root.media + assert media is not None and media.url == 'https://example.com/generated.mp4' diff --git a/py/samples/google-genai-media/README.md b/py/samples/google-genai-media/README.md index 859e50fcc1..8b2557e83c 100644 --- a/py/samples/google-genai-media/README.md +++ b/py/samples/google-genai-media/README.md @@ -27,3 +27,5 @@ Flows: `generate_speech`, `generate_image`, `generate_video`. - `googleai/veo-2.0-generate-001` The flow input includes Veo config fields such as `aspect_ratio`, `duration_seconds`, `resolution`, and `seed`. + +`generate_video` uses `generate_operation()` to start Veo and `check_operation()` to poll until the video URL is ready in `operation.output`. diff --git a/py/samples/google-genai-media/src/main.py b/py/samples/google-genai-media/src/main.py index 1792289694..1b52b1d2a4 100644 --- a/py/samples/google-genai-media/src/main.py +++ b/py/samples/google-genai-media/src/main.py @@ -17,16 +17,11 @@ """Google GenAI media - one simple example each for speech, image, and video.""" import asyncio -import time -from typing import Any, Literal -from genkit_google_genai import GoogleAI +from genkit_google_genai import GoogleAI, VeoConfig from pydantic import BaseModel, Field from genkit import Genkit -from genkit._core._background import lookup_background_action -from genkit._core._typing import Operation, Part, Role, TextPart -from genkit.model import Message, ModelRequest ai = Genkit(plugins=[GoogleAI()]) @@ -47,110 +42,52 @@ class ImageInput(BaseModel): class VideoInput(BaseModel): """Input for Veo.""" - model: Literal[ - 'googleai/veo-3.1-generate-preview', - 'googleai/veo-3.1-fast-generate-preview', - 'googleai/veo-3.1-generate-001', - 'googleai/veo-3.1-fast-generate-001', - 'googleai/veo-3.0-generate-001', - 'googleai/veo-3.0-fast-generate-001', - 'googleai/veo-2.0-generate-001', - ] = Field(default='googleai/veo-3.1-generate-preview', description='Veo model for generation') + model: str = Field(default='googleai/veo-3.1-generate-preview', description='Veo model for generation') prompt: str = Field( default='A paper airplane gliding through a bright classroom, cinematic slow motion', description='Video prompt', ) - aspect_ratio: str = Field(default='16:9', description='Video aspect ratio') - duration_seconds: int = Field(default=5, description='Video duration in seconds') - resolution: str | None = Field( - default=None, description='Output resolution (for supported models, e.g. "720p", "1080p")' + config: VeoConfig = Field( + default_factory=lambda: VeoConfig(aspect_ratio='16:9', duration_seconds=5), + description='Veo model configuration', ) - seed: int | None = Field(default=None, description='Optional RNG seed') - - -def _first_media_url(response: Any) -> str | None: - """Return the first media URL in a model response.""" - - message = getattr(response, 'message', None) - if not message: - return None - for part in message.content: - media = getattr(part.root, 'media', None) - if media and getattr(media, 'url', None): - return media.url - return None @ai.flow(name='generate_speech') -async def tts_speech_generator(input: SpeechInput) -> dict[str, str | None]: +async def tts_speech_generator(input: SpeechInput) -> str | None: """Turn text into speech with one TTS call.""" - response = await ai.generate( model='googleai/gemini-2.5-flash-preview-tts', prompt=input.text, config={'speech_config': {'voice_config': {'prebuilt_voice_config': {'voice_name': input.voice}}}}, ) - return {'model': 'googleai/gemini-2.5-flash-preview-tts', 'audio_url': _first_media_url(response)} + return response.media[0].url if response.media else None @ai.flow(name='generate_image') -async def imagen_image_generator(input: ImageInput) -> dict[str, str | None]: +async def imagen_image_generator(input: ImageInput) -> str | None: """Generate one image with Imagen.""" - response = await ai.generate( model='googleai/imagen-3.0-generate-002', prompt=input.prompt, config={'number_of_images': 1}, ) - return {'model': 'googleai/imagen-3.0-generate-002', 'image_url': _first_media_url(response)} - - -async def _poll_video(operation: Operation, model_name: str) -> Operation: - """Wait for a background video operation to finish.""" + return response.media[0].url if response.media else None - action = await lookup_background_action(ai.registry, f'/background-model/{model_name}') - if action is None: - raise ValueError(f'Veo background model not found: {model_name}') - started_at = time.monotonic() +@ai.flow(name='generate_video') +async def veo_video_generator(input: VideoInput) -> str | None: + """Generate one Veo video with generate_operation() and poll to completion.""" + operation = await ai.generate_operation( + model=input.model, + prompt=input.prompt, + config=input.config, + ) while not operation.done: - if time.monotonic() - started_at > 180: - raise TimeoutError('Timed out waiting for Veo output') await asyncio.sleep(3) - operation = await action.check(operation) - return operation + operation = await ai.check_operation(operation) - -@ai.flow(name='generate_video') -async def veo_video_generator(input: VideoInput) -> dict[str, str | int | None]: - """Generate one video by starting and polling a background model.""" - - action = await lookup_background_action(ai.registry, f'/background-model/{input.model}') - if action is None: - raise ValueError(f'Veo background model not found: {input.model}') - - operation = await action.start( - ModelRequest( - messages=[Message(role=Role.USER, content=[Part(root=TextPart(text=input.prompt))])], - config=input.model_dump(exclude_none=True, exclude={'prompt', 'model'}), - ) - ) - operation = await _poll_video(operation, input.model) - - video_url = None - if isinstance(operation.output, dict): - message = operation.output.get('message', {}) - content = message.get('content', []) - if content: - media = content[0].get('media', {}) - video_url = media.get('url') - - return { - 'model': input.model, - 'operation_id': operation.id, - 'video_url': video_url, - 'duration_seconds': input.duration_seconds, - } + return operation.output.media[0].url if operation.output and operation.output.media else None async def main() -> None: