From c0f88e05cdab0370313b22273df9d2935027439d Mon Sep 17 00:00:00 2001 From: Roo Code Date: Tue, 9 Sep 2025 06:14:49 +0000 Subject: [PATCH] feat: add support for Google Gemini embedding model gemini-embedding-exp-03-07 - Add gemini-embedding-exp-03-07 model variants with flexible dimensions (3072, 1536, 768) - Update GeminiEmbedder to support outputDimension parameter - Update OpenAICompatibleEmbedder to handle outputDimension in API calls - Set gemini-embedding-exp-03-07-3072 as new default model for Gemini provider - Preserve backward compatibility with existing text-embedding-004 model - Update tests to cover new model and outputDimension functionality Fixes #5621 --- .../embedders/__tests__/gemini.spec.ts | 25 +++++++++++++++++-- src/services/code-index/embedders/gemini.ts | 13 +++++++--- .../code-index/embedders/openai-compatible.ts | 8 +++++- src/shared/embeddingModels.ts | 5 +++- 4 files changed, 43 insertions(+), 8 deletions(-) diff --git a/src/services/code-index/embedders/__tests__/gemini.spec.ts b/src/services/code-index/embedders/__tests__/gemini.spec.ts index d41a4dc1e9..2ac694e5df 100644 --- a/src/services/code-index/embedders/__tests__/gemini.spec.ts +++ b/src/services/code-index/embedders/__tests__/gemini.spec.ts @@ -36,8 +36,9 @@ describe("GeminiEmbedder", () => { expect(MockedOpenAICompatibleEmbedder).toHaveBeenCalledWith( "https://generativelanguage.googleapis.com/v1beta/openai/", apiKey, - "gemini-embedding-001", + "gemini-embedding-exp-03-07-3072", 2048, + undefined, ) }) @@ -55,6 +56,26 @@ describe("GeminiEmbedder", () => { apiKey, "text-embedding-004", 2048, + undefined, + ) + }) + + it("should create an instance with specified model and outputDimension", () => { + // Arrange + const apiKey = "test-gemini-api-key" + const modelId = "gemini-embedding-exp-03-07-1536" + const outputDimension = 1536 + + // Act + embedder = new GeminiEmbedder(apiKey, modelId, outputDimension) + + // Assert + expect(MockedOpenAICompatibleEmbedder).toHaveBeenCalledWith( + "https://generativelanguage.googleapis.com/v1beta/openai/", + apiKey, + "gemini-embedding-exp-03-07-1536", + 2048, + 1536, ) }) @@ -104,7 +125,7 @@ describe("GeminiEmbedder", () => { const result = await embedder.createEmbeddings(texts) // Assert - expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "gemini-embedding-001") + expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "gemini-embedding-exp-03-07-3072") expect(result).toEqual(mockResponse) }) diff --git a/src/services/code-index/embedders/gemini.ts b/src/services/code-index/embedders/gemini.ts index 7e795875c9..bd590899f5 100644 --- a/src/services/code-index/embedders/gemini.ts +++ b/src/services/code-index/embedders/gemini.ts @@ -11,20 +11,24 @@ import { TelemetryService } from "@roo-code/telemetry" * * Supported models: * - text-embedding-004 (dimension: 768) - * - gemini-embedding-001 (dimension: 2048) + * - gemini-embedding-001 (dimension: 3072) + * - gemini-embedding-exp-03-07-3072 (dimension: 3072) + * - gemini-embedding-exp-03-07-1536 (dimension: 1536) + * - gemini-embedding-exp-03-07-768 (dimension: 768) */ export class GeminiEmbedder implements IEmbedder { private readonly openAICompatibleEmbedder: OpenAICompatibleEmbedder private static readonly GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/openai/" - private static readonly DEFAULT_MODEL = "gemini-embedding-001" + private static readonly DEFAULT_MODEL = "gemini-embedding-exp-03-07-3072" private readonly modelId: string /** * Creates a new Gemini embedder * @param apiKey The Gemini API key for authentication - * @param modelId The model ID to use (defaults to gemini-embedding-001) + * @param modelId The model ID to use (defaults to gemini-embedding-exp-03-07-3072) + * @param outputDimension Optional output dimension for flexible models */ - constructor(apiKey: string, modelId?: string) { + constructor(apiKey: string, modelId?: string, outputDimension?: number) { if (!apiKey) { throw new Error(t("embeddings:validation.apiKeyRequired")) } @@ -38,6 +42,7 @@ export class GeminiEmbedder implements IEmbedder { apiKey, this.modelId, GEMINI_MAX_ITEM_TOKENS, + outputDimension, ) } diff --git a/src/services/code-index/embedders/openai-compatible.ts b/src/services/code-index/embedders/openai-compatible.ts index 06c4ba5282..53ab245566 100644 --- a/src/services/code-index/embedders/openai-compatible.ts +++ b/src/services/code-index/embedders/openai-compatible.ts @@ -38,6 +38,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder { private readonly apiKey: string private readonly isFullUrl: boolean private readonly maxItemTokens: number + private readonly outputDimension?: number // Global rate limiting state shared across all instances private static globalRateLimitState = { @@ -55,8 +56,9 @@ export class OpenAICompatibleEmbedder implements IEmbedder { * @param apiKey The API key for authentication * @param modelId Optional model identifier (defaults to "text-embedding-3-small") * @param maxItemTokens Optional maximum tokens per item (defaults to MAX_ITEM_TOKENS) + * @param outputDimension Optional output dimension for flexible models */ - constructor(baseUrl: string, apiKey: string, modelId?: string, maxItemTokens?: number) { + constructor(baseUrl: string, apiKey: string, modelId?: string, maxItemTokens?: number, outputDimension?: number) { if (!baseUrl) { throw new Error(t("embeddings:validation.baseUrlRequired")) } @@ -74,6 +76,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder { // Cache the URL type check for performance this.isFullUrl = this.isFullEndpointUrl(baseUrl) this.maxItemTokens = maxItemTokens || MAX_ITEM_TOKENS + this.outputDimension = outputDimension } /** @@ -208,6 +211,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder { input: batchTexts, model: model, encoding_format: "base64", + ...(this.outputDimension && { dimensions: this.outputDimension }), }), }) @@ -270,6 +274,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder { // when processing numeric arrays, which breaks compatibility with models using larger dimensions. // By requesting base64 encoding, we bypass the package's parser and handle decoding ourselves. encoding_format: "base64", + ...(this.outputDimension && { dimensions: this.outputDimension }), })) as OpenAIEmbeddingResponse } @@ -369,6 +374,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder { input: testTexts, model: modelToUse, encoding_format: "base64", + ...(this.outputDimension && { dimensions: this.outputDimension }), })) as OpenAIEmbeddingResponse } diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts index 80c51a6b45..cee53993d4 100644 --- a/src/shared/embeddingModels.ts +++ b/src/shared/embeddingModels.ts @@ -49,6 +49,9 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = { gemini: { "text-embedding-004": { dimension: 768 }, "gemini-embedding-001": { dimension: 3072, scoreThreshold: 0.4 }, + "gemini-embedding-exp-03-07-3072": { dimension: 3072, scoreThreshold: 0.4 }, + "gemini-embedding-exp-03-07-1536": { dimension: 1536, scoreThreshold: 0.4 }, + "gemini-embedding-exp-03-07-768": { dimension: 768, scoreThreshold: 0.4 }, }, mistral: { "codestral-embed-2505": { dimension: 1536, scoreThreshold: 0.4 }, @@ -155,7 +158,7 @@ export function getDefaultModelId(provider: EmbedderProvider): string { } case "gemini": - return "gemini-embedding-001" + return "gemini-embedding-exp-03-07-3072" case "mistral": return "codestral-embed-2505"