feat: Modelgarden uses openai plugin

Author: zarinn3pal

py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/__init__.py

@@ -19,7 +19,7 @@
 
 from .handler import OpenAIModelHandler
 from .model import OpenAIModel
-from .model_info import SUPPORTED_OPENAI_MODELS
+from .model_info import SUPPORTED_OPENAI_COMPAT_MODELS, SUPPORTED_OPENAI_MODELS, PluginSource
 
 
 def package_name() -> str:
@@ -28,7 +28,9 @@ def package_name() -> str:
 
 __all__ = [
     'OpenAIModel',
+    'PluginSource',
     'SUPPORTED_OPENAI_MODELS',
+    'SUPPORTED_OPENAI_COMPAT_MODELS',
     'OpenAIModelHandler',
     'package_name',
 ]

py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/handler.py

@@ -25,8 +25,9 @@
 from genkit.ai import ActionRunContext, GenkitRegistry
 from genkit.plugins.compat_oai.models.model import OpenAIModel
 from genkit.plugins.compat_oai.models.model_info import (
+    SUPPORTED_OPENAI_COMPAT_MODELS,
     SUPPORTED_OPENAI_MODELS,
-    SUPPORTED_OPENAI_COMPAT_MODELS
+    PluginSource,
 )
 from genkit.plugins.compat_oai.typing import OpenAIConfig
 from genkit.types import (
@@ -38,17 +39,35 @@
 class OpenAIModelHandler:
     """Handles OpenAI API interactions for the Genkit plugin."""
 
-    def __init__(self, model: Any) -> None:
+    def __init__(self, model: Any, source: PluginSource = PluginSource.OPENAI) -> None:
         """Initializes the OpenAIModelHandler with a specified model.
 
         Args:
             model: An instance of a Model subclass representing the OpenAI model.
+            source: Helps distinguish if model handler is called from model-garden plugin.
+                    Default source is openai.
         """
         self._model = model
+        self._source = source
+
+    @staticmethod
+    def _get_supported_models(source: PluginSource) -> dict[str, Any]:
+        """Returns the supported models based on the plugin source.
+        Args:
+            source: Helps distinguish if model handler is called from model-garden plugin.
+                    Default source is openai.
+
+        Returns:
+            Openai models if source is openai. Merges supported openai models with openai-compat models if source is model-garden.
+
+        """
+        if source == PluginSource.MODEL_GARDEN:
+            return {**SUPPORTED_OPENAI_MODELS, **SUPPORTED_OPENAI_COMPAT_MODELS}
+        return SUPPORTED_OPENAI_MODELS
 
     @classmethod
     def get_model_handler(
-        cls, model: str, client: OpenAI, registry: GenkitRegistry
+        cls, model: str, client: OpenAI, registry: GenkitRegistry, source: PluginSource = PluginSource.OPENAI
     ) -> Callable[[GenerateRequest, ActionRunContext], GenerateResponse]:
         """Factory method to initialize the model handler for the specified OpenAI model.
 
@@ -62,18 +81,22 @@ def get_model_handler(
             model: The OpenAI model name.
             client: OpenAI client instance.
             registry: Genkit registry instance.
+            source: Helps distinguish if model handler is called from model-garden plugin.
+                    Default source is openai.
 
         Returns:
             A callable function that acts as an action handler.
 
         Raises:
             ValueError: If the specified model is not supported.
         """
-        if model not in SUPPORTED_OPENAI_MODELS:
+        supported_models = cls._get_supported_models(source)
+
+        if model not in supported_models:
             raise ValueError(f"Model '{model}' is not supported.")
 
         openai_model = OpenAIModel(model, client, registry)
-        return cls(openai_model).generate
+        return cls(openai_model, source).generate
 
     def _validate_version(self, version: str) -> None:
         """Validates whether the specified model version is supported.
@@ -84,7 +107,8 @@ def _validate_version(self, version: str) -> None:
         Raises:
             ValueError: If the specified model version is not supported.
         """
-        model_info = SUPPORTED_OPENAI_MODELS[self._model.name]
+        supported_models = self._get_supported_models(self._source)
+        model_info = supported_models[self._model.name]
         if version not in model_info.versions:
             raise ValueError(f"Model version '{version}' is not supported.")
 

py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model_info.py

@@ -17,12 +17,23 @@
 
 """OpenAI Compatible Models for Genkit."""
 
+from enum import StrEnum
+
 from genkit.plugins.compat_oai.typing import SupportedOutputFormat
 from genkit.types import (
     ModelInfo,
     Supports,
 )
 
+OPENAI = 'openai'
+MODEL_GARDEN = 'model-garden'
+
+
+class PluginSource(StrEnum):
+    OPENAI = 'openai'
+    MODEL_GARDEN = 'model-garden'
+
+
 GPT_3_5_TURBO = 'gpt-3.5-turbo'
 GPT_4 = 'gpt-4'
 GPT_4_TURBO = 'gpt-4-turbo'
@@ -136,32 +147,7 @@
             media=False,
             tools=True,
             systemRole=True,
-            output=['text', 'json']
-        ),
-    ),
-    LLAMA_3_2: ModelInfo(
-        versions=['meta/llama-3.2-90b-vision-instruct-maas'],
-        label='llama-3.2',
-        supports=Supports(
-            multiturn=True,
-            media=True,
-            tools=True,
-            systemRole=True,
-            output=['text', 'json']
-        ),
-    ),
-}
-
-SUPPORTED_OPENAI_COMPAT_MODELS: dict[str, ModelInfo] = {
-    LLAMA_3_1: ModelInfo(
-        versions=['meta/llama3-405b-instruct-maas'],
-        label='llama-3.1',
-        supports=Supports(
-            multiturn=True,
-            media=False,
-            tools=True,
-            systemRole=True,
-            output=['text', 'json']
+            output=[SupportedOutputFormat.JSON_MODE, SupportedOutputFormat.TEXT],
         ),
     ),
     LLAMA_3_2: ModelInfo(
@@ -172,7 +158,7 @@
             media=True,
             tools=True,
             systemRole=True,
-            output=['text', 'json']
+            output=[SupportedOutputFormat.JSON_MODE, SupportedOutputFormat.TEXT],
         ),
     ),
 }

