add googleai embed dimension and fix vertexai parameter

firebase · Dec 16, 2024 · aa46a65 · aa46a65
1 parent a4e3892
commit aa46a65
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 3 deletions.
diff --git 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 `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 GeminiEmbeddingConfig = z.infer<typeof GeminiEmbeddingConfigSchema>;
@@ -69,6 +77,8 @@ export const SUPPORTED_MODELS = {
   'text-embedding-004': textEmbedding004,
 };
 
+const customDimensionsSupportedModelNames = [textEmbedding004.name];
+
 export function defineGoogleAIEmbedder(
   ai: Genkit,
   name: string,
@@ -106,6 +116,11 @@ export function defineGoogleAIEmbedder(
       info: embedder.info!,
     },
     async (input, options) => {
+      const dimensions =
+        customDimensionsSupportedModelNames.includes(embedder.name) &&
+        options?.dimensions
+          ? options.dimensions
+          : undefined;
       const client = new GoogleGenerativeAI(apiKey!).getGenerativeModel({
         model:
           options?.version ||
@@ -122,6 +137,7 @@ export function defineGoogleAIEmbedder(
               role: '',
               parts: [{ text: doc.text }],
             },
+            outputDimensionality: dimensions,
           } 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
@@ -165,7 +165,7 @@ export function defineVertexAIEmbedder(
             title: options?.title,
           };
         }),
-        dimensions ? { outputDimensionality: dimensions } : {}
+        { outputDimensionality: dimensions }
       );
       return {
         embeddings: response.predictions.map((p) => ({

diff --git 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,7 +43,7 @@ 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;