Skip to content

chore(py/plugins/google-genai): trim Veo catalog to 3.1 and add lite variants #6300

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hilariie wants to merge 1 commit into from
Open

chore(py/plugins/google-genai): trim Veo catalog to 3.1 and add lite variants #6300

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion py/packages/genkit-google-genai/src/genkit_google_genai/google.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -996,7 +996,7 @@ class VertexAI(GoogleFamilyRefs, Plugin):
| Type | Action Kind | Example |
|---|---|---|
| Gemini / Gemma | MODEL | ``vertexai/gemini-flash-latest`` |
| 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:
Expand Down
15 changes: 6 additions & 9 deletions py/packages/genkit-google-genai/src/genkit_google_genai/models/veo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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',
]


Expand Down
26 changes: 14 additions & 12 deletions py/packages/genkit-google-genai/test/google_plugin_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -410,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()
Expand All @@ -420,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


Expand Down Expand Up @@ -815,7 +815,7 @@ async def test_vertexai_list_actions(vertexai_plugin_instance: VertexAI) -> None
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'

Expand All @@ -836,7 +836,7 @@ async def test_vertexai_list_actions(vertexai_plugin_instance: VertexAI) -> None
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
Expand Down Expand Up @@ -864,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

Expand All @@ -874,7 +874,7 @@ def mock_model(name: str) -> MagicMock:
# Gemini text model discovered despite supported_actions=None.
assert vertexai_name('gemini-2.5-pro') in names
# Veo discovered; Imagen is not advertised.
assert vertexai_name('veo-2.0-generate-001') in names
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.
Expand All @@ -889,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
Expand All @@ -901,11 +903,11 @@ 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
Expand All @@ -929,7 +931,7 @@ async def test_vertexai_list_known_models(vertexai_plugin_instance: VertexAI) ->
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'

Expand All @@ -946,7 +948,7 @@ async def test_vertexai_list_known_models(vertexai_plugin_instance: VertexAI) ->
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}
Expand Down
8 changes: 4 additions & 4 deletions py/packages/genkit-google-genai/test/models/family_name_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@
('gemini-2.5-flash-image', False),
('gemini-2.5-flash-image-preview', False),
('gemini-2.5-flash', False),
('veo-3.0-generate-001', False),
('veo-3.1-generate-001', False),
],
)
def test_is_unsupported_image_model_name(name: str, expected: bool) -> None:
Expand Down Expand Up @@ -122,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),
],
Expand Down Expand Up @@ -161,7 +161,7 @@ def test_is_lyria_model(name: str, expected: bool) -> None:
'virtual-try-on-001',
'imagen-3.0-generate-002',
'vertexai/imagen-4.0-generate-001',
'veo-3.0-generate-001',
'veo-3.1-generate-001',
'googleai/lyria-002',
'models/deep-research-pro-preview',
'googleai/deep-research-pro-preview',
Expand Down
17 changes: 15 additions & 2 deletions py/packages/genkit-google-genai/test/models/model_refs_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,8 @@ def test_enum_names_still_work(self) -> None:
assert GoogleAI.gemini_model(GoogleAIGeminiVersion.GEMINI_2_5_FLASH).name == 'googleai/gemini-2.5-flash'
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."""
Expand Down Expand Up @@ -154,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
Expand Down Expand Up @@ -184,7 +186,7 @@ def test_error_points_at_the_right_constructor(self) -> None:
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'):
Expand Down Expand Up @@ -273,6 +275,17 @@ def test_family_literals_match_their_catalog(self) -> None:
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."""
Expand Down
Loading
Loading