py/plugins/compat-oai/src/genkit/plugins/compat_oai/openai_plugin.py

@@ -24,8 +24,8 @@
     SUPPORTED_OPENAI_MODELS,
     OpenAIModelHandler,
 )
-from genkit.plugins.compat_oai.typing import OpenAIConfig
 from genkit.plugins.compat_oai.openai_client_handler import OpenAIClientHandler
+from genkit.plugins.compat_oai.typing import OpenAIConfig
 
 
 class OpenAI(Plugin):

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

@@ -14,7 +14,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from .model_garden import OpenAIFormatModelVersion, vertexai_name
+from .model_garden import vertexai_name
 from .modelgarden_plugin import VertexAIModelGarden
 
-__all__ = [OpenAIFormatModelVersion.__name__, vertexai_name, VertexAIModelGarden.__name__]
+__all__ = [vertexai_name, VertexAIModelGarden.__name__]

py/plugins/compat-oai/src/genkit/plugins/compat_oai/openai_client_handler.py renamed to py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/client.py

@@ -14,12 +14,22 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
-from openai import OpenAI as OpenAIClient
+from google.auth import default, transport
+from openai import OpenAI as _OpenAI
 
 
-class OpenAIClientHandler:
+class OpenAIClient:
     """Handles OpenAI API client Initialization."""
 
-    @staticmethod
-    def initialize_client(plugin_source: Optional[str] = None, **openai_params):
+    def __new__(cls, **openai_params) -> _OpenAI:
         """Initializes the OpenAIClient based on the plugin source."""
+        location = openai_params.get('location')
+        project_id = openai_params.get('project_id')
+        if project_id:
+            credentials, _ = default()
+        else:
+            credentials, project_id = default()
+
+        credentials.refresh(transport.requests.Request())
+        base_url = f'https://{location}-aiplatform.googleapis.com/v1beta1/projects/{project_id}/locations/{location}/endpoints/openapi'
+        return _OpenAI(api_key=credentials.token, base_url=base_url)

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

@@ -1,17 +1,17 @@
 # Copyright 2025 Google LLC
 # SPDX-License-Identifier: Apache-2.0
