feat(py/vertexai): Enchance VertexAI plugin

py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/__init__.py

@@ -8,7 +8,7 @@
 
 from genkit.plugins.vertex_ai.embedding import EmbeddingModels
 from genkit.plugins.vertex_ai.gemini import GeminiVersion
-from genkit.plugins.vertex_ai.imagen import ImagenVersion
+from genkit.plugins.vertex_ai.imagen import ImagenOptions, ImagenVersion
 from genkit.plugins.vertex_ai.plugin_api import VertexAI, vertexai_name
 
 
@@ -28,4 +28,5 @@ def package_name() -> str:
     EmbeddingModels.__name__,
     GeminiVersion.__name__,
     ImagenVersion.__name__,
+    ImagenOptions.__name__,
 ]

py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/embedding.py

@@ -83,7 +83,7 @@ def embedding_model(self) -> TextEmbeddingModel:
         """
         return TextEmbeddingModel.from_pretrained(self._version)
 
-    def handle_request(self, request: EmbedRequest) -> EmbedResponse:
+    def generate(self, request: EmbedRequest) -> EmbedResponse:
         """Handle an embedding request.
 
         Args:

py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/gemini.py

@@ -11,6 +11,7 @@
 from enum import StrEnum
 from typing import Any
 
+from genkit.core.action import ActionRunContext
 from genkit.core.typing import (
     GenerateRequest,
     GenerateResponse,
@@ -20,7 +21,8 @@
     Supports,
     TextPart,
 )
-from vertexai.generative_models import Content, GenerativeModel, Part
+from genkit.plugins.vertex_ai import request_builder
+from vertexai.generative_models import GenerativeModel
 
 
 class GeminiVersion(StrEnum):
@@ -101,32 +103,39 @@ def gemini_model(self) -> GenerativeModel:
         """
         return GenerativeModel(self._version)
 
-    def handle_request(self, request: GenerateRequest) -> GenerateResponse:
+    def generate(
+        self, request: GenerateRequest, ctx: ActionRunContext
+    ) -> GenerateResponse:
         """Handle a generation request using the Gemini model.
 
         Args:
             request: The generation request containing messages and parameters.
