diff --git a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model.py b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model.py index 8d00bcee0e..bf1f7315aa 100644 --- a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model.py +++ b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model.py @@ -241,6 +241,11 @@ def normalize_config(config: Any) -> OpenAIConfig: return config if isinstance(config, dict): + if config.get('topK'): + del config['topK'] + if config.get('topP'): + config['top_p'] = config['topP'] + del config['topP'] return OpenAIConfig(**config) raise ValueError(f'Expected request.config to be a dict or OpenAIConfig, got {type(config).__name__}.') diff --git a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model_info.py b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model_info.py index bcef1fe517..ff9e72ac5f 100644 --- a/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model_info.py +++ b/py/plugins/compat-oai/src/genkit/plugins/compat_oai/models/model_info.py @@ -46,7 +46,7 @@ class PluginSource(StrEnum): GPT_4O_MINI = 'gpt-4o-mini' O1_MINI = 'o1-mini' -LLAMA_3_1 = 'meta/llama3-405b-instruct-maas' +LLAMA_3_1 = 'meta/llama-3.1-405b-instruct-maas' LLAMA_3_2 = 'meta/llama-3.2-90b-vision-instruct-maas' SUPPORTED_OPENAI_MODELS: dict[str, ModelInfo] = { diff --git a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/modelgarden_plugin.py b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/modelgarden_plugin.py index 4093135b93..4746fefee6 100644 --- a/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/modelgarden_plugin.py +++ b/py/plugins/vertex-ai/src/genkit/plugins/vertex_ai/model_garden/modelgarden_plugin.py @@ -17,11 +17,16 @@ """ModelGarden API Compatible Plugin for Genkit.""" import os +from functools import cached_property from genkit.ai import GenkitRegistry, Plugin +from genkit.blocks.model import model_action_metadata +from genkit.core.action.types import ActionKind +from genkit.plugins.compat_oai.models import SUPPORTED_OPENAI_COMPAT_MODELS +from genkit.plugins.compat_oai.typing import OpenAIConfig from genkit.plugins.vertex_ai import constants as const -from .model_garden import MODELGARDEN_PLUGIN_NAME, ModelGarden +from .model_garden import MODELGARDEN_PLUGIN_NAME, ModelGarden, model_garden_name class VertexAIModelGarden(Plugin): @@ -35,8 +40,24 @@ class VertexAIModelGarden(Plugin): name = MODELGARDEN_PLUGIN_NAME - 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.""" + def __init__( + self, + project_id: str | None = None, + location: str | None = None, + models: list[str] | None = None, + ) -> None: + """Initializes the plugin and sets up its configuration. + + This constructor prepares the plugin by assigning the Google Cloud project ID, + location, and a list of models to be used. + + Args: + project_id: The Google Cloud project ID to use. If not provided, it attempts + to load from the `GCLOUD_PROJECT` environment variable. + location: The Google Cloud region to use for services. If not provided, + it defaults to `DEFAULT_REGION`. + models: An optional list of model names to register with the plugin. + """ 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 @@ -55,3 +76,66 @@ def initialize(self, ai: GenkitRegistry) -> None: registry=ai, ) model_proxy.define_model() + + def resolve_action( + self, + ai: GenkitRegistry, + kind: ActionKind, + name: str, + ) -> None: + """Resolves and action. + + Args: + ai: The Genkit registry. + kind: The kind of action to resolve. + name: The name of the action to resolve. + """ + if kind == ActionKind.MODEL: + self._resolve_model(ai=ai, name=name) + + def _resolve_model(self, ai: GenkitRegistry, name: str) -> None: + """Resolves and defines a Model Garden Vertex AI model within the Genkit registry. + + This internal method handles the logic for registering new models + of Vertex AI Model Garden that are compatible with OpenaI + based on the provided name. + It extracts a clean name, determines the model type, instantiates the + appropriate model class, and registers it with the Genkit AI registry. + + Args: + ai: The Genkit AI registry instance to define the model in. + name: The name of the model to resolve. This name might include a + prefix indicating it's from a specific plugin. + """ + clean_name = ( + name.replace(MODELGARDEN_PLUGIN_NAME + '/', '') if name.startswith(MODELGARDEN_PLUGIN_NAME) else name + ) + + model_proxy = ModelGarden( + model=clean_name, + location=self.location, + project_id=self.project_id, + registry=ai, + ) + model_proxy.define_model() + + @cached_property + def list_actions(self) -> list[dict[str, str]]: + """Generate a list of available actions or models. + + Returns: + list of actions dicts with the following shape: + { + 'name': str, + 'kind': ActionKind, + } + """ + actions_list = [] + for model, model_info in SUPPORTED_OPENAI_COMPAT_MODELS.items(): + actions_list.append( + model_action_metadata( + name=model_garden_name(model), info=model_info.model_dump(), config_schema=OpenAIConfig + ) + ) + + return actions_list diff --git a/py/samples/model-garden/src/main.py b/py/samples/model-garden/src/main.py index 4479bbf219..a0087964ac 100644 --- a/py/samples/model-garden/src/main.py +++ b/py/samples/model-garden/src/main.py @@ -25,7 +25,6 @@ plugins=[ VertexAIModelGarden( location='us-central1', - models=['meta/llama-3.2-90b-vision-instruct-maas'], ), ], )