diff --git a/js/plugins/vertexai/src/embedder.ts b/js/plugins/vertexai/src/embedder.ts index 28e8ad7ff..e445a74ca 100644 --- a/js/plugins/vertexai/src/embedder.ts +++ b/js/plugins/vertexai/src/embedder.ts @@ -39,6 +39,14 @@ export const VertexEmbeddingConfigSchema = z.object({ title: z.string().optional(), location: z.string().optional(), version: z.string().optional(), + /** + * The `dimensions` parameter allows you to specify the dimensionality of the embedding output. + * By default, the model generates embeddings with 768 dimensions. Models such as + * `text-embedding-004`, `text-embedding-005`, and `text-multilingual-embedding-002` + * allow the output dimensionality to be adjusted between 1 and 768. + * By selecting a smaller output dimensionality, users can save memory and storage space, leading to more efficient computations. + **/ + dimensions: z.number().min(1).max(768).optional(), }); export type VertexEmbeddingConfig = z.infer; @@ -83,6 +91,12 @@ export const SUPPORTED_EMBEDDER_MODELS: Record = { // ]), }; +const customDimensionsSupportedModelNames = [ + textEmbedding004.name, + textEmbedding005.name, + textMultilingualEmbedding002.name, +]; + interface EmbeddingInstance { task_type?: TaskType; content: string; @@ -137,6 +151,11 @@ export function defineVertexAIEmbedder( info: embedder.info!, }, async (input, options) => { + const dimensions = + customDimensionsSupportedModelNames.includes(embedder.name) && + options?.dimensions + ? options.dimensions + : undefined; const predictClient = predictClientFactory(options); const response = await predictClient( input.map((i) => { @@ -145,7 +164,8 @@ export function defineVertexAIEmbedder( task_type: options?.taskType, title: options?.title, }; - }) + }), + dimensions ? { outputDimensionality: dimensions } : {} ); return { embeddings: response.predictions.map((p) => ({