-import os
-from enum import StrEnum
-
-from genkit.ai.registry import GenkitRegistry
-from genkit.core.typing import (
-    ModelInfo,
-    Supports,
+from genkit.ai import GenkitRegistry
+from genkit.plugins.compat_oai.models import (
+    SUPPORTED_OPENAI_COMPAT_MODELS,
+    OpenAIModelHandler,
+    PluginSource,
 )
+from genkit.plugins.compat_oai.openai_client_handler import PluginSource
 from genkit.plugins.compat_oai.typing import OpenAIConfig
-from genkit.plugins.vertex_ai import constants as const
 
-from .openai_compatibility import OpenAICompatibleModel
+from .client import OpenAIClient
+
+OPENAI_COMPAT = 'openai-compat'
 
 
 def vertexai_name(name: str) -> str:
@@ -26,43 +26,33 @@ def vertexai_name(name: str) -> str:
     return f'vertexai/{name}'
 
 
-class OpenAIFormatModelVersion(StrEnum):
-    """Available versions of the llama model.
-
-    This enum defines the available versions of the llama model that
-    can be used through Vertex AI.
-    """
-
-    LLAMA_3_1 = 'llama-3.1'
-    LLAMA_3_2 = 'llama-3.2'
-
+class ModelGarden:
+    @staticmethod
+    def get_model_info(name: str) -> dict[str, str] | None:
+        """Returns model type and name for a given model.
 
-SUPPORTED_OPENAI_FORMAT_MODELS: dict[str, ModelInfo] = {
-    OpenAIFormatModelVersion.LLAMA_3_1: ModelInfo(
-        versions=['meta/llama3-405b-instruct-maas'],
-        label='llama-3.1',
-        supports=Supports(multiturn=True, media=False, tools=True, systemRole=True, output=['text', 'json']),
-    ),
-    OpenAIFormatModelVersion.LLAMA_3_2: ModelInfo(
-        versions=['meta/llama-3.2-90b-vision-instruct-maas'],
-        label='llama-3.2',
-        supports=Supports(multiturn=True, media=True, tools=True, systemRole=True, output=['text', 'json']),
-    ),
-}
+        Args:
+            name: Name of the model for which type and name are to be returned
 
+        """
+        if SUPPORTED_OPENAI_COMPAT_MODELS.get(name):
+            return {'name': SUPPORTED_OPENAI_COMPAT_MODELS.get(name).label, 'type': OPENAI_COMPAT}
 
-class ModelGarden:
     @classmethod
     def to_openai_compatible_model(cls, ai: GenkitRegistry, model: str, location: str, project_id: str):
-        if model not in SUPPORTED_OPENAI_FORMAT_MODELS:
+        if model not in SUPPORTED_OPENAI_COMPAT_MODELS:
             raise ValueError(f"Model '{model}' is not supported.")
-        model_version = SUPPORTED_OPENAI_FORMAT_MODELS[model].versions[0]
-        open_ai_compat = OpenAICompatibleModel(model_version, project_id, location)
-        supports = SUPPORTED_OPENAI_FORMAT_MODELS[model].supports.model_dump()
+        openai_params = {'location': location, 'project_id': project_id}
+        openai_client = OpenAIClient(**openai_params)
+        handler = OpenAIModelHandler.get_model_handler(
+            model=model, client=openai_client, registry=ai, source=PluginSource.MODEL_GARDEN
+        )
+
+        supports = SUPPORTED_OPENAI_COMPAT_MODELS[model].supports.model_dump()
 
         ai.define_model(
             name=f'vertexai/{model}',
-            fn=open_ai_compat.generate,
+            fn=handler,
             config_schema=OpenAIConfig,
             metadata={'model': {'supports': supports}},
         )

py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/modelgarden_plugin.py

@@ -5,36 +5,11 @@
 """ModelGarden API Compatible Plugin for Genkit."""
 
 import os
-import pdb
-from pprint import pprint
 
-from pydantic import BaseModel, ConfigDict
-
-from genkit.ai.plugin import Plugin
-from genkit.ai.registry import GenkitRegistry
+from genkit.ai import GenkitRegistry, Plugin
 from genkit.plugins.vertex_ai import constants as const
 
-from .model_garden import SUPPORTED_OPENAI_FORMAT_MODELS, ModelGarden
-
-
-class CommonPluginOptions(BaseModel):
-    model_config = ConfigDict(extra='forbid', populate_by_name=True)
-
-    project_id: str | None = None
-    location: str | None = None
-    models: list[str] | None = None
-
-
-def vertexai_name(name: str) -> str:
-    """Create a Vertex AI action name.
-
-    Args:
-        name: Base name for the action.
-
-    Returns:
-        The fully qualified Vertex AI action name.
-    """
-    return f'vertexai/{name}'
+from .model_garden import OPENAI_COMPAT, ModelGarden
 
 
 class VertexAIModelGarden(Plugin):
@@ -46,24 +21,25 @@ class VertexAIModelGarden(Plugin):
     registration of model actions.
     """
 
-    name = 'modelgarden'
+    name = 'vertex-ai-model-garden'
 
-    def __init__(self, **kwargs):
+    def __init__(self, project_id: str | None = None, location: str | None = None, models: list[str] | None = None):
         """Initialize the plugin by registering actions with the registry."""
-        self.plugin_options = CommonPluginOptions(
-            project_id=kwargs.get('project_id', os.getenv(const.GCLOUD_PROJECT)),
-            location=kwargs.get('location', const.DEFAULT_REGION),
-            models=kwargs.get('models'),
-        )
+        self.project_id = project_id if project_id is not None else os.getenv(const.GCLOUD_PROJECT)
+        self.location = location if location is not None else const.DEFAULT_REGION
+        self.models = models
 
     def initialize(self, ai: GenkitRegistry) -> None:
         """Handles actions for various openaicompatible models."""
-        for model in self.plugin_options.models:
-            openai_model = SUPPORTED_OPENAI_FORMAT_MODELS.get(model).label
-            if openai_model:
-                ModelGarden.to_openai_compatible_model(
-                    ai,
-                    model=openai_model,
-                    location=self.plugin_options.location,
-                    project_id=self.plugin_options.project_id,
-                )
+        models = self.models
+        if models:
+            for model in models:
+                model_info = ModelGarden.get_model_info(model)
+                if model_info:
+                    if model_info['type'] == OPENAI_COMPAT:
+                        ModelGarden.to_openai_compatible_model(
+                            ai,
+                            model=model_info['name'],
+                            location=self.location,
+                            project_id=self.project_id,
+                        )