fix(plugins/google-genai): address feedback

Author: CorieW

js/plugins/google-genai/src/googleai/embedder.ts

@@ -141,7 +141,7 @@ export function defineEmbedder(
       configSchema: ref.configSchema,
       info: ref.info,
     },
-    async (request, _) => {
+    async (request) => {
       const embedApiKey = calculateApiKey(
         pluginOptions?.apiKey,
         request.options?.apiKey

js/plugins/google-genai/src/googleai/index.ts

@@ -15,25 +15,26 @@
  */
 
 import { ActionMetadata, EmbedderReference, ModelReference, z } from 'genkit';
+import { logger } from 'genkit/logging';
 import {
   GenkitPluginV2,
   ResolvableAction,
   genkitPluginV2,
 } from 'genkit/plugin';
-import type { ActionType } from 'genkit/registry';
+import { ActionType } from 'genkit/registry';
 
 // These are namespaced because they all intentionally have
 // functions of the same name with the same arguments.
 // (All exports from these files are used here)
-import { logger } from 'genkit/logging';
-import { extractErrMsg } from '../common/utils.js';
-import { listModels } from './client.js';
 import * as embedder from './embedder.js';
 import * as gemini from './gemini.js';
 import * as imagen from './imagen.js';
+import * as veo from './veo.js';
+
+import { extractErrMsg } from '../common/utils.js';
+import { listModels } from './client.js';
 import { GoogleAIPluginOptions } from './types.js';
 import { calculateApiKey } from './utils.js';
-import * as veo from './veo.js';
 
 export { type EmbeddingConfig } from './embedder.js';
 export { type GeminiConfig, type GeminiTtsConfig } from './gemini.js';
@@ -111,7 +112,7 @@ export function googleAIPlugin(
   return genkitPluginV2({
     name: 'googleai',
     init: async () => await initializer(options),
-    resolve: async (actionType, actionName) =>
+    resolve: async (actionType: ActionType, actionName: string) =>
       await resolver(actionType, actionName, options || {}),
     list: async () => {
       if (listActionsCache) return listActionsCache;

js/plugins/google-genai/src/googleai/veo.ts

@@ -29,7 +29,6 @@ import {
   type ModelReference,
 } from 'genkit/model';
 import { backgroundModel } from 'genkit/plugin';
-
 import { veoCheckOperation, veoPredict } from './client.js';
 import {
   ClientOptions,

js/plugins/google-genai/src/vertexai/embedder.ts

@@ -108,13 +108,6 @@ export const KNOWN_MODELS = {
     supports: { input: ['text'] },
   }),
 } as const;
-export type KnownModels = keyof typeof KNOWN_MODELS;
-export type EmbedderModelName = `embedder=${string}`;
-export function isEmbedderModelName(
-  value?: string
-): value is EmbedderModelName {
-  return !!value?.includes('embedding');
-}
 
 export function createEmbedderRef(
   version: string,
@@ -165,7 +158,7 @@ export function defineKnownModels(
 export function defineEmbedder(
   name: string,
   clientOptions: ClientOptions,
-  pluginOptions?: VertexPluginOptions
+  _?: VertexPluginOptions
 ): EmbedderAction<any> {
   const ref = createEmbedderRef(name);
 
@@ -175,7 +168,7 @@ export function defineEmbedder(
       configSchema: ref.configSchema,
       info: ref.info!,
     },
-    async (request, options) => {
+    async (request) => {
       const embedContentRequest: EmbedContentRequest = {
         instances: request.input.map((doc: Document) =>
           toEmbeddingInstance(ref, doc, request.options)

js/plugins/google-genai/src/vertexai/index.ts

@@ -22,14 +22,14 @@
 
 import { EmbedderReference, ModelReference, z } from 'genkit';
 import { GenkitPluginV2, genkitPluginV2 } from 'genkit/plugin';
+import { ActionType } from 'genkit/registry';
 
 import * as embedder from './embedder.js';
 import * as gemini from './gemini.js';
 import * as imagen from './imagen.js';
 import * as lyria from './lyria.js';
 import * as veo from './veo.js';
 
-import { ActionType } from 'genkit/registry';
 import { listModels } from './client.js';
 import { VertexPluginOptions } from './types.js';
 import { getDerivedOptions } from './utils.js';
@@ -45,6 +45,8 @@ async function initializer(pluginOptions?: VertexPluginOptions) {
   const clientOptions = await getDerivedOptions(pluginOptions);
 
   return [
+    ...veo.defineKnownModels(clientOptions, pluginOptions),
+    ...lyria.defineKnownModels(clientOptions, pluginOptions),
     ...imagen.defineKnownModels(clientOptions, pluginOptions),
     ...gemini.defineKnownModels(clientOptions, pluginOptions),
     ...embedder.defineKnownModels(clientOptions, pluginOptions),
@@ -59,25 +61,37 @@ async function resolver(
   const clientOptions = await getDerivedOptions(pluginOptions);
   switch (actionType) {
     case 'model':
-      if (imagen.isImagenModelName(actionName)) {
+      if (lyria.isLyriaModelName(actionName)) {
+        return lyria.defineModel(actionName, clientOptions, pluginOptions);
+      } else if (imagen.isImagenModelName(actionName)) {
         return imagen.defineModel(actionName, clientOptions, pluginOptions);
+      } else if (veo.isVeoModelName(actionName)) {
+        // no-op (not gemini)
       } else {
         return gemini.defineModel(actionName, clientOptions, pluginOptions);
       }
+      break;
+    case 'background-model':
+      if (veo.isVeoModelName(actionName)) {
+        return veo.defineModel(actionName, clientOptions, pluginOptions);
+      }
+      break;
     case 'embedder':
       return embedder.defineEmbedder(actionName, clientOptions, pluginOptions);
-    default:
-      return undefined;
   }
+
+  return undefined;
 }
 
 async function listActions(options?: VertexPluginOptions) {
-  const clientOptions = await getDerivedOptions(options);
   try {
+    const clientOptions = await getDerivedOptions(options);
     const models = await listModels(clientOptions);
     return [
       ...gemini.listActions(models),
       ...imagen.listActions(models),
+      ...lyria.listActions(models),
+      ...veo.listActions(models),
       // We don't list embedders here
     ];
   } catch (e: unknown) {