From fb5b68a67621721da31f9eb9302c67030a20caca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebasti=C3=A1n=20Dar=C3=ADo=20Ram=C3=ADrez?=
 <gitlab@seba.work>
Date: Mon, 5 May 2025 10:08:16 -0400
Subject: [PATCH 1/2] adds gemini-embedding-exp-03-07 support

---
 .../GeminiOptions/index.jsx                   |   2 +-
 server/utils/EmbeddingEngines/gemini/index.js | 102 +++++++++++-------
 2 files changed, 62 insertions(+), 42 deletions(-)

diff --git a/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx b/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx
index c25a2b13f8d..fcafe0b0e08 100644
--- a/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx
+++ b/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx
@@ -27,7 +27,7 @@ export default function GeminiOptions({ settings }) {
             className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
           >
             <optgroup label="Available embedding models">
-              {["text-embedding-004"].map((model) => {
+              {["text-embedding-004", "gemini-embedding-exp-03-07"].map((model) => {
                 return (
                   <option
                     key={model}
diff --git a/server/utils/EmbeddingEngines/gemini/index.js b/server/utils/EmbeddingEngines/gemini/index.js
index a7d50c90663..43af061ac6a 100644
--- a/server/utils/EmbeddingEngines/gemini/index.js
+++ b/server/utils/EmbeddingEngines/gemini/index.js
@@ -1,33 +1,45 @@
 const { toChunks } = require("../../helpers");
 
+/**
+ * Gemini embeddings engine implementation for generating vector embeddings
+ */
 class GeminiEmbedder {
   constructor() {
-    if (!process.env.GEMINI_EMBEDDING_API_KEY)
+    if (!process.env.GEMINI_EMBEDDING_API_KEY) {
       throw new Error("No Gemini API key was set.");
+    }
 
     const { OpenAI: OpenAIApi } = require("openai");
+    
+    // Configure the embedding model and client
     this.model = process.env.EMBEDDING_MODEL_PREF || "text-embedding-004";
     this.openai = new OpenAIApi({
       apiKey: process.env.GEMINI_EMBEDDING_API_KEY,
-      // Even models that are v1 in gemini API can be used with v1beta/openai/ endpoint and nobody knows why.
+      // Even models that are v1 in gemini API can be used with v1beta/openai/ endpoint
       baseURL: "https://generativelanguage.googleapis.com/v1beta/openai/",
     });
 
+    // Embedding configuration
     this.maxConcurrentChunks = 4;
-
-    // https://ai.google.dev/gemini-api/docs/models/gemini#text-embedding-and-embedding
-    // TODO: May need to make this dynamic based on the model
-    this.embeddingMaxChunkLength = 2_048;
-    this.log(`Initialized with ${this.model}`);
+    
+    // Set max chunk length based on model - newer models support longer context
+    this.embeddingMaxChunkLength = this.model === "gemini-embedding-exp-03-07" 
+      ? 8_192 
+      : 2_048;
+      
+    this.log(`Initialized with ${this.model} (max tokens: ${this.embeddingMaxChunkLength})`);
   }
 
+  /**
+   * Helper to output consistent log messages
+   */
   log(text, ...args) {
     console.log(`\x1b[36m[GeminiEmbedder]\x1b[0m ${text}`, ...args);
   }
 
   /**
    * Embeds a single text input
-   * @param {string} textInput - The text to embed
+   * @param {string|string[]} textInput - The text to embed
    * @returns {Promise<Array<number>>} The embedding values
    */
   async embedTextInput(textInput) {
@@ -39,7 +51,7 @@ class GeminiEmbedder {
 
   /**
    * Embeds a list of text inputs
-   * @param {Array<string>} textInputs - The list of text to embed
+   * @param {Array<string>} textChunks - The list of text to embed
    * @returns {Promise<Array<Array<number>>>} The embedding values
    */
   async embedChunks(textChunks = []) {
@@ -47,8 +59,8 @@ class GeminiEmbedder {
 
     // Because there is a hard POST limit on how many chunks can be sent at once to OpenAI (~8mb)
     // we concurrently execute each max batch of text chunks possible.
-    // Refer to constructor maxConcurrentChunks for more info.
     const embeddingRequests = [];
+    
     for (const chunk of toChunks(textChunks, this.maxConcurrentChunks)) {
       embeddingRequests.push(
         new Promise((resolve) => {
@@ -60,47 +72,55 @@ class GeminiEmbedder {
             .then((result) => {
               resolve({ data: result?.data, error: null });
             })
-            .catch((e) => {
-              e.type =
-                e?.response?.data?.error?.code ||
-                e?.response?.status ||
+            .catch((error) => {
+              const errorType = 
+                error?.response?.data?.error?.code ||
+                error?.response?.status ||
                 "failed_to_embed";
-              e.message = e?.response?.data?.error?.message || e.message;
-              resolve({ data: [], error: e });
+                
+              const errorMessage = error?.response?.data?.error?.message || error.message;
+              
+              error.type = errorType;
+              error.message = errorMessage;
+              resolve({ data: [], error });
             });
         })
       );
     }
 
-    const { data = [], error = null } = await Promise.all(
-      embeddingRequests
-    ).then((results) => {
-      // If any errors were returned from OpenAI abort the entire sequence because the embeddings
-      // will be incomplete.
-      const errors = results
-        .filter((res) => !!res.error)
-        .map((res) => res.error)
-        .flat();
-      if (errors.length > 0) {
-        let uniqueErrors = new Set();
-        errors.map((error) =>
-          uniqueErrors.add(`[${error.type}]: ${error.message}`)
-        );
+    const { data = [], error = null } = await Promise.all(embeddingRequests)
+      .then((results) => {
+        // If any errors were returned from OpenAI abort the entire sequence because the embeddings
+        // will be incomplete.
+        const errors = results
+          .filter((res) => !!res.error)
+          .map((res) => res.error)
+          .flat();
+          
+        if (errors.length > 0) {
+          const uniqueErrors = new Set();
+          errors.forEach((error) =>
+            uniqueErrors.add(`[${error.type}]: ${error.message}`)
+          );
 
+          return {
+            data: [],
+            error: Array.from(uniqueErrors).join(", "),
+          };
+        }
+        
         return {
-          data: [],
-          error: Array.from(uniqueErrors).join(", "),
+          data: results.map((res) => res?.data || []).flat(),
+          error: null,
         };
-      }
-      return {
-        data: results.map((res) => res?.data || []).flat(),
-        error: null,
-      };
-    });
+      });
 
-    if (!!error) throw new Error(`OpenAI Failed to embed: ${error}`);
-    return data.length > 0 &&
-      data.every((embd) => embd.hasOwnProperty("embedding"))
+    if (!!error) {
+      throw new Error(`OpenAI Failed to embed: ${error}`);
+    }
+    
+    // Ensure we have valid embeddings before returning
+    return data.length > 0 && data.every((embd) => embd.hasOwnProperty("embedding"))
       ? data.map((embd) => embd.embedding)
       : null;
   }

From c3f65279be16564330b4039198ad550465a02809 Mon Sep 17 00:00:00 2001
From: timothycarambat <rambat1010@gmail.com>
Date: Mon, 5 May 2025 09:35:29 -0700
Subject: [PATCH 2/2] unset random changes, make model map

---
 .../GeminiOptions/index.jsx                   |  25 +++-
 server/utils/EmbeddingEngines/gemini/index.js | 107 ++++++++----------
 2 files changed, 67 insertions(+), 65 deletions(-)

diff --git a/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx b/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx
index fcafe0b0e08..18e1530ac6b 100644
--- a/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx
+++ b/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx
@@ -1,3 +1,18 @@
+const DEFAULT_MODELS = [
+  {
+    id: "embedding-001",
+    name: "Embedding 001",
+  },
+  {
+    id: "text-embedding-004",
+    name: "Text Embedding 004",
+  },
+  {
+    id: "gemini-embedding-exp-03-07",
+    name: "Gemini Embedding Exp 03 07",
+  },
+];
+
 export default function GeminiOptions({ settings }) {
   return (
     <div className="w-full flex flex-col gap-y-4">
@@ -27,14 +42,14 @@ export default function GeminiOptions({ settings }) {
             className="border-none bg-theme-settings-input-bg border-gray-500 text-white text-sm rounded-lg block w-full p-2.5"
           >
             <optgroup label="Available embedding models">
-              {["text-embedding-004", "gemini-embedding-exp-03-07"].map((model) => {
+              {DEFAULT_MODELS.map((model) => {
                 return (
                   <option
-                    key={model}
-                    value={model}
-                    selected={settings?.EmbeddingModelPref === model}
+                    key={model.id}
+                    value={model.id}
+                    selected={settings?.EmbeddingModelPref === model.id}
                   >
-                    {model}
+                    {model.name}
                   </option>
                 );
               })}
diff --git a/server/utils/EmbeddingEngines/gemini/index.js b/server/utils/EmbeddingEngines/gemini/index.js
index 43af061ac6a..8c505466728 100644
--- a/server/utils/EmbeddingEngines/gemini/index.js
+++ b/server/utils/EmbeddingEngines/gemini/index.js
@@ -1,38 +1,33 @@
 const { toChunks } = require("../../helpers");
 
-/**
- * Gemini embeddings engine implementation for generating vector embeddings
- */
+const MODEL_MAP = {
+  "embedding-001": 2048,
+  "text-embedding-004": 2048,
+  "gemini-embedding-exp-03-07": 8192,
+};
+
 class GeminiEmbedder {
   constructor() {
-    if (!process.env.GEMINI_EMBEDDING_API_KEY) {
+    if (!process.env.GEMINI_EMBEDDING_API_KEY)
       throw new Error("No Gemini API key was set.");
-    }
 
     const { OpenAI: OpenAIApi } = require("openai");
-    
-    // Configure the embedding model and client
     this.model = process.env.EMBEDDING_MODEL_PREF || "text-embedding-004";
     this.openai = new OpenAIApi({
       apiKey: process.env.GEMINI_EMBEDDING_API_KEY,
-      // Even models that are v1 in gemini API can be used with v1beta/openai/ endpoint
+      // Even models that are v1 in gemini API can be used with v1beta/openai/ endpoint and nobody knows why.
       baseURL: "https://generativelanguage.googleapis.com/v1beta/openai/",
     });
 
-    // Embedding configuration
     this.maxConcurrentChunks = 4;
-    
-    // Set max chunk length based on model - newer models support longer context
-    this.embeddingMaxChunkLength = this.model === "gemini-embedding-exp-03-07" 
-      ? 8_192 
-      : 2_048;
-      
-    this.log(`Initialized with ${this.model} (max tokens: ${this.embeddingMaxChunkLength})`);
+
+    // https://ai.google.dev/gemini-api/docs/models/gemini#text-embedding-and-embedding
+    this.embeddingMaxChunkLength = MODEL_MAP[this.model] || 2_048;
+    this.log(
+      `Initialized with ${this.model} - Max Size: ${this.embeddingMaxChunkLength}`
+    );
   }
 
-  /**
-   * Helper to output consistent log messages
-   */
   log(text, ...args) {
     console.log(`\x1b[36m[GeminiEmbedder]\x1b[0m ${text}`, ...args);
   }
@@ -51,7 +46,7 @@ class GeminiEmbedder {
 
   /**
    * Embeds a list of text inputs
-   * @param {Array<string>} textChunks - The list of text to embed
+   * @param {string[]} textChunks - The list of text to embed
    * @returns {Promise<Array<Array<number>>>} The embedding values
    */
   async embedChunks(textChunks = []) {
@@ -59,8 +54,8 @@ class GeminiEmbedder {
 
     // Because there is a hard POST limit on how many chunks can be sent at once to OpenAI (~8mb)
     // we concurrently execute each max batch of text chunks possible.
+    // Refer to constructor maxConcurrentChunks for more info.
     const embeddingRequests = [];
-    
     for (const chunk of toChunks(textChunks, this.maxConcurrentChunks)) {
       embeddingRequests.push(
         new Promise((resolve) => {
@@ -72,55 +67,47 @@ class GeminiEmbedder {
             .then((result) => {
               resolve({ data: result?.data, error: null });
             })
-            .catch((error) => {
-              const errorType = 
-                error?.response?.data?.error?.code ||
-                error?.response?.status ||
+            .catch((e) => {
+              e.type =
+                e?.response?.data?.error?.code ||
+                e?.response?.status ||
                 "failed_to_embed";
-                
-              const errorMessage = error?.response?.data?.error?.message || error.message;
-              
-              error.type = errorType;
-              error.message = errorMessage;
-              resolve({ data: [], error });
+              e.message = e?.response?.data?.error?.message || e.message;
+              resolve({ data: [], error: e });
             });
         })
       );
     }
 
-    const { data = [], error = null } = await Promise.all(embeddingRequests)
-      .then((results) => {
-        // If any errors were returned from OpenAI abort the entire sequence because the embeddings
-        // will be incomplete.
-        const errors = results
-          .filter((res) => !!res.error)
-          .map((res) => res.error)
-          .flat();
-          
-        if (errors.length > 0) {
-          const uniqueErrors = new Set();
-          errors.forEach((error) =>
-            uniqueErrors.add(`[${error.type}]: ${error.message}`)
-          );
+    const { data = [], error = null } = await Promise.all(
+      embeddingRequests
+    ).then((results) => {
+      // If any errors were returned from OpenAI abort the entire sequence because the embeddings
+      // will be incomplete.
+      const errors = results
+        .filter((res) => !!res.error)
+        .map((res) => res.error)
+        .flat();
+      if (errors.length > 0) {
+        let uniqueErrors = new Set();
+        errors.map((error) =>
+          uniqueErrors.add(`[${error.type}]: ${error.message}`)
+        );
 
-          return {
-            data: [],
-            error: Array.from(uniqueErrors).join(", "),
-          };
-        }
-        
         return {
-          data: results.map((res) => res?.data || []).flat(),
-          error: null,
+          data: [],
+          error: Array.from(uniqueErrors).join(", "),
         };
-      });
+      }
+      return {
+        data: results.map((res) => res?.data || []).flat(),
+        error: null,
+      };
+    });
 
-    if (!!error) {
-      throw new Error(`OpenAI Failed to embed: ${error}`);
-    }
-    
-    // Ensure we have valid embeddings before returning
-    return data.length > 0 && data.every((embd) => embd.hasOwnProperty("embedding"))
+    if (!!error) throw new Error(`Gemini Failed to embed: ${error}`);
+    return data.length > 0 &&
+      data.every((embd) => embd.hasOwnProperty("embedding"))
       ? data.map((embd) => embd.embedding)
       : null;
   }