Apply fixes based on review comments. and Update parameter evaluation.

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

Hide file tree

Showing 3 changed files with 7 additions and 25 deletions.
diff --git a/js/plugins/googleai/src/embedder.ts b/js/plugins/googleai/src/embedder.ts
@@ -37,13 +37,13 @@ export const GeminiEmbeddingConfigSchema = z.object({
   title: z.string().optional(),
   version: z.string().optional(),
   /**
-   * The `dimensions` parameter allows you to specify the dimensionality of the embedding output.
+   * 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.
    **/
-  dimensions: z.number().min(1).max(768).optional(),
+  outputDimensionality: z.number().min(1).max(768).optional(),
 });
 
 export type GeminiEmbeddingConfig = z.infer<typeof GeminiEmbeddingConfigSchema>;
@@ -77,8 +77,6 @@ export const SUPPORTED_MODELS = {
   'text-embedding-004': textEmbedding004,
 };
 
-const customDimensionsSupportedModelNames = [textEmbedding004.name];
-
 export function defineGoogleAIEmbedder(
   ai: Genkit,
   name: string,
@@ -116,11 +114,6 @@ 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 ||
@@ -137,7 +130,7 @@ export function defineGoogleAIEmbedder(
               role: '',
               parts: [{ text: doc.text }],
             },
-            outputDimensionality: dimensions,
+            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
@@ -40,13 +40,13 @@ export const VertexEmbeddingConfigSchema = z.object({
   location: z.string().optional(),
   version: z.string().optional(),
   /**
-   * The `dimensions` parameter allows you to specify the dimensionality of the embedding output.
+   * 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.
    **/
-  dimensions: z.number().min(1).max(768).optional(),
+  outputDimensionality: z.number().min(1).max(768).optional(),
 });
 
 export type VertexEmbeddingConfig = z.infer<typeof VertexEmbeddingConfigSchema>;
@@ -91,12 +91,6 @@ export const SUPPORTED_EMBEDDER_MODELS: Record<string, EmbedderReference> = {
   // ]),
 };
 
-const customDimensionsSupportedModelNames = [
-  textEmbedding004.name,
-  textEmbedding005.name,
-  textMultilingualEmbedding002.name,
-];
-
 interface EmbeddingInstance {
   task_type?: TaskType;
   content: string;
@@ -151,11 +145,6 @@ 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) => {
@@ -165,7 +154,7 @@ export function defineVertexAIEmbedder(
             title: options?.title,
           };
         }),
-        { outputDimensionality: dimensions }
+        { 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
@@ -50,7 +50,7 @@ export function predictModel<I = unknown, R = unknown, P = unknown>(
     const accessToken = await auth.getAccessToken();
     const req = {
       instances,
-      parameters: parameters || {},
+      parameters,
     };
 
     const response = await fetch(