From 9ccbd8d533eaa44c6ced1bd87c3a7e672e5d419a Mon Sep 17 00:00:00 2001 From: Jacob Cable Date: Mon, 11 Nov 2024 11:44:51 +0000 Subject: [PATCH] refactor(js/plugins/vertexai): extract vectorsearch --- js/plugins/vertexai/package.json | 9 +++++ js/plugins/vertexai/src/index.ts | 36 +------------------ .../{vector-search => vectorsearch}/index.ts | 12 +------ .../{vector-search => vectorsearch}/types.ts | 3 +- .../vector_search/bigquery.ts | 0 .../vector_search/firestore.ts | 0 .../vector_search/index.ts | 0 .../vector_search/indexers.ts | 12 +++++-- .../vector_search/query_public_endpoint.ts | 0 .../vector_search/retrievers.ts | 11 +++++- .../vector_search/types.ts | 6 ++-- .../vector_search/upsert_datapoints.ts | 0 .../vector_search/utils.ts | 0 js/testapps/model-tester/src/index.ts | 10 +++++- js/testapps/rag/src/genkit.ts | 9 +++-- .../src/index.ts | 16 +++++---- .../src/index.ts | 10 ++++-- .../src/index.ts | 12 +++++-- 18 files changed, 78 insertions(+), 68 deletions(-) rename js/plugins/vertexai/src/{vector-search => vectorsearch}/index.ts (85%) rename js/plugins/vertexai/src/{vector-search => vectorsearch}/types.ts (92%) rename js/plugins/vertexai/src/{vector-search => vectorsearch}/vector_search/bigquery.ts (100%) rename js/plugins/vertexai/src/{vector-search => vectorsearch}/vector_search/firestore.ts (100%) rename js/plugins/vertexai/src/{vector-search => vectorsearch}/vector_search/index.ts (100%) rename js/plugins/vertexai/src/{vector-search => vectorsearch}/vector_search/indexers.ts (93%) rename js/plugins/vertexai/src/{vector-search => vectorsearch}/vector_search/query_public_endpoint.ts (100%) rename js/plugins/vertexai/src/{vector-search => vectorsearch}/vector_search/retrievers.ts (93%) rename js/plugins/vertexai/src/{vector-search => vectorsearch}/vector_search/types.ts (97%) rename js/plugins/vertexai/src/{vector-search => vectorsearch}/vector_search/upsert_datapoints.ts (100%) rename js/plugins/vertexai/src/{vector-search => vectorsearch}/vector_search/utils.ts (100%) diff --git a/js/plugins/vertexai/package.json b/js/plugins/vertexai/package.json index 401b6c40e..2b4571b97 100644 --- a/js/plugins/vertexai/package.json +++ b/js/plugins/vertexai/package.json @@ -84,6 +84,12 @@ "import": "./lib/modelgarden/index.mjs", "types": "./lib/modelgarden/index.d.ts", "default": "./lib/modelgarden/index.js" + }, + "./vectorsearch": { + "require": "./lib/vectorsearch/index.js", + "import": "./lib/vectorsearch/index.mjs", + "types": "./lib/vectorsearch/index.d.ts", + "default": "./lib/vectorsearch/index.js" } }, "typesVersions": { @@ -96,6 +102,9 @@ ], "modelgarden": [ "./lib/modelgarden/index" + ], + "vectorsearch": [ + "./lib/vectorsearch/index" ] } } diff --git a/js/plugins/vertexai/src/index.ts b/js/plugins/vertexai/src/index.ts index 4454e24f5..1b6566e1a 100644 --- a/js/plugins/vertexai/src/index.ts +++ b/js/plugins/vertexai/src/index.ts @@ -40,22 +40,7 @@ import { imagen3Fast, imagenModel, } from './imagen.js'; -import { vertexAiIndexers, vertexAiRetrievers } from './vector-search/index.js'; export { PluginOptions } from './common/types.js'; -export { - DocumentIndexer, - DocumentRetriever, - Neighbor, - VectorSearchOptions, - getBigQueryDocumentIndexer, - getBigQueryDocumentRetriever, - getFirestoreDocumentIndexer, - getFirestoreDocumentRetriever, - vertexAiIndexerRef, - vertexAiIndexers, - vertexAiRetrieverRef, - vertexAiRetrievers, -} from './vector-search/index.js'; export { gemini10Pro, gemini15Flash, @@ -83,28 +68,9 @@ export function vertexAI(options?: PluginOptions): GenkitPlugin { defineGeminiModel(ai, name, vertexClientFactory, { projectId, location }) ); - const embedders = Object.keys(SUPPORTED_EMBEDDER_MODELS).map((name) => + Object.keys(SUPPORTED_EMBEDDER_MODELS).map((name) => defineVertexAIEmbedder(ai, name, authClient, { projectId, location }) ); - - if ( - options?.vectorSearchOptions && - options.vectorSearchOptions.length > 0 - ) { - const defaultEmbedder = embedders[0]; - - vertexAiIndexers(ai, { - pluginOptions: options, - authClient, - defaultEmbedder, - }); - - vertexAiRetrievers(ai, { - pluginOptions: options, - authClient, - defaultEmbedder, - }); - } }); } diff --git a/js/plugins/vertexai/src/vector-search/index.ts b/js/plugins/vertexai/src/vectorsearch/index.ts similarity index 85% rename from js/plugins/vertexai/src/vector-search/index.ts rename to js/plugins/vertexai/src/vectorsearch/index.ts index 340d73de5..048e18b72 100644 --- a/js/plugins/vertexai/src/vector-search/index.ts +++ b/js/plugins/vertexai/src/vectorsearch/index.ts @@ -18,14 +18,6 @@ import { Genkit } from 'genkit'; import { GenkitPlugin, genkitPlugin } from 'genkit/plugin'; import { getDerivedParams } from '../common/index.js'; import { PluginOptions } from './types.js'; -// import { -// SUPPORTED_EMBEDDER_MODELS, -// defineVertexAIEmbedder, -// textEmbedding004, -// textEmbeddingGecko003, -// textEmbeddingGeckoMultilingual001, -// textMultilingualEmbedding002, -// } from '../embedder.js'; import { vertexAiIndexers, vertexAiRetrievers } from './vector_search/index.js'; export { PluginOptions } from '../common/types.js'; export { @@ -45,7 +37,7 @@ export { /** * Add Google Cloud Vertex AI to Genkit. Includes Gemini and Imagen models and text embedder. */ -export function vertexAI(options?: PluginOptions): GenkitPlugin { +export function vertexAIVectorSearch(options?: PluginOptions): GenkitPlugin { return genkitPlugin('vertexai', async (ai: Genkit) => { const { projectId, location, vertexClientFactory, authClient } = await getDerivedParams(options); @@ -68,5 +60,3 @@ export function vertexAI(options?: PluginOptions): GenkitPlugin { } }); } - -export default vertexAI; diff --git a/js/plugins/vertexai/src/vector-search/types.ts b/js/plugins/vertexai/src/vectorsearch/types.ts similarity index 92% rename from js/plugins/vertexai/src/vector-search/types.ts rename to js/plugins/vertexai/src/vectorsearch/types.ts index 63ffa7a3f..b3d790ff3 100644 --- a/js/plugins/vertexai/src/vector-search/types.ts +++ b/js/plugins/vertexai/src/vectorsearch/types.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { z } from 'genkit'; +import { EmbedderReference, z } from 'genkit'; import { CommonPluginOptions } from '../common/types'; import { VectorSearchOptions } from './vector_search'; @@ -22,6 +22,7 @@ import { VectorSearchOptions } from './vector_search'; export interface VectorSearchOptionsConfig { /** Configure Vertex AI vector search index options */ vectorSearchOptions?: VectorSearchOptions[]; + embedder?: EmbedderReference; } export interface PluginOptions diff --git a/js/plugins/vertexai/src/vector-search/vector_search/bigquery.ts b/js/plugins/vertexai/src/vectorsearch/vector_search/bigquery.ts similarity index 100% rename from js/plugins/vertexai/src/vector-search/vector_search/bigquery.ts rename to js/plugins/vertexai/src/vectorsearch/vector_search/bigquery.ts diff --git a/js/plugins/vertexai/src/vector-search/vector_search/firestore.ts b/js/plugins/vertexai/src/vectorsearch/vector_search/firestore.ts similarity index 100% rename from js/plugins/vertexai/src/vector-search/vector_search/firestore.ts rename to js/plugins/vertexai/src/vectorsearch/vector_search/firestore.ts diff --git a/js/plugins/vertexai/src/vector-search/vector_search/index.ts b/js/plugins/vertexai/src/vectorsearch/vector_search/index.ts similarity index 100% rename from js/plugins/vertexai/src/vector-search/vector_search/index.ts rename to js/plugins/vertexai/src/vectorsearch/vector_search/index.ts diff --git a/js/plugins/vertexai/src/vector-search/vector_search/indexers.ts b/js/plugins/vertexai/src/vectorsearch/vector_search/indexers.ts similarity index 93% rename from js/plugins/vertexai/src/vector-search/vector_search/indexers.ts rename to js/plugins/vertexai/src/vectorsearch/vector_search/indexers.ts index 66a00e913..aabd23a1f 100644 --- a/js/plugins/vertexai/src/vector-search/vector_search/indexers.ts +++ b/js/plugins/vertexai/src/vectorsearch/vector_search/indexers.ts @@ -58,7 +58,6 @@ export function vertexAiIndexers( params: VertexVectorSearchOptions ): IndexerAction[] { const vectorSearchOptions = params.pluginOptions.vectorSearchOptions; - const defaultEmbedder = params.defaultEmbedder; const indexers: IndexerAction[] = []; if (!vectorSearchOptions || vectorSearchOptions.length === 0) { @@ -67,7 +66,14 @@ export function vertexAiIndexers( for (const vectorSearchOption of vectorSearchOptions) { const { documentIndexer, indexId } = vectorSearchOption; - const embedder = vectorSearchOption.embedder ?? defaultEmbedder; + const embedderReference = + vectorSearchOption.embedder ?? params.defaultEmbedder; + + if (!embedderReference) { + throw new Error( + 'Embedder reference is required to define Vertex AI retriever' + ); + } const embedderOptions = vectorSearchOption.embedderOptions; const indexer = ai.defineIndexer( @@ -87,7 +93,7 @@ export function vertexAiIndexers( } const embeddings = await ai.embedMany({ - embedder, + embedder: embedderReference, content: docs, options: embedderOptions, }); diff --git a/js/plugins/vertexai/src/vector-search/vector_search/query_public_endpoint.ts b/js/plugins/vertexai/src/vectorsearch/vector_search/query_public_endpoint.ts similarity index 100% rename from js/plugins/vertexai/src/vector-search/vector_search/query_public_endpoint.ts rename to js/plugins/vertexai/src/vectorsearch/vector_search/query_public_endpoint.ts diff --git a/js/plugins/vertexai/src/vector-search/vector_search/retrievers.ts b/js/plugins/vertexai/src/vectorsearch/vector_search/retrievers.ts similarity index 93% rename from js/plugins/vertexai/src/vector-search/vector_search/retrievers.ts rename to js/plugins/vertexai/src/vectorsearch/vector_search/retrievers.ts index 67f47f33d..0f8a64024 100644 --- a/js/plugins/vertexai/src/vector-search/vector_search/retrievers.ts +++ b/js/plugins/vertexai/src/vectorsearch/vector_search/retrievers.ts @@ -57,8 +57,17 @@ export function vertexAiRetrievers( configSchema: VertexAIVectorRetrieverOptionsSchema.optional(), }, async (content, options) => { + const embedderReference = + vectorSearchOption.embedder ?? defaultEmbedder; + + if (!embedderReference) { + throw new Error( + 'Embedder reference is required to define Vertex AI retriever' + ); + } + const queryEmbeddings = await ai.embed({ - embedder, + embedder: embedderReference, options: embedderOptions, content, }); diff --git a/js/plugins/vertexai/src/vector-search/vector_search/types.ts b/js/plugins/vertexai/src/vectorsearch/vector_search/types.ts similarity index 97% rename from js/plugins/vertexai/src/vector-search/vector_search/types.ts rename to js/plugins/vertexai/src/vectorsearch/vector_search/types.ts index ce0929847..7034d1489 100644 --- a/js/plugins/vertexai/src/vector-search/vector_search/types.ts +++ b/js/plugins/vertexai/src/vectorsearch/vector_search/types.ts @@ -16,7 +16,7 @@ import * as aiplatform from '@google-cloud/aiplatform'; import { z } from 'genkit'; -import { EmbedderArgument } from 'genkit/embedder'; +import { EmbedderReference } from 'genkit/embedder'; import { CommonRetrieverOptionsSchema, Document } from 'genkit/retriever'; import { GoogleAuth } from 'google-auth-library'; import { PluginOptions } from '../types'; @@ -27,7 +27,7 @@ export interface VertexVectorSearchOptions< > { pluginOptions: PluginOptions; authClient: GoogleAuth; - defaultEmbedder: EmbedderArgument; + defaultEmbedder?: EmbedderReference; } export type IIndexDatapoint = @@ -184,6 +184,6 @@ export interface VectorSearchOptions< documentRetriever: DocumentRetriever; documentIndexer: DocumentIndexer; // Embedder and default options to use for indexing and retrieval - embedder?: EmbedderArgument; + embedder?: EmbedderReference; embedderOptions?: z.infer; } diff --git a/js/plugins/vertexai/src/vector-search/vector_search/upsert_datapoints.ts b/js/plugins/vertexai/src/vectorsearch/vector_search/upsert_datapoints.ts similarity index 100% rename from js/plugins/vertexai/src/vector-search/vector_search/upsert_datapoints.ts rename to js/plugins/vertexai/src/vectorsearch/vector_search/upsert_datapoints.ts diff --git a/js/plugins/vertexai/src/vector-search/vector_search/utils.ts b/js/plugins/vertexai/src/vectorsearch/vector_search/utils.ts similarity index 100% rename from js/plugins/vertexai/src/vector-search/vector_search/utils.ts rename to js/plugins/vertexai/src/vectorsearch/vector_search/utils.ts diff --git a/js/testapps/model-tester/src/index.ts b/js/testapps/model-tester/src/index.ts index cc9fb5361..7c1fe1fcb 100644 --- a/js/testapps/model-tester/src/index.ts +++ b/js/testapps/model-tester/src/index.ts @@ -15,7 +15,12 @@ */ import { googleAI } from '@genkit-ai/googleai'; -import { claude3Sonnet, llama31, vertexAI } from '@genkit-ai/vertexai'; +import { vertexAI } from '@genkit-ai/vertexai'; +import { + claude3Sonnet, + llama31, + vertexAIModelGarden, +} from '@genkit-ai/vertexai/modelgarden'; import * as clc from 'colorette'; import { genkit } from 'genkit'; import { testModels } from 'genkit/testing'; @@ -27,6 +32,9 @@ export const ai = genkit({ googleAI(), vertexAI({ location: 'us-central1', + }), + vertexAIModelGarden({ + location: 'us-central1', modelGarden: { models: [claude3Sonnet, llama31], }, diff --git a/js/testapps/rag/src/genkit.ts b/js/testapps/rag/src/genkit.ts index 5e2cd4163..b8436f752 100644 --- a/js/testapps/rag/src/genkit.ts +++ b/js/testapps/rag/src/genkit.ts @@ -17,12 +17,12 @@ import { devLocalVectorstore } from '@genkit-ai/dev-local-vectorstore'; import { genkitEval, GenkitMetric } from '@genkit-ai/evaluator'; import { gemini15Flash, googleAI } from '@genkit-ai/googleai'; +import { textEmbedding004, vertexAI } from '@genkit-ai/vertexai'; import { claude3Sonnet, llama31, - textEmbedding004, - vertexAI, -} from '@genkit-ai/vertexai'; + vertexAIModelGarden, +} from '@genkit-ai/vertexai/modelgarden'; import { genkit } from 'genkit'; import { chroma } from 'genkitx-chromadb'; import { langchain } from 'genkitx-langchain'; @@ -76,6 +76,9 @@ export const ai = genkit({ }), vertexAI({ location: 'us-central1', + }), + vertexAIModelGarden({ + location: 'us-central1', modelGarden: { models: [claude3Sonnet, llama31], }, diff --git a/js/testapps/vertexai-vector-search-bigquery/src/index.ts b/js/testapps/vertexai-vector-search-bigquery/src/index.ts index 59c7763ee..a2b13b6ee 100644 --- a/js/testapps/vertexai-vector-search-bigquery/src/index.ts +++ b/js/testapps/vertexai-vector-search-bigquery/src/index.ts @@ -18,16 +18,16 @@ import { Document, genkit, z } from 'genkit'; // important imports for this sample: +import { textEmbedding004, vertexAI } from '@genkit-ai/vertexai'; import { + DocumentIndexer, + DocumentRetriever, 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 { BIGQUERY_DATASET, @@ -81,6 +81,11 @@ const ai = genkit({ googleAuth: { scopes: ['https://www.googleapis.com/auth/cloud-platform'], }, + }), + vertexAIVectorSearch({ + location: LOCATION, + projectId: PROJECT_ID, + embedder: textEmbedding004, vectorSearchOptions: [ { publicDomainName: VECTOR_SEARCH_PUBLIC_DOMAIN_NAME, @@ -95,7 +100,6 @@ const ai = genkit({ ], }); -// // Define indexing flow export const indexFlow = ai.defineFlow( { name: 'indexFlow', diff --git a/js/testapps/vertexai-vector-search-custom/src/index.ts b/js/testapps/vertexai-vector-search-custom/src/index.ts index 96a708d4a..581fa5079 100644 --- a/js/testapps/vertexai-vector-search-custom/src/index.ts +++ b/js/testapps/vertexai-vector-search-custom/src/index.ts @@ -18,14 +18,15 @@ import { Document, genkit, z } from 'genkit'; // important imports for this sample: +import { textEmbedding004, vertexAI } from '@genkit-ai/vertexai'; 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 { @@ -148,6 +149,10 @@ const ai = genkit({ googleAuth: { scopes: ['https://www.googleapis.com/auth/cloud-platform'], }, + }), + vertexAIVectorSearch({ + location: LOCATION, + projectId: PROJECT_ID, vectorSearchOptions: [ { publicDomainName: VECTOR_SEARCH_PUBLIC_DOMAIN_NAME, @@ -156,6 +161,7 @@ const ai = genkit({ deployedIndexId: VECTOR_SEARCH_DEPLOYED_INDEX_ID, documentRetriever: localDocumentRetriever, documentIndexer: localDocumentIndexer, + embedder: textEmbedding004, }, ], }), diff --git a/js/testapps/vertexai-vector-search-firestore/src/index.ts b/js/testapps/vertexai-vector-search-firestore/src/index.ts index fda773418..bad5aba88 100644 --- a/js/testapps/vertexai-vector-search-firestore/src/index.ts +++ b/js/testapps/vertexai-vector-search-firestore/src/index.ts @@ -19,15 +19,18 @@ import { initializeApp } from 'firebase-admin/app'; import { Document, genkit, z } from 'genkit'; // important imports for this sample: + +import { textEmbedding004, vertexAI } from '@genkit-ai/vertexai'; + import { DocumentIndexer, DocumentRetriever, getFirestoreDocumentIndexer, getFirestoreDocumentRetriever, - vertexAI, vertexAiIndexerRef, vertexAiRetrieverRef, -} from '@genkit-ai/vertexai'; + vertexAIVectorSearch, +} from '@genkit-ai/vertexai/vectorsearch'; // // Environment variables set with dotenv for simplicity of sample import { getFirestore } from 'firebase-admin/firestore'; @@ -80,6 +83,10 @@ const ai = genkit({ googleAuth: { scopes: ['https://www.googleapis.com/auth/cloud-platform'], }, + }), + vertexAIVectorSearch({ + projectId: PROJECT_ID, + location: LOCATION, vectorSearchOptions: [ { publicDomainName: VECTOR_SEARCH_PUBLIC_DOMAIN_NAME, @@ -88,6 +95,7 @@ const ai = genkit({ deployedIndexId: VECTOR_SEARCH_DEPLOYED_INDEX_ID, documentRetriever: firestoreDocumentRetriever, documentIndexer: firestoreDocumentIndexer, + embedder: textEmbedding004, }, ], }),