feat(py/vertexai): Enhance VertexAI plugin (#2184)

Author: Irillit

py/packages/genkit/src/genkit/ai/embedding.py

@@ -2,6 +2,7 @@
 # SPDX-License-Identifier: Apache-2.0
 
 from collections.abc import Callable
+from typing import Any
 
 from pydantic import BaseModel
 
@@ -14,6 +15,7 @@ class EmbedRequest(BaseModel):
     """
 
     documents: list[str]
+    options: dict[str, Any] | None = None
 
 
 class EmbedResponse(BaseModel):

py/packages/genkit/src/genkit/veneer/veneer.py

@@ -435,21 +435,27 @@ def generate_stream(
         return (stream, stream.closed)
 
     async def embed(
-        self, model: str | None = None, documents: list[str] | None = None
+        self,
+        model: str | None = None,
+        documents: list[str] | None = None,
+        options: dict[str, Any] | None = None,
     ) -> EmbedResponse:
         """Calculates embeddings for documents.
 
         Args:
             model: Optional embedder model name to use.
             documents: Texts to embed.
+            options: embedding options
 
         Returns:
             The generated response with embeddings.
         """
         embed_action = self.registry.lookup_action(ActionKind.EMBEDDER, model)
 
         return (
-            await embed_action.arun(EmbedRequest(documents=documents))
+            await embed_action.arun(
+                EmbedRequest(documents=documents, options=options)
+            )
         ).response
 
 

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

@@ -6,9 +6,12 @@
 enabling the use of Vertex AI models and services within the Genkit framework.
 """
 
-from genkit.plugins.vertex_ai.embedding import EmbeddingModels
+from genkit.plugins.vertex_ai.embedding import (
+    EmbeddingModels,
+    EmbeddingsTaskType,
+)
 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
 
 
@@ -26,6 +29,8 @@ def package_name() -> str:
     VertexAI.__name__,
     vertexai_name.__name__,
     EmbeddingModels.__name__,
+    EmbeddingsTaskType.__name__,
     GeminiVersion.__name__,
     ImagenVersion.__name__,
+    ImagenOptions.__name__,
 ]

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

@@ -28,7 +28,7 @@ class EmbeddingModels(StrEnum):
     TEXT_EMBEDDING_002_MULTILINGUAL = 'text-multilingual-embedding-002'
 
 
-class TaskType(StrEnum):
+class EmbeddingsTaskType(StrEnum):
     """Task types supported by Vertex AI.
 
     Attributes:
@@ -47,24 +47,21 @@ class TaskType(StrEnum):
     CLUSTERING = 'CLUSTERING'
     RETRIEVAL_DOCUMENT = 'RETRIEVAL_DOCUMENT'
     RETRIEVAL_QUERY = 'RETRIEVAL_QUERY'
+    QUESTION_ANSWERING = 'QUESTION_ANSWERING'
+    FACT_VERIFICATION = 'FACT_VERIFICATION'
+    CODE_RETRIEVAL_QUERY = 'CODE_RETRIEVAL_QUERY'
 
 
 class Embedder:
-    """Embedder for Vertex AI.
+    """Embedder for Vertex AI."""
 
-    Attributes:
-        version: The version of the embedding model to use.
-        task: The task type to use for the embedding.
-        dimensionality: The dimensionality of the embedding.
-    """
-
-    TASK = TaskType.RETRIEVAL_QUERY
+    TASK_KEY = 'task'
+    DEFAULT_TASK = EmbeddingsTaskType.RETRIEVAL_QUERY
 
     # By default, the model generates embeddings with 768 dimensions.
     # Models such as `text-embedding-004`, `text-embedding-005`,
     # and `text-multilingual-embedding-002`allow the output dimensionality
     # to be adjusted between 1 and 768.
-    DIMENSIONALITY = 768
 
     def __init__(self, version: EmbeddingModels):
         """Initialize the embedder.
@@ -81,9 +78,11 @@ def embedding_model(self) -> TextEmbeddingModel:
         Returns:
             The embedding model.
         """
+
+        # TODO: pass additional parameters
         return TextEmbeddingModel.from_pretrained(self._version)
 
-    def handle_request(self, request: EmbedRequest) -> EmbedResponse:
+    def generate(self, request: EmbedRequest) -> EmbedResponse:
         """Handle an embedding request.
 
         Args:
