refactor(js/plugins/vertexai): extract reranker to subplugin

js/plugins/vertexai/package.json

@@ -66,9 +66,22 @@
       "import": "./lib/index.mjs",
       "types": "./lib/index.d.ts",
       "default": "./lib/index.js"
+    },
+    "./reranker": {
+      "require": "./lib/reranker/index.js",
+      "import": "./lib/reranker/index.mjs",
+      "types": "./lib/reranker/index.d.ts",
+      "default": "./lib/reranker/index.js"
     }
   },
   "typesVersions": {
-    "*": {}
+    "*": {
+      ".": [
+        "./lib/index.d.ts"
+      ],
+      "./reranker": [
+        "./lib/reranker/index.d.ts"
+      ]
+    }
   }
 }

js/plugins/vertexai/src/index.ts

@@ -70,7 +70,6 @@ import {
   modelGardenOpenaiCompatibleModel,
   SUPPORTED_OPENAI_FORMAT_MODELS,
 } from './model_garden.js';
-import { vertexAiRerankers, VertexRerankerConfig } from './reranker.js';
 import {
   VectorSearchOptions,
   vertexAiIndexers,
@@ -138,8 +137,6 @@ export interface PluginOptions {
   };
   /** Configure Vertex AI vector search index options */
   vectorSearchOptions?: VectorSearchOptions<z.ZodTypeAny, any, any>[];
-  /** Configure reranker options */
-  rerankOptions?: VertexRerankerConfig[];
 }
 
 /**
@@ -246,21 +243,12 @@ export const vertexAI: Plugin<[PluginOptions] | []> = genkitPlugin(
       });
     }
 
-    const rerankOptions = {
-      pluginOptions: options,
-      authClient,
-      projectId,
-    };
-
-    const rerankers = await vertexAiRerankers(rerankOptions);
-
     return {
       models,
       embedders,
       evaluators: vertexEvaluators(authClient, metrics, projectId, location),
       retrievers,
       indexers,
-      rerankers,
     };
   }
 );

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

@@ -0,0 +1,63 @@
+// Copyright 2024 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import { genkitPlugin, Plugin } from 'genkit';
+import { authenticate } from '../common/auth.js';
+import { confError, DEFAULT_LOCATION } from '../common/global.js';
+import { BasePluginOptions } from '../common/types.js';
+import { vertexAiRerankers, VertexRerankerConfig } from './reranker.js';
+
+export interface PluginOptions extends BasePluginOptions {
+  /** Configure reranker options */
+  options: VertexRerankerConfig[];
+}
+
+/**
+ *  Plugin for Vertex AI Reranker
+ */
+export const vertexAIReranker: Plugin<[PluginOptions]> = genkitPlugin(
+  'vertexaiReranker',
+  async (options: PluginOptions) => {
+    // Authenticate with Google Cloud
+    const authOptions = options?.googleAuth;
+    const authClient = authenticate(authOptions);
+
+    const projectId = options?.projectId || (await authClient.getProjectId());
+    const location = options?.location || DEFAULT_LOCATION;
+
+    if (!location) {
+      throw confError('location', 'GCLOUD_LOCATION');
+    }
+    if (!projectId) {
+      throw confError('project', 'GCLOUD_PROJECT');
+    }
+
+    const rerankOptions = {
+      pluginOptions: options,
+      authClient,
+      projectId,
+    };
+
+    const rerankers = await vertexAiRerankers(rerankOptions);
+
+    return {
+      rerankers,
+      // We add this here just to satisfy the types. The plugin return types will soon be removed, there
+      // is no need to add rerankers to that return type.
+      models: [],
+    };
+  }
+);
+
+export default vertexAIReranker;

js/plugins/vertexai/src/reranker.ts → js/plugins/vertexai/src/reranker/reranker.ts

@@ -22,7 +22,7 @@ import {
   rerankerRef,
 } from 'genkit/reranker';
 import { GoogleAuth } from 'google-auth-library';
-import { PluginOptions } from '.';
+import { PluginOptions } from './';
 
 const DEFAULT_MODEL = 'semantic-ranker-512@latest';
 
@@ -63,7 +63,7 @@ export interface VertexRerankPluginOptions {
 
 export interface VertexRerankOptions {
   authClient: GoogleAuth;
-  pluginOptions?: PluginOptions;
+  pluginOptions: PluginOptions;
 }
 
 /**
@@ -82,11 +82,11 @@ export async function vertexAiRerankers(
     return [];
   }
   const pluginOptions = params.pluginOptions;
-  if (!params.pluginOptions.rerankOptions) {
+  if (!params.pluginOptions.options) {
     return [];
   }
 
-  const rerankOptions = params.pluginOptions.rerankOptions;
+  const rerankOptions = params.pluginOptions.options;
   const rerankers: RerankerAction<z.ZodTypeAny>[] = [];
 
   if (!rerankOptions || rerankOptions.length === 0) {
@@ -99,7 +99,7 @@ export async function vertexAiRerankers(
   for (const rerankOption of rerankOptions) {
     const reranker = defineReranker(
       {
-        name: `vertexai/${rerankOption.name || rerankOption.model}`,
+        name: `vertexaiReranker/${rerankOption.name || rerankOption.model}`,
         configSchema: VertexAIRerankerOptionsSchema.optional(),
       },
       async (query, documents, _options) => {
@@ -154,7 +154,7 @@ export const vertexAiRerankerRef = (params: {
   displayName?: string;
 }) => {
   return rerankerRef({
-    name: `vertexai/${name}`,
+    name: `vertexaiReranker/${name}`,
     info: {
       label: params.displayName ?? `Vertex AI Reranker`,
     },

js/testapps/vertexai-reranker/src/index.ts

@@ -18,19 +18,19 @@
 
 import { Document, genkit, z } from 'genkit';
 // important imports for this sample:
-import { vertexAI } from '@genkit-ai/vertexai';
+import { vertexAIReranker } from '@genkit-ai/vertexai/reranker';
 import { LOCATION, PROJECT_ID } from './config';
 
 // Configure Genkit with Vertex AI plugin
 const ai = genkit({
   plugins: [
-    vertexAI({
+    vertexAIReranker({
       projectId: PROJECT_ID,
       location: LOCATION,
       googleAuth: {
         scopes: ['https://www.googleapis.com/auth/cloud-platform'],
       },
-      rerankOptions: [
+      options: [
         {
           model: 'vertexai/semantic-ranker-512',
         },

js/testapps/vertexai-reranker/tsconfig.json

@@ -2,7 +2,7 @@
   "compileOnSave": true,
   "include": ["src"],
   "compilerOptions": {
-    "module": "commonjs",
+    "module": "NodeNext",
     "noImplicitReturns": true,
     "outDir": "lib",
     "sourceMap": true,