From 67b61514c24fc192efe46e8875faf2eafeef2817 Mon Sep 17 00:00:00 2001
From: jiiyaaa__ <66894607+river0g@users.noreply.github.com>
Date: Tue, 17 Dec 2024 06:10:28 +0900
Subject: [PATCH] feat(js/plugins): added custom dimensions parameter to
 text-embedding (#1529)

---
 js/plugins/googleai/src/embedder.ts |  9 +++++++++
 js/plugins/vertexai/src/embedder.ts | 11 ++++++++++-
 js/plugins/vertexai/src/predict.ts  |  6 +++---
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/js/plugins/googleai/src/embedder.ts b/js/plugins/googleai/src/embedder.ts
index 00adfb22f..31e27afa0 100644
--- a/js/plugins/googleai/src/embedder.ts
+++ b/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 };
diff --git a/js/plugins/vertexai/src/embedder.ts b/js/plugins/vertexai/src/embedder.ts
index 28e8ad7ff..8e81599ac 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 `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) => ({
diff --git a/js/plugins/vertexai/src/predict.ts b/js/plugins/vertexai/src/predict.ts
index 0d4f39302..93bf8277c 100644
--- a/js/plugins/vertexai/src/predict.ts
+++ b/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(