fix(plugins/google-genai): tests

Author: CorieW

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

@@ -483,7 +483,8 @@ export function defineModel(
       configSchema: ref.configSchema,
       use: middleware,
     },
-    async (request, { streamingRequested, sendChunk, abortSignal }) => {
+    async (request, options) => {
+      const { streamingRequested, sendChunk, abortSignal } = options || {};
       const clientOpt = { ...clientOptions, signal: abortSignal };
 
       // Make a copy so that modifying the request will not produce side-effects

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

@@ -16,14 +16,17 @@
 
 import { EmbedderReference, ModelReference, z } from 'genkit';
 import { GenkitPluginV2, genkitPluginV2 } from 'genkit/plugin';
+import type { 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 { listModels } from './client.js';
 import * as embedder from './embedder.js';
 import * as gemini from './gemini.js';
 import * as imagen from './imagen.js';
 import { GoogleAIPluginOptions } from './types.js';
+import { getApiKeyFromEnvVar } from './utils.js';
 import * as veo from './veo.js';
 
 export { type EmbeddingConfig } from './embedder.js';
@@ -37,6 +40,9 @@ export { type GoogleAIPluginOptions };
 export function googleAIPlugin(
   options?: GoogleAIPluginOptions
 ): GenkitPluginV2 {
+  // Cache for list function - per plugin instance
+  let listCache: any[] | null = null;
+
   return genkitPluginV2({
     name: 'googleai',
     init: async () => {
@@ -47,32 +53,85 @@ export function googleAIPlugin(
         ...embedder.defineKnownModels(options),
       ];
     },
+    resolve: async (actionType: ActionType, name: string) => {
+      // Extract the model name without the plugin prefix
+      const modelName = name.replace(/^googleai\//, '');
+
+      if (actionType === 'model') {
+        if (veo.isVeoModelName(modelName)) {
+          return veo.defineModel(modelName, options);
+        }
+        if (imagen.isImagenModelName(modelName)) {
+          return imagen.defineModel(modelName, options);
+        }
+        // For gemini, tts, gemma, and unknown model families
+        return gemini.defineModel(modelName, options);
+      } else if (actionType === 'embedder') {
+        return embedder.defineEmbedder(modelName, options);
+      } else if (actionType === 'background-model') {
+        if (veo.isVeoModelName(modelName)) {
+          return veo.defineModel(modelName, options);
+        }
+      }
+
+      return undefined;
+    },
+    list: async () => {
+      // Return cached result if available
+      if (listCache !== null) {
+        return listCache;
+      }
+
+      const apiKey = options?.apiKey || getApiKeyFromEnvVar();
+
+      if (!apiKey) {
+        return [];
+      }
+
+      try {
+        const models = await listModels(apiKey, options);
+        const result = [
+          ...imagen.listActions(models),
+          ...gemini.listActions(models),
+          ...veo.listActions(models),
+          ...embedder.listActions(models),
+        ];
+
+        // Cache the result
+        listCache = result;
+        return result;
+      } catch (error) {
+        // Handle API errors gracefully by returning empty array
+        console.warn('GoogleAI plugin: Failed to fetch models list:', error);
+        return [];
+      }
+    },
   });
 }
 
 export type GoogleAIPlugin = {
   (pluginOptions?: GoogleAIPluginOptions): GenkitPluginV2;
-  model(
+  createModelRef(
     name: gemini.KnownGemmaModels | (gemini.GemmaModelName & {}),
     config: gemini.GemmaConfig
   ): ModelReference<gemini.GemmaConfigSchemaType>;
-  model(
+  createModelRef(
     name: gemini.KnownTtsModels | (gemini.TTSModelName & {}),
     config: gemini.GeminiTtsConfig
   ): ModelReference<gemini.GeminiTtsConfigSchemaType>;
-  model(
+  createModelRef(
     name: gemini.KnownGeminiModels | (gemini.GeminiModelName & {}),
     config?: gemini.GeminiConfig
   ): ModelReference<gemini.GeminiConfigSchemaType>;
-  model(
+  createModelRef(
     name: imagen.KnownModels | (imagen.ImagenModelName & {}),
     config?: imagen.ImagenConfig
   ): ModelReference<imagen.ImagenConfigSchemaType>;
-  model(
+  createModelRef(
     name: veo.KnownModels | (veo.VeoModelName & {}),
     config?: veo.VeoConfig
   ): ModelReference<veo.VeoConfigSchemaType>;
-  model(name: string, config?: any): ModelReference<z.ZodTypeAny>;
+  createModelRef(name: string, config?: any): ModelReference<z.ZodTypeAny>;
 
   embedder(
     name: string,

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

@@ -14,7 +14,12 @@
  * limitations under the License.
  */
 
-import { z, type Document } from 'genkit';
+import {
+  ActionMetadata,
+  embedderActionMetadata,
+  z,
+  type Document,
+} from 'genkit';
 import {
   EmbedderInfo,
   embedderRef,
@@ -29,6 +34,7 @@ import {
   EmbeddingInstance,
   EmbeddingPrediction,
   EmbeddingResult,
+  Model,
   TaskTypeSchema,
   VertexPluginOptions,
   isMultimodalEmbeddingPrediction,
@@ -108,6 +114,13 @@ 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 createModelRef(
   version: string,
@@ -146,6 +159,21 @@ export function createModelRef(
   });
 }
 
+// Takes a full list of models, filters for current Veo models only
+// and returns a modelActionMetadata for each.
+export function listActions(models: Model[]): ActionMetadata[] {
+  return models
+    .filter((m: Model) => isEmbedderModelName(m.name))
+    .map((m: Model) => {
+      const ref = createModelRef(m.name);
+      return embedderActionMetadata({
+        name: ref.name,
+        info: ref.info,
+        configSchema: ref.configSchema,
+      });
+    });
+}
+
 export function defineKnownModels(
   clientOptions: ClientOptions,
   pluginOptions?: VertexPluginOptions

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

@@ -446,7 +446,8 @@ export function defineModel(
       configSchema: ref.configSchema,
       use: middlewares,
     },
-    async (request, { streamingRequested, sendChunk, abortSignal }) => {
+    async (request, options) => {
+      const { streamingRequested, sendChunk, abortSignal } = options || {};
       let clientOpt = { ...clientOptions, signal: abortSignal };
 
       // Make a copy of messages to avoid side-effects

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

@@ -29,6 +29,8 @@ 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';
 
@@ -43,6 +45,9 @@ export { type VeoConfig } from './veo.js';
  * Add Google Cloud Vertex AI to Genkit. Includes Gemini and Imagen models and text embedder.
  */
 function vertexAIPlugin(options?: VertexPluginOptions): GenkitPluginV2 {
+  // Cache for list function - per plugin instance
+  let listCache: any[] | null = null;
+
   return genkitPluginV2({
     name: 'vertexai',
     init: async () => {
@@ -55,28 +60,66 @@ function vertexAIPlugin(options?: VertexPluginOptions): GenkitPluginV2 {
         ...embedder.defineKnownModels(clientOptions, options),
       ];
     },
+    resolve: async (actionType: ActionType, name: string) => {
+      const clientOptions = await getDerivedOptions(options);
+      if (actionType === 'model') {
+        return gemini.defineModel(name, clientOptions);
+      } else if (actionType === 'embedder') {
+        return embedder.defineEmbedder(name, clientOptions);
+      } else if (actionType === 'background-model') {
+        return veo.defineModel(name, clientOptions);
+      }
+      return undefined;
+    },
+    list: async () => {
+      // Return cached result if available
+      if (listCache !== null) {
+        return listCache;
+      }
+
+      try {
+        // Allow undefined options for testing, but try to get derived options
+        const clientOptions = await getDerivedOptions(options);
+        const models = await listModels(clientOptions);
+        const result = [
+          ...imagen.listActions(models),
+          ...lyria.listActions(models),
+          ...veo.listActions(models),
+          ...gemini.listActions(models),
+          // Note: embedders are excluded from list() for VertexAI plugin
+        ];
+
+        // Cache the result
+        listCache = result;
+        return result;
+      } catch (error) {
+        // Handle errors gracefully by returning empty array
+        console.warn('VertexAI plugin: Failed to fetch models list:', error);
+        return [];
+      }
+    },
   });
 }
 
 export type VertexAIPlugin = {
   (pluginOptions?: VertexPluginOptions): GenkitPluginV2;
-  model(
+  createModelRef(
     name: gemini.KnownModels | (gemini.GeminiModelName & {}),
     config?: gemini.GeminiConfig
   ): ModelReference<gemini.GeminiConfigSchemaType>;
-  model(
+  createModelRef(
     name: imagen.KnownModels | (imagen.ImagenModelName & {}),
     config?: imagen.ImagenConfig
   ): ModelReference<imagen.ImagenConfigSchemaType>;
-  model(
+  createModelRef(
     name: lyria.KnownModels | (lyria.LyriaModelName & {}),
     config: lyria.LyriaConfig
   ): ModelReference<lyria.LyriaConfigSchemaType>;
-  model(
+  createModelRef(
     name: veo.KnownModels | (veo.VeoModelName & {}),
     config: veo.VeoConfig
   ): ModelReference<veo.VeoConfigSchemaType>;
-  model(name: string, config?: any): ModelReference<z.ZodTypeAny>;
+  createModelRef(name: string, config?: any): ModelReference<z.ZodTypeAny>;
 
   embedder(
     name: string,
@@ -90,7 +133,7 @@ export type VertexAIPlugin = {
  */
 export const vertexAI = vertexAIPlugin as VertexAIPlugin;
 // provide generic implementation for the model function overloads.
-(vertexAI as any).model = (
+(vertexAI as any).createModelRef = (
   name: string,
   config?: any
 ): ModelReference<z.ZodTypeAny> => {