feat(js/plugins): added custom dimensions parameter to text-embedding (…

…#1529)

js/plugins/googleai/src/embedder.ts

@@ -36,6 +36,14 @@ export const GeminiEmbeddingConfigSchema = z.object({
   taskType: TaskTypeSchema.optional(),
   title: z.string().optional(),
   version: z.string().optional(),
+  /**
+   * The `outputDimensionality` 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.
+   **/
+  outputDimensionality: z.number().min(1).max(768).optional(),
 });
 
 export type GeminiEmbeddingConfig = z.infer<typeof GeminiEmbeddingConfigSchema>;
@@ -122,6 +130,7 @@ export function defineGoogleAIEmbedder(
               role: '',
               parts: [{ text: doc.text }],
             },
+            outputDimensionality: options?.outputDimensionality,
           } as EmbedContentRequest);
           const values = response.embedding.values;
           return { embedding: values };

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 `outputDimensionality` 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.
+   **/
+  outputDimensionality: z.number().min(1).max(768).optional(),
 });
 
 export type VertexEmbeddingConfig = z.infer<typeof VertexEmbeddingConfigSchema>;
@@ -145,7 +153,8 @@ export function defineVertexAIEmbedder(
             task_type: options?.taskType,
             title: options?.title,
           };
-        })
+        }),
+        { outputDimensionality: options?.outputDimensionality }
       );
       return {
         embeddings: response.predictions.map((p) => ({

js/plugins/vertexai/src/predict.ts

@@ -33,7 +33,7 @@ interface PredictionResponse<R> {
 
 export type PredictClient<I = unknown, R = unknown, P = unknown> = (
   instances: I[],
-  parameters?: P
+  parameters: P
 ) => Promise<PredictionResponse<R>>;
 
 export function predictModel<I = unknown, R = unknown, P = unknown>(
@@ -43,14 +43,14 @@ export function predictModel<I = unknown, R = unknown, P = unknown>(
 ): PredictClient<I, R, P> {
   return async (
     instances: I[],
-    parameters?: P
+    parameters: P
   ): Promise<PredictionResponse<R>> => {
     const fetch = (await import('node-fetch')).default;
 
     const accessToken = await auth.getAccessToken();
     const req = {
       instances,
-      parameters: parameters || {},
+      parameters,
     };
 
     const response = await fetch(