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';
@@ -59,7 +60,7 @@ async function resolver(
   switch (actionType) {
     case 'model':
       if (veo.isVeoModelName(actionName)) {
-        // no-op
+        // no-op (not gemini)
       } else if (imagen.isImagenModelName(actionName)) {
         return imagen.defineModel(actionName, options);
       } else {
@@ -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) {

js/plugins/google-genai/tests/googleai/veo_test.ts

@@ -19,7 +19,6 @@ import { Operation } from 'genkit';
 import { GenerateRequest } from 'genkit/model';
 import { afterEach, beforeEach, describe, it } from 'node:test';
 import * as sinon from 'sinon';
-
 import { getGenkitClientHeader } from '../../src/common/utils.js';
 import { getGoogleAIUrl } from '../../src/googleai/client.js';
 import { VeoOperation, VeoPredictRequest } from '../../src/googleai/types.js';
@@ -236,8 +235,8 @@ describe('Google AI Veo', () => {
         };
         mockFetchResponse(mockOp);
 
-        const modelAction = captureModelRunner({ apiKey: defaultApiKey });
-        const result = await modelAction.start(request);
+        const { start } = captureModelRunner({ apiKey: defaultApiKey });
+        const result = await start(request);
 
         sinon.assert.calledOnce(fetchStub);
         const fetchArgs = fetchStub.lastCall.args;
@@ -276,12 +275,12 @@ describe('Google AI Veo', () => {
         const apiVersion = 'v1test';
         const baseUrl = 'https://test.example.com';
 
-        const modelAction = captureModelRunner({
+        const { start } = captureModelRunner({
           apiKey: defaultApiKey,
           apiVersion,
           baseUrl,
         });
-        await modelAction.start(request);
+        await start(request);
 
         sinon.assert.calledOnce(fetchStub);
         const fetchArgs = fetchStub.lastCall.args;
@@ -298,8 +297,8 @@ describe('Google AI Veo', () => {
         envStub.value({ GOOGLE_API_KEY: envKey });
         mockFetchResponse({ name: 'operations/start789', done: false });
 
-        const modelAction = captureModelRunner({ apiKey: undefined });
-        await modelAction.start(request);
+        const { start } = captureModelRunner({ apiKey: undefined });
+        await start(request);
 
         sinon.assert.calledOnce(fetchStub);
         const fetchArgs = fetchStub.lastCall.args;
@@ -310,9 +309,9 @@ describe('Google AI Veo', () => {
         const errorBody = { error: { message: 'Invalid argument', code: 400 } };
         mockFetchResponse(errorBody, 400);
 
-        const modelAction = captureModelRunner({ apiKey: defaultApiKey });
+        const { start } = captureModelRunner({ apiKey: defaultApiKey });
         await assert.rejects(
-          modelAction.start(request),
+          start(request),
           /Error fetching from .*models\/veo-test-model:predictLongRunning.* Invalid argument/
         );
       });
@@ -334,8 +333,8 @@ describe('Google AI Veo', () => {
         };
         mockFetchResponse(mockResponse);
 
-        const modelAction = captureModelRunner({ apiKey: defaultApiKey });
-        const result = await modelAction.check(pendingOp);
+        const { check } = captureModelRunner({ apiKey: defaultApiKey });
+        const result = await check(pendingOp);
 
         sinon.assert.calledOnce(fetchStub);
         const fetchArgs = fetchStub.lastCall.args;
@@ -365,12 +364,12 @@ describe('Google AI Veo', () => {
         const apiVersion = 'v1test';
         const baseUrl = 'https://test.example.com';
 
-        const modelAction = captureModelRunner({
+        const { check } = captureModelRunner({
           apiKey: defaultApiKey,
           apiVersion,
           baseUrl,
         });
-        await modelAction.check(pendingOp);
+        await check(pendingOp);
 
         sinon.assert.calledOnce(fetchStub);
         const fetchArgs = fetchStub.lastCall.args;
@@ -385,9 +384,9 @@ describe('Google AI Veo', () => {
         const errorBody = { error: { message: 'Not found', code: 404 } };
         mockFetchResponse(errorBody, 404);
 
-        const modelAction = captureModelRunner({ apiKey: defaultApiKey });
+        const { check } = captureModelRunner({ apiKey: defaultApiKey });
         await assert.rejects(
-          modelAction.check(pendingOp),
+          check(pendingOp),
           /Error fetching from .*operations\/check123.* Not found/
         );
       });

js/plugins/google-genai/tests/vertexai/veo_test.ts

@@ -20,7 +20,6 @@ import { GenerateRequest } from 'genkit/model';
 import { GoogleAuth } from 'google-auth-library';
 import { afterEach, beforeEach, describe, it } from 'node:test';
 import * as sinon from 'sinon';
-
 import { getGenkitClientHeader } from '../../src/common/utils.js';
 import { getVertexAIUrl } from '../../src/vertexai/client.js';
 import {
@@ -134,8 +133,8 @@ describe('Vertex AI Veo', () => {
         };
         mockFetchResponse(mockOp);
 
-        const modelAction = captureModelRunner(defaultRegionalClientOptions);
-        const result = await modelAction.start(request);
+        const { start } = captureModelRunner(defaultRegionalClientOptions);
+        const result = await start(request);
 
         sinon.assert.calledOnce(fetchStub);
         const fetchArgs = fetchStub.lastCall.args;
@@ -168,9 +167,9 @@ describe('Vertex AI Veo', () => {
         const errorBody = { error: { message: 'Invalid arg', code: 400 } };
         mockFetchResponse(errorBody, 400);
 
-        const modelAction = captureModelRunner(defaultRegionalClientOptions);
+        const { start } = captureModelRunner(defaultRegionalClientOptions);
         await assert.rejects(
-          modelAction.start(request),
+          start(request),
           /Error fetching from .*predictLongRunning.* Invalid arg/
         );
       });
@@ -184,8 +183,8 @@ describe('Vertex AI Veo', () => {
           ...defaultRegionalClientOptions,
           signal: abortSignal,
         };
-        const modelAction = captureModelRunner(clientOptionsWithSignal);
-        await modelAction.start(request);
+        const { start } = captureModelRunner(clientOptionsWithSignal);
+        await start(request);
 
         sinon.assert.calledOnce(fetchStub);
         const fetchOptions = fetchStub.lastCall.args[1];
@@ -220,8 +219,8 @@ describe('Vertex AI Veo', () => {
         };
         mockFetchResponse(mockResponse);
 
-        const modelAction = captureModelRunner(defaultRegionalClientOptions);
-        const result = await modelAction.check(pendingOp);
+        const { check } = captureModelRunner(defaultRegionalClientOptions);
+        const result = await check(pendingOp);
 
         sinon.assert.calledOnce(fetchStub);
         const fetchArgs = fetchStub.lastCall.args;
@@ -252,9 +251,9 @@ describe('Vertex AI Veo', () => {
         const errorBody = { error: { message: 'Not found', code: 404 } };
         mockFetchResponse(errorBody, 404);
 
-        const modelAction = captureModelRunner(defaultRegionalClientOptions);
+        const { check } = captureModelRunner(defaultRegionalClientOptions);
         await assert.rejects(
-          modelAction.check(pendingOp),
+          check(pendingOp),
           /Error fetching from .*fetchPredictOperation.* Not found/
         );
       });