From e5c77f63ced7b4633b0e157acb87a6161171f3f2 Mon Sep 17 00:00:00 2001 From: Chaiwat Saithongcum Date: Tue, 24 Dec 2024 14:21:29 +0700 Subject: [PATCH 1/3] Add support for Google Generative AI (Gemini) embedder --- .../GeminiOptions/index.jsx | 50 +++++++++++++++++++ .../EmbeddingPreference/index.jsx | 9 ++++ server/.env.example | 4 ++ server/utils/EmbeddingEngines/gemini/index.js | 40 +++++++++++++++ server/utils/helpers/index.js | 3 ++ server/utils/helpers/updateENV.js | 1 + 6 files changed, 107 insertions(+) create mode 100644 frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx create mode 100644 server/utils/EmbeddingEngines/gemini/index.js diff --git a/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx b/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx new file mode 100644 index 00000000000..03def80c239 --- /dev/null +++ b/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx @@ -0,0 +1,50 @@ +export default function GeminiOptions({ settings }) { + return ( +
+
+
+ + +
+
+ + +
+
+
+ ); +} diff --git a/frontend/src/pages/GeneralSettings/EmbeddingPreference/index.jsx b/frontend/src/pages/GeneralSettings/EmbeddingPreference/index.jsx index 77853e0a999..a1833b0cc6f 100644 --- a/frontend/src/pages/GeneralSettings/EmbeddingPreference/index.jsx +++ b/frontend/src/pages/GeneralSettings/EmbeddingPreference/index.jsx @@ -6,6 +6,7 @@ import showToast from "@/utils/toast"; import AnythingLLMIcon from "@/media/logo/anything-llm-icon.png"; import OpenAiLogo from "@/media/llmprovider/openai.png"; import AzureOpenAiLogo from "@/media/llmprovider/azure.png"; +import GemeniAiLogo from "@/media/llmprovider/gemini.png"; import LocalAiLogo from "@/media/llmprovider/localai.png"; import OllamaLogo from "@/media/llmprovider/ollama.png"; import LMStudioLogo from "@/media/llmprovider/lmstudio.png"; @@ -19,6 +20,7 @@ import PreLoader from "@/components/Preloader"; import ChangeWarningModal from "@/components/ChangeWarning"; import OpenAiOptions from "@/components/EmbeddingSelection/OpenAiOptions"; import AzureAiOptions from "@/components/EmbeddingSelection/AzureAiOptions"; +import GeminiOptions from "@/components/EmbeddingSelection/GeminiOptions"; import LocalAiOptions from "@/components/EmbeddingSelection/LocalAiOptions"; import NativeEmbeddingOptions from "@/components/EmbeddingSelection/NativeEmbeddingOptions"; import OllamaEmbeddingOptions from "@/components/EmbeddingSelection/OllamaOptions"; @@ -59,6 +61,13 @@ const EMBEDDERS = [ options: (settings) => , description: "The enterprise option of OpenAI hosted on Azure services.", }, + { + name: "Gemini", + value: "gemini", + logo: GemeniAiLogo, + options: (settings) => , + description: "Run powerful embedding models from Google AI.", + }, { name: "Local AI", value: "localai", diff --git a/server/.env.example b/server/.env.example index 3346fc397d5..b8e538c2464 100644 --- a/server/.env.example +++ b/server/.env.example @@ -128,6 +128,10 @@ SIG_SALT='salt' # Please generate random string at least 32 chars long. # AZURE_OPENAI_KEY= # EMBEDDING_MODEL_PREF='my-embedder-model' # This is the "deployment" on Azure you want to use for embeddings. Not the base model. Valid base model is text-embedding-ada-002 +# EMBEDDING_ENGINE='gemini' +# GEMINI_API_KEY= +# EMBEDDING_MODEL_PREF='text-embedding-004' + # EMBEDDING_ENGINE='localai' # EMBEDDING_BASE_PATH='http://localhost:8080/v1' # EMBEDDING_MODEL_PREF='text-embedding-ada-002' diff --git a/server/utils/EmbeddingEngines/gemini/index.js b/server/utils/EmbeddingEngines/gemini/index.js new file mode 100644 index 00000000000..748383cc699 --- /dev/null +++ b/server/utils/EmbeddingEngines/gemini/index.js @@ -0,0 +1,40 @@ +const { toChunks } = require("../../helpers"); + +class geminiEmbedder { + constructor() { + if (!process.env.GEMINI_API_KEY) throw new Error("No Gemini API key was set."); + const { GoogleGenerativeAI: GenerativeAI } = require("@google/generative-ai"); + this.gemini = new GenerativeAI(process.env.GEMINI_API_KEY); + this.model = process.env.EMBEDDING_MODEL_PREF || "text-embedding-004"; + + // Limit of how many strings we can process in a single pass to stay with resource or network limits + this.maxConcurrentChunks = 1; // Gemini's limit per request is 1 + + // https://ai.google.dev/gemini-api/docs/models/gemini#text-embedding-and-embedding + this.embeddingMaxChunkLength = 2_048; + } + + async embedTextInput(textInput) { + const genAI = this.gemini; + const model = genAI.getGenerativeModel({ model: this.model }); + + const result = await model.embedContent(textInput); + return result.embedding.values || []; + } + + async embedChunks(textChunks = []) { + const genAI = this.gemini; + const model = genAI.getGenerativeModel({ model: this.model }); + + const embeddings = []; + for (const chunk of toChunks(textChunks, this.maxConcurrentChunks)) { + const result = await model.embedContent(chunk); + embeddings.push(result.embedding.values); + } + return embeddings || []; + } +} + +module.exports = { + geminiEmbedder, +}; diff --git a/server/utils/helpers/index.js b/server/utils/helpers/index.js index 55d190f4fdb..be17625db88 100644 --- a/server/utils/helpers/index.js +++ b/server/utils/helpers/index.js @@ -222,6 +222,9 @@ function getEmbeddingEngineSelection() { AzureOpenAiEmbedder, } = require("../EmbeddingEngines/azureOpenAi"); return new AzureOpenAiEmbedder(); + case "gemini": + const { geminiEmbedder } = require("../EmbeddingEngines/gemini"); + return new geminiEmbedder(); case "localai": const { LocalAiEmbedder } = require("../EmbeddingEngines/localAi"); return new LocalAiEmbedder(); diff --git a/server/utils/helpers/updateENV.js b/server/utils/helpers/updateENV.js index da30b6ee0dd..dab2c82c98a 100644 --- a/server/utils/helpers/updateENV.js +++ b/server/utils/helpers/updateENV.js @@ -759,6 +759,7 @@ function supportedEmbeddingModel(input = "") { const supported = [ "openai", "azure", + "gemini", "localai", "native", "ollama", From 7d022fe3465e5fa96c91670e74d8e5cf2f1e8f02 Mon Sep 17 00:00:00 2001 From: timothycarambat Date: Tue, 31 Dec 2024 09:25:57 -0800 Subject: [PATCH 2/3] Add missing example in docker Fix UI key elements in options Add Gemini to data handling section Patch issues with chunk handling during embedding --- docker/.env.example | 4 ++ .../GeminiOptions/index.jsx | 11 ++-- .../Steps/DataHandling/index.jsx | 8 +++ server/.env.example | 4 ++ server/models/systemSettings.js | 5 +- server/utils/EmbeddingEngines/gemini/index.js | 56 ++++++++++++------- server/utils/helpers/index.js | 6 +- server/utils/helpers/updateENV.js | 6 ++ 8 files changed, 66 insertions(+), 34 deletions(-) diff --git a/docker/.env.example b/docker/.env.example index ee53c718bc6..19c04dfa855 100644 --- a/docker/.env.example +++ b/docker/.env.example @@ -170,6 +170,10 @@ GID='1000' # GENERIC_OPEN_AI_EMBEDDING_API_KEY='sk-123abc' # GENERIC_OPEN_AI_EMBEDDING_MAX_CONCURRENT_CHUNKS=500 +# EMBEDDING_ENGINE='gemini' +# GEMINI_EMBEDDING_API_KEY= +# EMBEDDING_MODEL_PREF='text-embedding-004' + ########################################### ######## Vector Database Selection ######## ########################################### diff --git a/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx b/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx index 03def80c239..c25a2b13f8d 100644 --- a/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx +++ b/frontend/src/components/EmbeddingSelection/GeminiOptions/index.jsx @@ -8,10 +8,10 @@ export default function GeminiOptions({ settings }) { - {[ - "text-embedding-004", - "embedding-001", - ].map((model) => { + {["text-embedding-004"].map((model) => { return (