+            ctx: additional context
 
         Returns:
             The model's response to the generation request.
         """
-        messages: list[Content] = []
-        for m in request.messages:
-            parts: list[Part] = []
-            for p in m.content:
-                if p.root.text is not None:
-                    parts.append(Part.from_text(p.root.text))
-                else:
-                    raise Exception('unsupported part type')
-            messages.append(Content(role=m.role.value, parts=parts))
-        response = self.gemini_model.generate_content(contents=messages)
+
+        is_streaming = ctx.is_streaming
+        messages = request_builder.build_messages(request)
+        response = self.gemini_model.generate_content(
+            contents=messages, stream=is_streaming
+        )
+
+        if is_streaming:
+            for chunk in response:
+                ctx.send_chunk(chunk=chunk)
+
         return GenerateResponse(
             message=Message(
                 role=Role.MODEL,
                 content=[TextPart(text=response.text)],
             )
         )
 
+    def chat(self):
+        pass
+
     @property
     def model_metadata(self) -> dict[str, dict[str, Any]]:
         supports = SUPPORTED_MODELS[self._version].supports.model_dump()

py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/imagen.py

@@ -2,8 +2,9 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from enum import StrEnum
-from typing import Any
+from typing import Any, Literal
 
+from genkit.core.action import ActionRunContext
 from genkit.core.typing import (
     GenerateRequest,
     GenerateResponse,
@@ -14,6 +15,8 @@
     Role,
     Supports,
 )
+from genkit.plugins.vertex_ai import request_builder
+from pydantic import BaseModel
 from vertexai.preview.vision_models import ImageGenerationModel
 
 
@@ -23,6 +26,20 @@ class ImagenVersion(StrEnum):
     IMAGEN2 = 'imagegeneration@006'
 
 
+class ImagenOptions(BaseModel):
+    prompt: str
+    number_of_images: int
+    language: Literal[
+        'auto', 'en', 'es', 'hi', 'ja', 'ko', 'pt', 'zh-TW', 'zh', 'zh-CN'
+    ]
+    aspect_ratio: Literal['1:1', '9:16', '16:9', '3:4', '4:3']
+    safety_filter_level: Literal[
+        'block_most', 'block_some', 'block_few', 'block_fewest'
+    ]
+    person_generation: Literal['dont_allow', 'allow_adult', 'allow_all']
+    negative_prompt: bool = False
+
+
 SUPPORTED_MODELS = {
     ImagenVersion.IMAGEN3: ModelInfo(
         label='Vertex AI - Imagen3',
@@ -67,16 +84,23 @@ def __init__(self, version):
     def model(self) -> ImageGenerationModel:
         return ImageGenerationModel.from_pretrained(self._version)
 
-    def handle_request(self, request: GenerateRequest) -> GenerateResponse:
-        parts: list[str] = []
-        for m in request.messages:
-            for p in m.content:
-                if p.root.text is not None:
-                    parts.append(p.root.text)
-                else:
-                    raise Exception('unsupported part type')
+    def generate(
+        self, request: GenerateRequest, ctx: ActionRunContext
+    ) -> GenerateResponse:
+        """Handle a generation request using the Imagen model.
+
+        Args:
+            request: The generation request containing messages and parameters.
+            ctx: additional context.
+
+        Returns:
+            The model's response to the generation request.
+        """
+        prompt = request_builder.build_prompt(request)
+
+        if request.config:
+            pass
 
-        prompt = ' '.join(parts)
         images = self.model.generate_images(
             prompt=prompt,
             number_of_images=1,

py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/plugin_api.py

@@ -74,7 +74,7 @@ def initialize(self, registry: Registry) -> None:
             registry.register_action(
                 kind=ActionKind.MODEL,
                 name=vertexai_name(model_version),
-                fn=gemini.handle_request,
+                fn=gemini.generate,
                 metadata=gemini.model_metadata,
             )
 
@@ -83,7 +83,7 @@ def initialize(self, registry: Registry) -> None:
             registry.register_action(
                 kind=ActionKind.EMBEDDER,
                 name=vertexai_name(embed_model),
-                fn=embedder.handle_request,
+                fn=embedder.generate,
                 metadata=embedder.model_metadata,
             )
 
@@ -92,6 +92,6 @@ def initialize(self, registry: Registry) -> None:
             registry.register_action(
                 kind=ActionKind.MODEL,
                 name=vertexai_name(imagen_version),
-                fn=imagen.handle_request,
+                fn=imagen.generate,
                 metadata=imagen.model_metadata,
             )

py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/request_builder.py

@@ -0,0 +1,33 @@
+# Copyright 2025 Google LLC
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+
+from genkit.core.typing import GenerateRequest, TextPart
+from vertexai.generative_models import Content, Part
+
+LOG = logging.getLogger(__name__)
+
+
+def build_prompt(request: GenerateRequest) -> str:
+    prompt = []
+    for message in request.messages:
+        for text_part in message.content:
+            if isinstance(text_part.root, TextPart):
+                prompt.append(text_part.root.text)
+            else:
+                LOG.error('Non-text messages are not supported')
+    return ' '.join(prompt)
+
+
+def build_messages(request: GenerateRequest) -> list[Content]:
+    messages: list[Content] = []
+    for message in request.messages:
+        parts: list[Part] = []
+        for text_part in message.content:
+            if isinstance(text_part.root, TextPart):
+                parts.append(Part.from_text(text_part.root.text))
+            else:
+                LOG.error('Non-text messages are not supported')
+        messages.append(Content(role=message.role.value, parts=parts))
+    return messages

py/plugins/vertex-ai/tests/request_builder.py

@@ -0,0 +1,62 @@
+# Copyright 2025 Google LLC
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+from genkit.core.typing import GenerateRequest, Message, Role, TextPart
+from genkit.plugins.vertex_ai import request_builder
+
+MULTILINE_CONTENT = [
+    'Hi!',
+    'I have a question for you.',
+    'Where can I read a Genkit documentation?',
+]
+
+
+@pytest.fixture
+def setup_request():
+    request = GenerateRequest(
+        messages=[
+            Message(
+                role=Role.USER,
+                content=[TextPart(text=x) for x in MULTILINE_CONTENT],
+            ),
+        ],
+    )
+
+    return request
+
+
+def test_create_prompt(setup_request):
+    request = setup_request
+    result = request_builder.build_prompt(request)
+    expected = ' '.join(MULTILINE_CONTENT)
+    assert isinstance(result, str)
+    assert result == expected
+
+
+def test_built_gemini_message_multiple_parts(setup_request):
+    request = setup_request
+    result = request_builder.build_messages(request)
+    assert isinstance(result, list)
+    assert isinstance(result[0].parts, list)
+    assert len(result[0].parts) == len(MULTILINE_CONTENT)
+
+    for part, text in zip(result[0].parts, MULTILINE_CONTENT):
+        assert part.text == text
+
+
+def test_built_gemini_message_multiple_messages():
+    request = GenerateRequest(
+        messages=[
+            Message(
+                role=Role.USER,
+                content=[TextPart(text=text)],
+            )
+            for text in MULTILINE_CONTENT
+        ],
+    )
+    result = request_builder.build_messages(request)
+    assert isinstance(result, list)
+    assert len(result) == len(MULTILINE_CONTENT)
+    for message, text in zip(result, MULTILINE_CONTENT):
+        assert message.parts[0].text == text

py/plugins/vertex-ai/tests/test_embeddings.py

@@ -0,0 +1,26 @@
+# Copyright 2025 Google LLC
+# SPDX-License-Identifier: Apache-2.0
+
+import pytest
+from genkit.ai.embedding import EmbedRequest, EmbedResponse
+from genkit.plugins.vertex_ai.embedding import Embedder, EmbeddingModels
+
+
+@pytest.mark.parametrize('version', [x for x in EmbeddingModels])
+def test_generate_text_response(mocker, version):
+    """Tests generate method for embeddings."""
+    mocked_respond = []
+    request = EmbedRequest(documents=['Text1', 'Text2', 'Text3'])
+    embedder = Embedder(version)
+    genai_model_mock = mocker.MagicMock()
+    model_response_mock = mocker.MagicMock()
+    model_response_mock.text = mocked_respond
+    genai_model_mock.generate_content.return_value = model_response_mock
+    mocker.patch(
+        'genkit.plugins.vertex_ai.embedding.Embedder.embedding_model',
+        genai_model_mock,
+    )
+
+    response = embedder.generate(request)
+    assert isinstance(response, EmbedResponse)
+    assert response.embeddings == mocked_respond

py/plugins/vertex-ai/tests/test_gemini.py

@@ -4,6 +4,7 @@
 """Test Gemini models."""
 
 import pytest
+from genkit.core.action import ActionRunContext
 from genkit.core.typing import (
     GenerateRequest,
     GenerateResponse,
@@ -36,7 +37,8 @@ def test_generate_text_response(mocker, version):
         'genkit.plugins.vertex_ai.gemini.Gemini.gemini_model', genai_model_mock
     )
 
-    response = gemini.handle_request(request)
+    ctx = ActionRunContext()
+    response = gemini.generate(request, ctx)
     assert isinstance(response, GenerateResponse)
     assert response.message.content[0].root.text == mocked_respond
 

py/plugins/vertex-ai/tests/test_imagen.py

@@ -7,6 +7,7 @@
 """Test Gemini models."""
 
 import pytest
+from genkit.core.action import ActionRunContext
 from genkit.core.typing import (
     GenerateRequest,
     GenerateResponse,
@@ -41,7 +42,8 @@ def test_generate(mocker, version):
         'genkit.plugins.vertex_ai.imagen.Imagen.model', genai_model_mock
     )
 
-    response = imagen.handle_request(request)
+    ctx = ActionRunContext()
+    response = imagen.generate(request, ctx)
     assert isinstance(response, GenerateResponse)
     assert isinstance(response.message.content[0].root.media, Media1)
     assert response.message.content[0].root.media.url == mocked_respond

py/samples/hello/README.md

@@ -13,5 +13,5 @@ source .venv/bin/activate
 TODO
 
 ```bash
-genkit start -- uv run --directory py samples/hello/main.py
+genkit start -- uv run --directory py samples/hello/src/hello.py
 ```

py/samples/hello/src/hello.py

@@ -201,6 +201,8 @@ async def main() -> None:
     print(
         await embed_docs(['banana muffins? ', 'banana bread? banana muffins?'])
     )
+    await streaming_async_flow()
+    streaming_sync_flow()
 
 
 if __name__ == '__main__':