chore(plugins/vertexai): address feedback and format

Author: CorieW

js/plugins/vertexai/src/embedder.ts

@@ -20,10 +20,10 @@ import {
   type EmbedderAction,
   type EmbedderReference,
 } from 'genkit/embedder';
+import { embedder } from 'genkit/plugin';
 import type { GoogleAuth } from 'google-auth-library';
 import type { PluginOptions } from './common/types.js';
 import { predictModel, type PredictClient } from './predict.js';
-import { embedder } from 'genkit/plugin';
 
 export const TaskTypeSchema = z.enum([
   'RETRIEVAL_DOCUMENT',
@@ -304,12 +304,15 @@ export function defineVertexAIEmbedder(
       configSchema: embedderRef.configSchema,
       info: embedderRef.info!,
     },
-    async (input, options) => {
+    async (request, options) => {
       const predictClient = predictClientFactory(embedderRef.config);
       const response = await predictClient(
-        input.input.map((doc: Document) => {
+        request.input.map((doc: Document) => {
           let instance: EmbeddingInstance;
-          if (isMultiModal(embedderRef) && checkValidDocument(embedderRef, doc)) {
+          if (
+            isMultiModal(embedderRef) &&
+            checkValidDocument(embedderRef, doc)
+          ) {
             instance = {};
             if (doc.text) {
               instance.text = doc.text;

js/plugins/vertexai/src/evaluation/evaluator_factory.ts

@@ -16,11 +16,11 @@
 
 import { type Action, type z } from 'genkit';
 import type { BaseEvalDataPoint, Score } from 'genkit/evaluator';
+import { evaluator } from 'genkit/plugin';
 import { runInNewSpan } from 'genkit/tracing';
 import type { GoogleAuth } from 'google-auth-library';
 import { getGenkitClientHeader } from '../common/index.js';
 import type { VertexAIEvaluationMetricType } from './evaluation.js';
-import { evaluator } from 'genkit/plugin';
 
 export class EvaluatorFactory {
   constructor(

js/plugins/vertexai/src/evaluation/index.ts

@@ -28,7 +28,8 @@ export function vertexAIEvaluation(options: PluginOptions): GenkitPluginV2 {
   return genkitPluginV2({
     name: 'vertexAIEvaluation',
     init: async () => {
-      const { projectId, location, authClient } = await getDerivedParams(options);
+      const { projectId, location, authClient } =
+        await getDerivedParams(options);
 
       return vertexEvaluators(authClient, options.metrics, projectId, location);
     },

js/plugins/vertexai/src/gemini.ts

@@ -56,13 +56,13 @@ import {
   downloadRequestMedia,
   simulateSystemPrompt,
 } from 'genkit/model/middleware';
+import { model } from 'genkit/plugin';
 import { runInNewSpan } from 'genkit/tracing';
 import { GoogleAuth } from 'google-auth-library';
 import { getGenkitClientHeader } from './common/index.js';
 import type { PluginOptions } from './common/types.js';
 import { handleCacheIfNeeded } from './context-caching/index.js';
 import { extractCacheConfig } from './context-caching/utils.js';
-import { model } from 'genkit/plugin';
 
 // Extra type guard to keep the compiler happy and avoid a cast to any. The
 // legacy Gemini SDK is no longer maintained, and doesn't have updated types.

js/plugins/vertexai/src/imagen.ts

@@ -14,21 +14,21 @@
  * limitations under the License.
  */
 
-import { z, type Genkit } from 'genkit';
+import { z } from 'genkit';
 import {
   GenerationCommonConfigSchema,
-  getBasicUsageStats,
   modelRef as createModelRef,
+  getBasicUsageStats,
   type CandidateData,
   type GenerateRequest,
   type ModelAction,
   type ModelInfo,
   type ModelReference,
 } from 'genkit/model';
+import { model } from 'genkit/plugin';
 import type { GoogleAuth } from 'google-auth-library';
 import type { PluginOptions } from './common/types.js';
 import { predictModel, type PredictClient } from './predict.js';
-import { model } from 'genkit/plugin';
 
 /**
  * See https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/imagen-api.

js/plugins/vertexai/src/index.ts

@@ -28,7 +28,11 @@ import {
   type ModelReference,
   type z,
 } from 'genkit';
-import { genkitPluginV2, ResolvableAction, type GenkitPluginV2 } from 'genkit/plugin';
+import {
+  ResolvableAction,
+  genkitPluginV2,
+  type GenkitPluginV2,
+} from 'genkit/plugin';
 import type { ActionType } from 'genkit/registry';
 import { getDerivedParams } from './common/index.js';
 import type { PluginOptions } from './common/types.js';
@@ -116,20 +120,22 @@ async function initializer(options?: PluginOptions) {
 
   const actions: ResolvableAction[] = [];
 
-  Object.keys(SUPPORTED_IMAGEN_MODELS).map((name) =>
-    actions.push(defineImagenModel(name, authClient, { projectId, location }))
-  );
-  Object.keys(SUPPORTED_GEMINI_MODELS).map((name) =>
-    actions.push(defineGeminiKnownModel(
-      name,
-      vertexClientFactory,
-      {
-        projectId,
-        location,
-      },
-      options?.experimental_debugTraces
-    ))
-  );
+  for (const name of Object.keys(SUPPORTED_IMAGEN_MODELS)) {
+    actions.push(defineImagenModel(name, authClient, { projectId, location }));
+  }
+  for (const name of Object.keys(SUPPORTED_GEMINI_MODELS)) {
+    actions.push(
+      defineGeminiKnownModel(
+        name,
+        vertexClientFactory,
+        {
+          projectId,
+          location,
+        },
+        options?.experimental_debugTraces
+      )
+    );
+  }
   if (options?.models) {
     for (const modelOrRef of options?.models) {
       const modelName =
@@ -139,22 +145,26 @@ async function initializer(options?: PluginOptions) {
             modelOrRef.name.split('/')[1];
       const modelRef =
         typeof modelOrRef === 'string' ? gemini(modelOrRef) : modelOrRef;
-      actions.push(defineGeminiModel({
-        modelName: modelRef.name,
-        version: modelName,
-        modelInfo: modelRef.info,
-        vertexClientFactory,
-        options: {
-          projectId,
-          location,
-        },
-        debugTraces: options.experimental_debugTraces,
-      }));
+      actions.push(
+        defineGeminiModel({
+          modelName: modelRef.name,
+          version: modelName,
+          modelInfo: modelRef.info,
+          vertexClientFactory,
+          options: {
+            projectId,
+            location,
+          },
+          debugTraces: options.experimental_debugTraces,
+        })
+      );
     }
   }
 
   Object.keys(SUPPORTED_EMBEDDER_MODELS).map((name) =>
-    actions.push(defineVertexAIEmbedder(name, authClient, { projectId, location }))
+    actions.push(
+      defineVertexAIEmbedder(name, authClient, { projectId, location })
+    )
   );
 
   return actions;
@@ -185,8 +195,7 @@ async function resolveModel(
     await getDerivedParams(options);
 
   if (actionName.startsWith('imagen')) {
-    defineImagenModel(actionName, authClient, { projectId, location });
-    return;
+    return defineImagenModel(actionName, authClient, { projectId, location });
   }
 
   const modelRef = gemini(actionName);
@@ -209,7 +218,10 @@ async function resolveEmbedder(
 ): Promise<ResolvableAction | undefined> {
   const { projectId, location, authClient } = await getDerivedParams(options);
 
-  return defineVertexAIEmbedder(actionName, authClient, { projectId, location })
+  return defineVertexAIEmbedder(actionName, authClient, {
+    projectId,
+    location,
+  });
 }
 
 // Vertex AI list models still returns these and the API does not indicate in any way
@@ -268,14 +280,14 @@ function vertexAIPlugin(options?: PluginOptions): GenkitPluginV2 {
   let listActionsCache;
   return genkitPluginV2({
     name: 'vertexai',
-    init: async () => initializer(),
+    init: async () => initializer(options),
     resolve: async (actionType: ActionType, actionName: string) =>
-      await resolver(actionType, actionName),
+      await resolver(actionType, actionName, options),
     list: async () => {
       if (listActionsCache) return listActionsCache;
       listActionsCache = await listActions(options);
       return listActionsCache;
-    }
+    },
   });
 }
 

js/plugins/vertexai/src/modelgarden/anthropic.ts

@@ -44,8 +44,8 @@ import {
   modelRef,
   type ModelAction,
 } from 'genkit/model';
-import { getGenkitClientHeader } from '../common/index.js';
 import { model } from 'genkit/plugin';
+import { getGenkitClientHeader } from '../common/index.js';
 
 export const AnthropicConfigSchema = GenerationCommonConfigSchema.extend({
   location: z.string().optional(),

js/plugins/vertexai/src/modelgarden/index.ts

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-import type { Action, Genkit } from 'genkit';
+import type { Action } from 'genkit';
 import { genkitPluginV2, type GenkitPluginV2 } from 'genkit/plugin';
 import { getDerivedParams } from '../common/index.js';
 import { SUPPORTED_ANTHROPIC_MODELS, anthropicModel } from './anthropic.js';
@@ -32,7 +32,8 @@ export function vertexAIModelGarden(options: PluginOptions): GenkitPluginV2 {
   return genkitPluginV2({
     name: 'vertexAIModelGarden',
     init: async () => {
-      const { projectId, location, authClient } = await getDerivedParams(options);
+      const { projectId, location, authClient } =
+        await getDerivedParams(options);
 
       const actions: Action[] = [];
 
@@ -55,13 +56,15 @@ export function vertexAIModelGarden(options: PluginOptions): GenkitPluginV2 {
           ([_, value]) => value.name === m.name
         );
         if (openaiModel) {
-          actions.push(modelGardenOpenaiCompatibleModel(
-            openaiModel[0],
-            projectId,
-            location,
-            authClient,
-            options.openAiBaseUrlTemplate
-          ));
+          actions.push(
+            modelGardenOpenaiCompatibleModel(
+              openaiModel[0],
+              projectId,
+              location,
+              authClient,
+              options.openAiBaseUrlTemplate
+            )
+          );
           return;
         }
         throw new Error(`Unsupported model garden model: ${m.name}`);

js/plugins/vertexai/src/modelgarden/mistral.ts

@@ -45,8 +45,8 @@ import {
   modelRef,
   type ModelAction,
 } from 'genkit/model';
-import { getGenkitClientHeader } from '../common/index.js';
 import { model } from 'genkit/plugin';
+import { getGenkitClientHeader } from '../common/index.js';
 
 /**
  * See https://docs.mistral.ai/api/#tag/chat/operation/chat_completion_v1_chat_completions_post

js/plugins/vertexai/src/rerankers/index.ts

@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 
-import type { Genkit } from 'genkit';
 import { genkitPluginV2, type GenkitPluginV2 } from 'genkit/plugin';
 import type { CommonPluginOptions } from '../common/types.js';
 import type { RerankerOptions } from './types.js';
@@ -32,7 +31,8 @@ export function vertexAIRerankers(options: PluginOptions): GenkitPluginV2 {
   return genkitPluginV2({
     name: 'vertexAIRerankers',
     init: async () => {
-      const { projectId, location, authClient } = await getDerivedParams(options);
+      const { projectId, location, authClient } =
+        await getDerivedParams(options);
 
       return vertexAiRerankers({
         projectId,