Add gemini-embedding-exp-03-07 model support (#3767)

* adds gemini-embedding-exp-03-07 support

* unset random changes, make model map

---------

Co-authored-by: Timothy Carambat <rambat1010@gmail.com>
This commit is contained in:
Sebastián Darío Ramírez 2025-05-05 12:37:01 -04:00 committed by GitHub
parent c402e5cea1
commit 4f23e44756
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 33 additions and 11 deletions

View File

@ -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"].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>
);
})}

View File

@ -1,5 +1,11 @@
const { toChunks } = require("../../helpers");
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)
@ -16,9 +22,10 @@ class GeminiEmbedder {
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}`);
this.embeddingMaxChunkLength = MODEL_MAP[this.model] || 2_048;
this.log(
`Initialized with ${this.model} - Max Size: ${this.embeddingMaxChunkLength}`
);
}
log(text, ...args) {
@ -27,7 +34,7 @@ class GeminiEmbedder {
/**
* 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 +46,7 @@ class GeminiEmbedder {
/**
* Embeds a list of text inputs
* @param {Array<string>} textInputs - 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 = []) {
@ -98,7 +105,7 @@ class GeminiEmbedder {
};
});
if (!!error) throw new Error(`OpenAI Failed to embed: ${error}`);
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)