set embedder output dimensions for LocalAI and Gemini (gemini-embedding-001) (#4980)

frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx

@@ -1,3 +1,6 @@
+import { Info } from "@phosphor-icons/react";
+import { Tooltip } from "react-tooltip";
+
 const DEFAULT_MODELS = [
   {
     id: "gemini-embedding-001",
@@ -7,6 +10,7 @@ const DEFAULT_MODELS = [
 
 export default function GeminiOptions({ settings }) {
   return (
+    <div className="w-full flex flex-col gap-y-6">
       <div className="w-full flex flex-col gap-y-4">
         <div className="w-full flex items-center gap-[36px] mt-1.5">
           <div className="flex flex-col w-60">
@@ -18,7 +22,9 @@ export default function GeminiOptions({ settings }) {
               name="GeminiEmbeddingApiKey"
               className="border-none bg-theme-settings-input-bg text-white placeholder:text-theme-settings-input-placeholder text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-full p-2.5"
               placeholder="Gemini API Key"
-            defaultValue={settings?.GeminiEmbeddingApiKey ? "*".repeat(20) : ""}
+              defaultValue={
+                settings?.GeminiEmbeddingApiKey ? "*".repeat(20) : ""
+              }
               required={true}
               autoComplete="off"
               spellCheck={false}
@@ -50,5 +56,48 @@ export default function GeminiOptions({ settings }) {
           </div>
         </div>
       </div>
+      <div className="flex flex-col w-60">
+        <div
+          data-tooltip-id="embedding-output-dimensions-tooltip"
+          className="flex gap-x-1 items-center mb-3"
+        >
+          <label className="text-white text-sm font-semibold block">
+            Output dimensions
+          </label>
+          <Info
+            size={16}
+            className="text-theme-text-secondary cursor-pointer"
+          />
+          <Tooltip
+            id="embedding-output-dimensions-tooltip"
+            place="top"
+            delayShow={300}
+            className="tooltip !text-xs !opacity-100"
+            style={{
+              maxWidth: "250px",
+              whiteSpace: "normal",
+              wordWrap: "break-word",
+            }}
+          >
+            The number of dimensions the resulting output embeddings should have
+            if it supports multiple dimensions output.
+            <br />
+            <br /> Leave blank to use the default dimensions for the selected
+            model.
+          </Tooltip>
+        </div>
+        <input
+          type="number"
+          name="EmbeddingOutputDimensions"
+          className="border-none bg-theme-settings-input-bg text-white placeholder:text-theme-settings-input-placeholder text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-full p-2.5"
+          placeholder="Assume default dimensions"
+          min={1}
+          onScroll={(e) => e.target.blur()}
+          defaultValue={settings?.EmbeddingOutputDimensions}
+          required={false}
+          autoComplete="off"
+        />
+      </div>
+    </div>
   );
 }

frontend/src/components/EmbeddingSelection/LocalAiOptions/index.jsx

@@ -30,20 +30,70 @@ export default function LocalAiOptions({ settings }) {
           apiKey={apiKey}
           basePath={basePath.value}
         />
+        <div className="flex flex-col w-60">
+          <div className="flex flex-col gap-y-1 mb-2">
+            <div className="flex gap-x-1 items-center">
+              <label className="text-white text-sm font-semibold block">
+                Local AI API Key
+              </label>
+              <Info
+                size={16}
+                data-tooltip-id="localai-api-key-tooltip"
+                className="text-theme-text-secondary cursor-pointer"
+              />
+              <Tooltip
+                id="localai-api-key-tooltip"
+                place="top"
+                delayShow={300}
+                className="tooltip !text-xs !opacity-100"
+                style={{
+                  maxWidth: "250px",
+                  whiteSpace: "normal",
+                  wordWrap: "break-word",
+                }}
+              >
+                The API key for the LocalAI server (if applicable).
+              </Tooltip>
+            </div>
+          </div>
+          <input
+            type="password"
+            name="LocalAiApiKey"
+            className="border-none bg-theme-settings-input-bg text-white placeholder:text-theme-settings-input-placeholder text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-full p-2.5"
+            placeholder="sk-mysecretkey"
+            defaultValue={settings?.LocalAiApiKey ? "*".repeat(20) : ""}
+            autoComplete="off"
+            spellCheck={false}
+            onChange={(e) => setApiKeyValue(e.target.value)}
+            onBlur={() => setApiKey(apiKeyValue)}
+          />
+        </div>
+      </div>
+      <div className="w-full flex items-center gap-[36px] mt-1.5">
         <div className="flex flex-col w-60">
           <div
             data-tooltip-place="top"
             data-tooltip-id="max-embedding-chunk-length-tooltip"
             className="flex gap-x-1 items-center mb-3"
           >
+            <label className="text-white text-sm font-semibold block">
+              Max embedding chunk length
+            </label>
             <Info
               size={16}
               className="text-theme-text-secondary cursor-pointer"
             />
-            <label className="text-white text-sm font-semibold block">
+            <Tooltip
-              Max embedding chunk length
+              id="max-embedding-chunk-length-tooltip"
-            </label>
+              place="top"
-            <Tooltip id="max-embedding-chunk-length-tooltip">
+              delayShow={300}
+              className="tooltip !text-xs !opacity-100"
+              style={{
+                maxWidth: "250px",
+                whiteSpace: "normal",
+                wordWrap: "break-word",
+              }}
+            >
               Maximum length of text chunks, in characters, for embedding.
             </Tooltip>
           </div>
@@ -59,23 +109,47 @@ export default function LocalAiOptions({ settings }) {
             autoComplete="off"
           />
         </div>
+
         <div className="flex flex-col w-60">
-          <div className="flex flex-col gap-y-1 mb-2">
+          <div
-            <label className="text-white text-sm font-semibold flex items-center gap-x-2">
+            data-tooltip-id="embedding-output-dimensions-tooltip"
-              Local AI API Key{" "}
+            className="flex gap-x-1 items-center mb-3"
-              <p className="!text-xs !italic !font-thin">optional</p>
+          >
+            <label className="text-white text-sm font-semibold block">
+              Output dimensions
             </label>
+            <Info
+              size={16}
+              className="text-theme-text-secondary cursor-pointer"
+            />
+            <Tooltip
+              id="embedding-output-dimensions-tooltip"
+              place="top"
+              delayShow={300}
+              className="tooltip !text-xs !opacity-100"
+              style={{
+                maxWidth: "250px",
+                whiteSpace: "normal",
+                wordWrap: "break-word",
+              }}
+            >
+              The number of dimensions the resulting output embeddings should
+              have if it supports multiple dimensions output.
+              <br />
+              <br /> Leave blank to use the default dimensions for the selected
+              model.
+            </Tooltip>
           </div>
           <input
-            type="password"
+            type="number"
-            name="LocalAiApiKey"
+            name="EmbeddingOutputDimensions"
             className="border-none bg-theme-settings-input-bg text-white placeholder:text-theme-settings-input-placeholder text-sm rounded-lg focus:outline-primary-button active:outline-primary-button outline-none block w-full p-2.5"
-            placeholder="sk-mysecretkey"
+            placeholder="Assume default dimensions"
-            defaultValue={settings?.LocalAiApiKey ? "*".repeat(20) : ""}
+            min={1}
+            onScroll={(e) => e.target.blur()}
+            defaultValue={settings?.EmbeddingOutputDimensions}
+            required={false}
             autoComplete="off"
-            spellCheck={false}
-            onChange={(e) => setApiKeyValue(e.target.value)}
-            onBlur={() => setApiKey(apiKeyValue)}
           />
         </div>
       </div>

server/models/systemSettings.js

@@ -233,6 +233,8 @@ const SystemSettings = {
         embeddingEngine === "native"
           ? NativeEmbedder._getEmbeddingModel()
           : process.env.EMBEDDING_MODEL_PREF,
+      EmbeddingOutputDimensions:
+        process.env.EMBEDDING_OUTPUT_DIMENSIONS || null,
       EmbeddingModelMaxChunkLength:
         process.env.EMBEDDING_MODEL_MAX_CHUNK_LENGTH,
       OllamaEmbeddingBatchSize: process.env.OLLAMA_EMBEDDING_BATCH_SIZE || 1,

server/utils/EmbeddingEngines/gemini/index.js

@@ -23,7 +23,10 @@ class GeminiEmbedder {
     // 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}`
+      `Initialized with ${this.model} - Max Size: ${this.embeddingMaxChunkLength}` +
+        (this.outputDimensions
+          ? ` - Output Dimensions: ${this.outputDimensions}`
+          : " Assuming default output dimensions")
     );
   }
 
@@ -31,6 +34,16 @@ class GeminiEmbedder {
     console.log(`\x1b[36m[${this.className}]\x1b[0m ${text}`, ...args);
   }
 
+  get outputDimensions() {
+    if (
+      process.env.EMBEDDING_OUTPUT_DIMENSIONS &&
+      !isNaN(process.env.EMBEDDING_OUTPUT_DIMENSIONS) &&
+      process.env.EMBEDDING_OUTPUT_DIMENSIONS > 0
+    )
+      return parseInt(process.env.EMBEDDING_OUTPUT_DIMENSIONS);
+    return null;
+  }
+
   /**
    * Embeds a single text input
    * @param {string|string[]} textInput - The text to embed
@@ -62,6 +75,7 @@ class GeminiEmbedder {
             .create({
               model: this.model,
               input: chunk,
+              dimensions: this.outputDimensions,
             })
             .then((result) => {
               resolve({ data: result?.data, error: null });

server/utils/EmbeddingEngines/localAi/index.js

@@ -7,7 +7,9 @@ class LocalAiEmbedder {
     if (!process.env.EMBEDDING_MODEL_PREF)
       throw new Error("No embedding model was set.");
 
+    this.className = "LocalAiEmbedder";
     const { OpenAI: OpenAIApi } = require("openai");
+    this.model = process.env.EMBEDDING_MODEL_PREF;
     this.openai = new OpenAIApi({
       baseURL: process.env.EMBEDDING_BASE_PATH,
       apiKey: process.env.LOCAL_AI_API_KEY ?? null,
@@ -16,6 +18,27 @@ class LocalAiEmbedder {
     // Limit of how many strings we can process in a single pass to stay with resource or network limits
     this.maxConcurrentChunks = 50;
     this.embeddingMaxChunkLength = maximumChunkLength();
+
+    this.log(
+      `Initialized with ${this.model} - Max Size: ${this.embeddingMaxChunkLength}` +
+        (this.outputDimensions
+          ? ` - Output Dimensions: ${this.outputDimensions}`
+          : " Assuming default output dimensions")
+    );
+  }
+
+  log(text, ...args) {
+    console.log(`\x1b[36m[${this.className}]\x1b[0m ${text}`, ...args);
+  }
+
+  get outputDimensions() {
+    if (
+      process.env.EMBEDDING_OUTPUT_DIMENSIONS &&
+      !isNaN(process.env.EMBEDDING_OUTPUT_DIMENSIONS) &&
+      process.env.EMBEDDING_OUTPUT_DIMENSIONS > 0
+    )
+      return parseInt(process.env.EMBEDDING_OUTPUT_DIMENSIONS);
+    return null;
   }
 
   async embedTextInput(textInput) {
@@ -32,8 +55,9 @@ class LocalAiEmbedder {
         new Promise((resolve) => {
           this.openai.embeddings
             .create({
-              model: process.env.EMBEDDING_MODEL_PREF,
+              model: this.model,
               input: chunk,
+              dimensions: this.outputDimensions,
             })
             .then((result) => {
               resolve({ data: result?.data, error: null });

server/utils/helpers/updateENV.js

@@ -307,6 +307,10 @@ const KEY_MAPPING = {
     envKey: "EMBEDDING_MODEL_MAX_CHUNK_LENGTH",
     checks: [nonZero],
   },
+  EmbeddingOutputDimensions: {
+    envKey: "EMBEDDING_OUTPUT_DIMENSIONS",
+    checks: [],
+  },
   OllamaEmbeddingBatchSize: {
     envKey: "OLLAMA_EMBEDDING_BATCH_SIZE",
     checks: [nonZero],