Split vector search into sub-plugin of vertexai plugin

js/plugins/vertexai/package.json

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

js/plugins/vertexai/src/index.ts

@@ -15,9 +15,8 @@
  */
 
 import { VertexAI } from '@google-cloud/vertexai';
-import { genkitPlugin, Plugin, z } from 'genkit';
+import { genkitPlugin, Plugin } from 'genkit';
 import { GenerateRequest, ModelReference } from 'genkit/model';
-import { IndexerAction, RetrieverAction } from 'genkit/retriever';
 import { GoogleAuthOptions } from 'google-auth-library';
 import {
   anthropicModel,
@@ -71,25 +70,6 @@ import {
   SUPPORTED_OPENAI_FORMAT_MODELS,
 } from './model_garden.js';
 import { vertexAiRerankers, VertexRerankerConfig } from './reranker.js';
-import {
-  VectorSearchOptions,
-  vertexAiIndexers,
-  vertexAiRetrievers,
-} from './vector-search';
-export {
-  DocumentIndexer,
-  DocumentRetriever,
-  getBigQueryDocumentIndexer,
-  getBigQueryDocumentRetriever,
-  getFirestoreDocumentIndexer,
-  getFirestoreDocumentRetriever,
-  Neighbor,
-  VectorSearchOptions,
-  vertexAiIndexerRef,
-  vertexAiIndexers,
-  vertexAiRetrieverRef,
-  vertexAiRetrievers,
-} from './vector-search';
 export {
   claude35Sonnet,
   claude3Haiku,
@@ -136,8 +116,6 @@ export interface PluginOptions {
     models: ModelReference<any>[];
     openAiBaseUrlTemplate?: string;
   };
-  /** Configure Vertex AI vector search index options */
-  vectorSearchOptions?: VectorSearchOptions<z.ZodTypeAny, any, any>[];
   /** Configure reranker options */
   rerankOptions?: VertexRerankerConfig[];
 }
@@ -224,28 +202,6 @@ export const vertexAI: Plugin<[PluginOptions] | []> = genkitPlugin(
       textEmbeddingGeckoEmbedder(name, authClient, { projectId, location })
     );
 
-    let indexers: IndexerAction<z.ZodTypeAny>[] = [];
-    let retrievers: RetrieverAction<z.ZodTypeAny>[] = [];
-
-    if (
-      options?.vectorSearchOptions &&
-      options.vectorSearchOptions.length > 0
-    ) {
-      const defaultEmbedder = embedders[0];
-
-      indexers = vertexAiIndexers({
-        pluginOptions: options,
-        authClient,
-        defaultEmbedder,
-      });
-
-      retrievers = vertexAiRetrievers({
-        pluginOptions: options,
-        authClient,
-        defaultEmbedder,
-      });
-    }
-
     const rerankOptions = {
       pluginOptions: options,
       authClient,
@@ -258,8 +214,6 @@ export const vertexAI: Plugin<[PluginOptions] | []> = genkitPlugin(
       models,
       embedders,
       evaluators: vertexEvaluators(authClient, metrics, projectId, location),
-      retrievers,
-      indexers,
       rerankers,
     };
   }

js/plugins/vertexai/src/vector-search/index.ts

@@ -14,17 +14,26 @@
  * limitations under the License.
  */
 
-export {
+import { genkitPlugin, Plugin, z } from 'genkit';
+import { IndexerAction, RetrieverAction } from 'genkit/retriever';
+import { authenticate } from '../common/auth.js';
+import { confError, DEFAULT_LOCATION } from '../common/global.js';
+import { BasePluginOptions } from '../common/types';
+import {
+  SUPPORTED_EMBEDDER_MODELS,
+  textEmbeddingGeckoEmbedder,
+} from '../embedder.js';
+import {
   getBigQueryDocumentIndexer,
   getBigQueryDocumentRetriever,
 } from './bigquery';
-export {
+import {
   getFirestoreDocumentIndexer,
   getFirestoreDocumentRetriever,
 } from './firestore';
-export { vertexAiIndexerRef, vertexAiIndexers } from './indexers';
-export { vertexAiRetrieverRef, vertexAiRetrievers } from './retrievers';
-export {
+import { vertexAiIndexerRef, vertexAiIndexers } from './indexers';
+import { vertexAiRetrieverRef, vertexAiRetrievers } from './retrievers';
+import {
   DocumentIndexer,
   DocumentRetriever,
   Neighbor,
@@ -34,3 +43,79 @@ export {
   VertexAIVectorRetrieverOptions,
   VertexAIVectorRetrieverOptionsSchema,
 } from './types';
+
+export {
+  DocumentIndexer,
+  DocumentRetriever,
+  getBigQueryDocumentIndexer,
+  getBigQueryDocumentRetriever,
+  getFirestoreDocumentIndexer,
+  getFirestoreDocumentRetriever,
+  Neighbor,
+  VectorSearchOptions,
+  vertexAiIndexerRef,
+  vertexAiIndexers,
+  vertexAiRetrieverRef,
+  vertexAiRetrievers,
+  VertexAIVectorIndexerOptions,
+  VertexAIVectorIndexerOptionsSchema,
+  VertexAIVectorRetrieverOptions,
+  VertexAIVectorRetrieverOptionsSchema,
+};
+
+export interface PluginOptions extends BasePluginOptions {
+  /** Configure Vertex AI vector search index options */
+  options?: VectorSearchOptions<z.ZodTypeAny, any, any>[];
+}
+
+/**
+ *  Plugin for Vertex AI Vector Search
+ */
+export const vertexAIVectorSearch: Plugin<[PluginOptions] | []> = genkitPlugin(
+  'vertexaiVectorSearch',
+  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 embedders = Object.keys(SUPPORTED_EMBEDDER_MODELS).map((name) =>
+      textEmbeddingGeckoEmbedder(name, authClient, { projectId, location })
+    );
+
+    let indexers: IndexerAction<z.ZodTypeAny>[] = [];
+    let retrievers: RetrieverAction<z.ZodTypeAny>[] = [];
+
+    if (options?.options && options.options.length > 0) {
+      const defaultEmbedder = embedders[0];
+
+      indexers = vertexAiIndexers({
+        pluginOptions: options,
+        authClient,
+        defaultEmbedder,
+      });
+
+      retrievers = vertexAiRetrievers({
+        pluginOptions: options,
+        authClient,
+        defaultEmbedder,
+      });
+    }
+
+    return {
+      indexers,
+      retrievers,
+    };
+  }
+);
+
+export default vertexAIVectorSearch;

js/plugins/vertexai/src/vector-search/indexers.ts

@@ -37,7 +37,7 @@ export const vertexAiIndexerRef = (params: {
   displayName?: string;
 }) => {
   return indexerRef({
-    name: `vertexai/${params.indexId}`,
+    name: `vertexaiVectorSearch/${params.indexId}`,
     info: {
       label: params.displayName ?? `Vertex AI - ${params.indexId}`,
     },
@@ -57,7 +57,7 @@ export const vertexAiIndexerRef = (params: {
 export function vertexAiIndexers<EmbedderCustomOptions extends z.ZodTypeAny>(
   params: VertexVectorSearchOptions<EmbedderCustomOptions>
 ): IndexerAction<z.ZodTypeAny>[] {
-  const vectorSearchOptions = params.pluginOptions.vectorSearchOptions;
+  const vectorSearchOptions = params.pluginOptions.options;
   const defaultEmbedder = params.defaultEmbedder;
   const indexers: IndexerAction<z.ZodTypeAny>[] = [];
 
@@ -72,7 +72,7 @@ export function vertexAiIndexers<EmbedderCustomOptions extends z.ZodTypeAny>(
 
     const indexer = defineIndexer(
       {
-        name: `vertexai/${indexId}`,
+        name: `vertexaiVectorSearch/${indexId}`,
         configSchema: VertexAIVectorIndexerOptionsSchema.optional(),
       },
       async (docs, options) => {

js/plugins/vertexai/src/vector-search/retrievers.ts

@@ -37,7 +37,7 @@ const DEFAULT_K = 10;
 export function vertexAiRetrievers<EmbedderCustomOptions extends z.ZodTypeAny>(
   params: VertexVectorSearchOptions<EmbedderCustomOptions>
 ): RetrieverAction<z.ZodTypeAny>[] {
-  const vectorSearchOptions = params.pluginOptions.vectorSearchOptions;
+  const vectorSearchOptions = params.pluginOptions.options;
   const defaultEmbedder = params.defaultEmbedder;
 
   const retrievers: RetrieverAction<z.ZodTypeAny>[] = [];
@@ -53,7 +53,7 @@ export function vertexAiRetrievers<EmbedderCustomOptions extends z.ZodTypeAny>(
 
     const retriever = defineRetriever(
       {
-        name: `vertexai/${indexId}`,
+        name: `vertexaiVectorSearch/${indexId}`,
         configSchema: VertexAIVectorRetrieverOptionsSchema.optional(),
       },
       async (content, options) => {
@@ -127,7 +127,7 @@ export const vertexAiRetrieverRef = (params: {
   displayName?: string;
 }) => {
   return retrieverRef({
-    name: `vertexai/${params.indexId}`,
+    name: `vertexaiVectorSearch/${params.indexId}`,
     info: {
       label: params.displayName ?? `ertex AI - ${params.indexId}`,
     },

js/plugins/vertexai/src/vector-search/types.ts

@@ -19,7 +19,7 @@ import { z } from 'genkit';
 import { EmbedderArgument } from 'genkit/embedder';
 import { CommonRetrieverOptionsSchema, Document } from 'genkit/retriever';
 import { GoogleAuth } from 'google-auth-library';
-import { PluginOptions } from '..';
+import { PluginOptions } from '../vector-search/';
 
 // This internal interface will be passed to the vertexIndexers and vertexRetrievers functions
 export interface VertexVectorSearchOptions<

js/plugins/vertexai/tests/vector-search/bigquery_test.ts

@@ -18,7 +18,7 @@ import { BigQuery } from '@google-cloud/bigquery';
 import { Document } from 'genkit/retriever';
 import assert from 'node:assert';
 import { describe, it } from 'node:test';
-import { getBigQueryDocumentRetriever } from '../../src';
+import { getBigQueryDocumentRetriever } from '../../src/vector-search';
 
 class MockBigQuery {
   query: Function;

js/testapps/vertexai-vector-search-bigquery/src/index.ts

@@ -21,12 +21,12 @@ import { Document, genkit, z } from 'genkit';
 import {
   getBigQueryDocumentIndexer,
   getBigQueryDocumentRetriever,
-  vertexAI,
+  vertexAIVectorSearch,
   vertexAiIndexerRef,
   vertexAiRetrieverRef,
   type DocumentIndexer,
   type DocumentRetriever,
-} from '@genkit-ai/vertexai';
+} from '@genkit-ai/vertexai/vectorsearch';
 
 // // Environment variables set with dotenv for simplicity of sample
 import {
@@ -75,13 +75,13 @@ const bigQueryDocumentIndexer: DocumentIndexer = getBigQueryDocumentIndexer(
 // Configure Genkit with Vertex AI plugin
 const ai = genkit({
   plugins: [
-    vertexAI({
+    vertexAIVectorSearch({
       projectId: PROJECT_ID,
       location: LOCATION,
       googleAuth: {
         scopes: ['https://www.googleapis.com/auth/cloud-platform'],
       },
-      vectorSearchOptions: [
+      options: [
         {
           publicDomainName: VECTOR_SEARCH_PUBLIC_DOMAIN_NAME,
           indexEndpointId: VECTOR_SEARCH_INDEX_ENDPOINT_ID,

js/testapps/vertexai-vector-search-custom/src/index.ts

@@ -19,13 +19,13 @@
 import { Document, genkit, z } from 'genkit';
 // important imports for this sample:
 import {
-  vertexAI,
+  vertexAIVectorSearch,
   vertexAiIndexerRef,
   vertexAiRetrieverRef,
   type DocumentIndexer,
   type DocumentRetriever,
   type Neighbor,
-} from '@genkit-ai/vertexai';
+} from '@genkit-ai/vertexai/vectorsearch';
 
 // // Environment variables set with dotenv for simplicity of sample
 import {
@@ -142,13 +142,13 @@ const localDocumentRetriever: DocumentRetriever = async (
 // Configure Genkit with Vertex AI plugin
 const ai = genkit({
   plugins: [
-    vertexAI({
+    vertexAIVectorSearch({
       projectId: PROJECT_ID,
       location: LOCATION,
       googleAuth: {
         scopes: ['https://www.googleapis.com/auth/cloud-platform'],
       },
-      vectorSearchOptions: [
+      options: [
         {
           publicDomainName: VECTOR_SEARCH_PUBLIC_DOMAIN_NAME,
           indexEndpointId: VECTOR_SEARCH_INDEX_ENDPOINT_ID,

js/testapps/vertexai-vector-search-firestore/src/index.ts

@@ -24,10 +24,10 @@ import {
   DocumentRetriever,
   getFirestoreDocumentIndexer,
   getFirestoreDocumentRetriever,
-  vertexAI,
+  vertexAIVectorSearch,
   vertexAiIndexerRef,
   vertexAiRetrieverRef,
-} from '@genkit-ai/vertexai';
+} from '@genkit-ai/vertexai/vectorsearch';
 
 // // Environment variables set with dotenv for simplicity of sample
 import { getFirestore } from 'firebase-admin/firestore';
@@ -74,13 +74,13 @@ const firestoreDocumentIndexer: DocumentIndexer = getFirestoreDocumentIndexer(
 // Configure Genkit with Vertex AI plugin
 const ai = genkit({
   plugins: [
-    vertexAI({
+    vertexAIVectorSearch({
       projectId: PROJECT_ID,
       location: LOCATION,
       googleAuth: {
         scopes: ['https://www.googleapis.com/auth/cloud-platform'],
       },
-      vectorSearchOptions: [
+      options: [
         {
           publicDomainName: VECTOR_SEARCH_PUBLIC_DOMAIN_NAME,
           indexEndpointId: VECTOR_SEARCH_INDEX_ENDPOINT_ID,