add custom dimensions

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<typeof VertexEmbeddingConfigSchema>;
@@ -83,6 +91,12 @@ export const SUPPORTED_EMBEDDER_MODELS: Record<string, EmbedderReference> = {
   // ]),
 };
 
+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) => ({