@@ -92,11 +91,18 @@ def handle_request(self, request: EmbedRequest) -> EmbedResponse:
         Returns:
             The embedding response.
         """
-        inputs = [
-            TextEmbeddingInput(text, self.TASK) for text in request.documents
-        ]
-        vertexai_embeddings = self.embedding_model.get_embeddings(inputs)
+        options = request.options
+        task = options.get(self.TASK_KEY) if options else self.DEFAULT_TASK
+        if task not in EmbeddingsTaskType:
+            raise ValueError(f'Unsupported task {task} for VertexAI.')
+
+        del options[self.TASK_KEY]
+        inputs = [TextEmbeddingInput(text, task) for text in request.documents]
+        vertexai_embeddings = self.embedding_model.get_embeddings(
+            inputs, **options
+        )
         embeddings = [embedding.values for embedding in vertexai_embeddings]
+
         return EmbedResponse(embeddings=embeddings)
 
     @property

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

@@ -8,19 +8,29 @@
 definitions and a client class for making requests to Gemini models.
 """
 
+import logging
 from enum import StrEnum
 from typing import Any
 
+import vertexai.generative_models as genai
+from genkit.core.action import ActionRunContext
 from genkit.core.typing import (
+    CustomPart,
+    DataPart,
     GenerateRequest,
     GenerateResponse,
+    GenerateResponseChunk,
+    MediaPart,
     Message,
     ModelInfo,
     Role,
     Supports,
     TextPart,
+    ToolRequestPart,
+    ToolResponsePart,
 )
-from vertexai.generative_models import Content, GenerativeModel, Part
+
+LOG = logging.getLogger(__name__)
 
 
 class GeminiVersion(StrEnum):
@@ -83,7 +93,7 @@ class Gemini:
     handling message formatting and response processing.
     """
 
-    def __init__(self, version: str):
+    def __init__(self, version: str | GeminiVersion):
         """Initialize a Gemini client.
 
         Args:
@@ -92,38 +102,93 @@ def __init__(self, version: str):
         """
         self._version = version
 
+    def is_multimode(self):
+        return SUPPORTED_MODELS[self._version].supports.media
+
+    def build_messages(self, request: GenerateRequest) -> list[genai.Content]:
+        """Builds a list of VertexAI content from a request.
+
+        Args:
+            - request: a packed request for the model
+
+        Returns:
+            - a list of VertexAI GenAI Content for the request
+        """
+        messages: list[genai.Content] = []
+        for message in request.messages:
+            parts: list[genai.Part] = []
+            for part in message.content:
+                if isinstance(part.root, TextPart):
+                    parts.append(genai.Part.from_text(part.root.text))
+                elif isinstance(part.root, MediaPart):
+                    if not self.is_multimode():
+                        LOG.error(
+                            f'The model {self._version} does not'
+                            f' support multimode input'
+                        )
+                        continue
+                    parts.append(
+                        genai.Part.from_uri(
+                            mime_type=part.root.media.content_type,
+                            uri=part.root.media.url,
+                        )
+                    )
+                elif isinstance(part.root, ToolRequestPart | ToolResponsePart):
+                    LOG.warning('Tools are not supported yet')
+                elif isinstance(part.root, CustomPart):
+                    # TODO: handle CustomPart
+                    LOG.warning('The code part is not supported yet.')
+                else:
+                    LOG.error('The type is not supported')
+            messages.append(genai.Content(role=message.role.value, parts=parts))
+
+        return messages
+
     @property
-    def gemini_model(self) -> GenerativeModel:
+    def gemini_model(self) -> genai.GenerativeModel:
         """Get the Vertex AI GenerativeModel instance.
 
         Returns:
             A configured GenerativeModel instance for the specified version.
         """
-        return GenerativeModel(self._version)
+        return genai.GenerativeModel(self._version)
 
-    def handle_request(self, request: GenerateRequest) -> GenerateResponse:
+    def generate(
+        self, request: GenerateRequest, ctx: ActionRunContext
+    ) -> GenerateResponse | None:
         """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)
+
+        messages = self.build_messages(request)
+        response = self.gemini_model.generate_content(
+            contents=messages, stream=ctx.is_streaming
+        )
+
+        text_response = ''
+        if ctx.is_streaming:
+            for chunk in response:
+                # TODO: Support other types of output
+                ctx.send_chunk(
+                    GenerateResponseChunk(
+                        role=Role.MODEL,
+                        content=[TextPart(text=chunk.text)],
+                    )
+                )
+
+        else:
+            text_response = response.text
+
         return GenerateResponse(
             message=Message(
                 role=Role.MODEL,
-                content=[TextPart(text=response.text)],
+                content=[TextPart(text=text_response)],
             )
         )