From 65b7236899c27a1cbb9b136bdc94ae4f32e412b2 Mon Sep 17 00:00:00 2001
From: Jopo-JP <21254390+Jopo-JP@users.noreply.github.com>
Date: Sun, 13 Jul 2025 15:49:40 +0200
Subject: [PATCH 01/10] feat: Add support for new gemini embedding model.
(#5621)
---
src/core/webview/webviewMessageHandler.ts | 25 +++++++------
.../embedders/__tests__/gemini.spec.ts | 36 +++++++++++++++++++
src/services/code-index/embedders/gemini.ts | 11 ++++--
.../code-index/embedders/openai-compatible.ts | 29 ++++++++++-----
.../code-index/interfaces/embedder.ts | 2 +-
src/shared/embeddingModels.ts | 3 ++
6 files changed, 84 insertions(+), 22 deletions(-)
diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts
index e70b39df8f..283ba7da6c 100644
--- a/src/core/webview/webviewMessageHandler.ts
+++ b/src/core/webview/webviewMessageHandler.ts
@@ -2065,11 +2065,14 @@ export const webviewMessageHandler = async (
}
case "requestIndexingStatus": {
- const status = provider.codeIndexManager!.getCurrentStatus()
- provider.postMessageToWebview({
- type: "indexingStatusUpdate",
- values: status,
- })
+ const manager = provider.codeIndexManager
+ if (manager) {
+ const status = manager.getCurrentStatus()
+ provider.postMessageToWebview({
+ type: "indexingStatusUpdate",
+ values: status,
+ })
+ }
break
}
case "requestCodeIndexSecretStatus": {
@@ -2094,8 +2097,8 @@ export const webviewMessageHandler = async (
}
case "startIndexing": {
try {
- const manager = provider.codeIndexManager!
- if (manager.isFeatureEnabled && manager.isFeatureConfigured) {
+ const manager = provider.codeIndexManager
+ if (manager && manager.isFeatureEnabled && manager.isFeatureConfigured) {
if (!manager.isInitialized) {
await manager.initialize(provider.contextProxy)
}
@@ -2109,9 +2112,11 @@ export const webviewMessageHandler = async (
}
case "clearIndexData": {
try {
- const manager = provider.codeIndexManager!
- await manager.clearIndexData()
- provider.postMessageToWebview({ type: "indexCleared", values: { success: true } })
+ const manager = provider.codeIndexManager
+ if (manager) {
+ await manager.clearIndexData()
+ provider.postMessageToWebview({ type: "indexCleared", values: { success: true } })
+ }
} catch (error) {
provider.log(`Error clearing index data: ${error instanceof Error ? error.message : String(error)}`)
provider.postMessageToWebview({
diff --git a/src/services/code-index/embedders/__tests__/gemini.spec.ts b/src/services/code-index/embedders/__tests__/gemini.spec.ts
index 378e6e7d95..7a35a40d1c 100644
--- a/src/services/code-index/embedders/__tests__/gemini.spec.ts
+++ b/src/services/code-index/embedders/__tests__/gemini.spec.ts
@@ -114,4 +114,40 @@ describe("GeminiEmbedder", () => {
await expect(embedder.validateConfiguration()).rejects.toThrow("Validation failed")
})
})
+
+ describe("createEmbeddings", () => {
+ let mockCreateEmbeddings: any
+
+ beforeEach(() => {
+ mockCreateEmbeddings = vitest.fn()
+ MockedOpenAICompatibleEmbedder.prototype.createEmbeddings = mockCreateEmbeddings
+ embedder = new GeminiEmbedder("test-api-key")
+ })
+
+ it("should use default model when none is provided", async () => {
+ // Arrange
+ const texts = ["text1", "text2"]
+ mockCreateEmbeddings.mockResolvedValue({ embeddings: [], usage: { promptTokens: 0, totalTokens: 0 } })
+
+ // Act
+ await embedder.createEmbeddings(texts)
+
+ // Assert
+ expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "text-embedding-004", undefined)
+ })
+
+ it("should pass model and dimension to the OpenAICompatibleEmbedder", async () => {
+ // Arrange
+ const texts = ["text1", "text2"]
+ const model = "custom-model"
+ const options = { dimension: 1536 }
+ mockCreateEmbeddings.mockResolvedValue({ embeddings: [], usage: { promptTokens: 0, totalTokens: 0 } })
+
+ // Act
+ await embedder.createEmbeddings(texts, model, options)
+
+ // Assert
+ expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, model, options)
+ })
+ })
})
diff --git a/src/services/code-index/embedders/gemini.ts b/src/services/code-index/embedders/gemini.ts
index fcca4c0fda..b1581c852f 100644
--- a/src/services/code-index/embedders/gemini.ts
+++ b/src/services/code-index/embedders/gemini.ts
@@ -44,10 +44,15 @@ export class GeminiEmbedder implements IEmbedder {
* @param model Optional model identifier (ignored - always uses text-embedding-004)
* @returns Promise resolving to embedding response
*/
- async createEmbeddings(texts: string[], model?: string): Promise {
+ async createEmbeddings(
+ texts: string[],
+ model?: string,
+ options?: { dimension?: number },
+ ): Promise {
try {
- // Always use the fixed Gemini model, ignoring any passed model parameter
- return await this.openAICompatibleEmbedder.createEmbeddings(texts, GeminiEmbedder.GEMINI_MODEL)
+ // Use the provided model or the fixed Gemini model
+ const modelToUse = model || GeminiEmbedder.GEMINI_MODEL
+ return await this.openAICompatibleEmbedder.createEmbeddings(texts, modelToUse, options)
} catch (error) {
TelemetryService.instance.captureEvent(TelemetryEventName.CODE_INDEX_ERROR, {
error: error instanceof Error ? error.message : String(error),
diff --git a/src/services/code-index/embedders/openai-compatible.ts b/src/services/code-index/embedders/openai-compatible.ts
index d882e78313..205675e9a4 100644
--- a/src/services/code-index/embedders/openai-compatible.ts
+++ b/src/services/code-index/embedders/openai-compatible.ts
@@ -71,7 +71,11 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
* @param model Optional model identifier
* @returns Promise resolving to embedding response
*/
- async createEmbeddings(texts: string[], model?: string): Promise {
+ async createEmbeddings(
+ texts: string[],
+ model?: string,
+ options?: { dimension?: number },
+ ): Promise {
const modelToUse = model || this.defaultModelId
// Apply model-specific query prefix if required
@@ -139,7 +143,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
}
if (currentBatch.length > 0) {
- const batchResult = await this._embedBatchWithRetries(currentBatch, modelToUse)
+ const batchResult = await this._embedBatchWithRetries(currentBatch, modelToUse, options)
allEmbeddings.push(...batchResult.embeddings)
usage.promptTokens += batchResult.usage.promptTokens
usage.totalTokens += batchResult.usage.totalTokens
@@ -181,7 +185,18 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
url: string,
batchTexts: string[],
model: string,
+ options?: { dimension?: number },
): Promise {
+ const body: Record = {
+ input: batchTexts,
+ model: model,
+ encoding_format: "base64",
+ }
+
+ if (options?.dimension) {
+ body.dimensions = options.dimension
+ }
+
const response = await fetch(url, {
method: "POST",
headers: {
@@ -191,11 +206,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
"api-key": this.apiKey,
Authorization: `Bearer ${this.apiKey}`,
},
- body: JSON.stringify({
- input: batchTexts,
- model: model,
- encoding_format: "base64",
- }),
+ body: JSON.stringify(body),
})
if (!response || !response.ok) {
@@ -234,6 +245,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
private async _embedBatchWithRetries(
batchTexts: string[],
model: string,
+ options?: { dimension?: number },
): Promise<{ embeddings: number[][]; usage: { promptTokens: number; totalTokens: number } }> {
// Use cached value for performance
const isFullUrl = this.isFullUrl
@@ -244,7 +256,7 @@ export class OpenAICompatibleEmbedder implements IEmbedder {
if (isFullUrl) {
// Use direct HTTP request for full endpoint URLs
- response = await this.makeDirectEmbeddingRequest(this.baseUrl, batchTexts, model)
+ response = await this.makeDirectEmbeddingRequest(this.baseUrl, batchTexts, model, options)
} else {
// Use OpenAI SDK for base URLs
response = (await this.embeddingsClient.embeddings.create({
@@ -254,6 +266,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",
+ ...(options?.dimension && { dimensions: options.dimension }),
})) as OpenAIEmbeddingResponse
}
diff --git a/src/services/code-index/interfaces/embedder.ts b/src/services/code-index/interfaces/embedder.ts
index 0a74446d5e..bb7c5283bb 100644
--- a/src/services/code-index/interfaces/embedder.ts
+++ b/src/services/code-index/interfaces/embedder.ts
@@ -9,7 +9,7 @@ export interface IEmbedder {
* @param model Optional model ID to use for embeddings
* @returns Promise resolving to an EmbeddingResponse
*/
- createEmbeddings(texts: string[], model?: string): Promise
+ createEmbeddings(texts: string[], model?: string, options?: { dimension?: number }): Promise
/**
* Validates the embedder configuration by testing connectivity and credentials.
diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts
index 4c6bc24319..e70ba23f53 100644
--- a/src/shared/embeddingModels.ts
+++ b/src/shared/embeddingModels.ts
@@ -48,6 +48,9 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
},
gemini: {
"text-embedding-004": { dimension: 768 },
+ // ADD: New model with a default dimension.
+ // The actual dimension will be passed from the configuration at runtime.
+ "gemini-embedding-exp-03-07": { dimension: 768 },
},
}
From 761b2ea8423a13708e927e145507c67751ab1781 Mon Sep 17 00:00:00 2001
From: Jopo-JP <21254390+Jopo-JP@users.noreply.github.com>
Date: Tue, 15 Jul 2025 12:39:37 +0200
Subject: [PATCH 02/10] feat: Enhance Gemini embedder with configurable
dimensions and validation for embedding models
---
packages/types/src/codebase-index.ts | 21 +++-
src/shared/embeddingModels.ts | 4 +
.../src/components/chat/CodeIndexPopover.tsx | 109 ++++++++++++++++--
3 files changed, 121 insertions(+), 13 deletions(-)
diff --git a/packages/types/src/codebase-index.ts b/packages/types/src/codebase-index.ts
index 0ad19d8676..d9922a53ed 100644
--- a/packages/types/src/codebase-index.ts
+++ b/packages/types/src/codebase-index.ts
@@ -42,11 +42,24 @@ export type CodebaseIndexConfig = z.infer
* CodebaseIndexModels
*/
+const modelProfileSchema = z.object({
+ /** The fixed dimension for the model, or a fallback for models with variable dimensions. */
+ dimension: z.number(),
+ scoreThreshold: z.number().optional(),
+ queryPrefix: z.string().optional(),
+ /** The minimum dimension supported by a variable-dimension model. */
+ minDimension: z.number().optional(),
+ /** The maximum dimension supported by a variable-dimension model. */
+ maxDimension: z.number().optional(),
+ /** The default dimension for a variable-dimension model, used for UI presentation. */
+ defaultDimension: z.number().optional(),
+})
+
export const codebaseIndexModelsSchema = z.object({
- openai: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
- ollama: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
- "openai-compatible": z.record(z.string(), z.object({ dimension: z.number() })).optional(),
- gemini: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
+ openai: z.record(z.string(), modelProfileSchema).optional(),
+ ollama: z.record(z.string(), modelProfileSchema).optional(),
+ "openai-compatible": z.record(z.string(), modelProfileSchema).optional(),
+ gemini: z.record(z.string(), modelProfileSchema).optional(),
})
export type CodebaseIndexModels = z.infer
diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts
index f387480c65..0a6890bf47 100644
--- a/src/shared/embeddingModels.ts
+++ b/src/shared/embeddingModels.ts
@@ -5,9 +5,13 @@
export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" // Add other providers as needed
export interface EmbeddingModelProfile {
+ /** The fixed dimension for the model, or a fallback for models with variable dimensions. */
dimension: number
scoreThreshold?: number // Model-specific minimum score threshold for semantic search
queryPrefix?: string // Optional prefix required by the model for queries
+ minDimension?: number // The minimum dimension supported by a variable-dimension model.
+ maxDimension?: number // The maximum dimension supported by a variable-dimension model.
+ defaultDimension?: number // The default dimension for a variable-dimension model, used for UI presentation.
// Add other model-specific properties if needed, e.g., context window size
}
diff --git a/webview-ui/src/components/chat/CodeIndexPopover.tsx b/webview-ui/src/components/chat/CodeIndexPopover.tsx
index b5742cc623..be379ffb69 100644
--- a/webview-ui/src/components/chat/CodeIndexPopover.tsx
+++ b/webview-ui/src/components/chat/CodeIndexPopover.tsx
@@ -71,7 +71,7 @@ interface LocalCodeIndexSettings {
}
// Validation schema for codebase index settings
-const createValidationSchema = (provider: EmbedderProvider, t: any) => {
+const createValidationSchema = (provider: EmbedderProvider, t: any, models: any) => {
const baseSchema = z.object({
codebaseIndexEnabled: z.boolean(),
codebaseIndexQdrantUrl: z
@@ -115,12 +115,32 @@ const createValidationSchema = (provider: EmbedderProvider, t: any) => {
})
case "gemini":
- return baseSchema.extend({
- codebaseIndexGeminiApiKey: z.string().min(1, t("settings:codeIndex.validation.geminiApiKeyRequired")),
- codebaseIndexEmbedderModelId: z
- .string()
- .min(1, t("settings:codeIndex.validation.modelSelectionRequired")),
- })
+ return baseSchema
+ .extend({
+ codebaseIndexGeminiApiKey: z
+ .string()
+ .min(1, t("settings:codeIndex.validation.geminiApiKeyRequired")),
+ codebaseIndexEmbedderModelId: z
+ .string()
+ .min(1, t("settings:codeIndex.validation.modelSelectionRequired")),
+ codebaseIndexEmbedderModelDimension: z.number().optional(),
+ })
+ .refine(
+ (data) => {
+ const model = models?.gemini?.[data.codebaseIndexEmbedderModelId || ""]
+ if (model?.minDimension && model?.maxDimension && data.codebaseIndexEmbedderModelDimension) {
+ return (
+ data.codebaseIndexEmbedderModelDimension >= model.minDimension &&
+ data.codebaseIndexEmbedderModelDimension <= model.maxDimension
+ )
+ }
+ return true
+ },
+ {
+ message: t("settings:codeIndex.validation.invalidDimension"),
+ path: ["codebaseIndexEmbedderModelDimension"],
+ },
+ )
default:
return baseSchema
@@ -188,7 +208,13 @@ export const CodeIndexPopover: React.FC = ({
codebaseIndexEmbedderBaseUrl: codebaseIndexConfig.codebaseIndexEmbedderBaseUrl || "",
codebaseIndexEmbedderModelId: codebaseIndexConfig.codebaseIndexEmbedderModelId || "",
codebaseIndexEmbedderModelDimension:
- codebaseIndexConfig.codebaseIndexEmbedderModelDimension || undefined,
+ // This order is critical to prevent a UI race condition. After saving,
+ // the component's local `currentSettings` is updated immediately, while
+ // the global `codebaseIndexConfig` might still be stale. Prioritizing
+ // `currentSettings` ensures the UI reflects the saved value instantly.
+ currentSettings.codebaseIndexEmbedderModelDimension ||
+ codebaseIndexConfig.codebaseIndexEmbedderModelDimension ||
+ undefined,
codebaseIndexSearchMaxResults:
codebaseIndexConfig.codebaseIndexSearchMaxResults ?? CODEBASE_INDEX_DEFAULTS.DEFAULT_SEARCH_RESULTS,
codebaseIndexSearchMinScore:
@@ -349,7 +375,7 @@ export const CodeIndexPopover: React.FC = ({
// Validation function
const validateSettings = (): boolean => {
- const schema = createValidationSchema(currentSettings.codebaseIndexEmbedderProvider, t)
+ const schema = createValidationSchema(currentSettings.codebaseIndexEmbedderProvider, t, codebaseIndexModels)
// Prepare data for validation
const dataToValidate: any = {}
@@ -916,6 +942,71 @@ export const CodeIndexPopover: React.FC = ({
)}
+ {(() => {
+ const selectedModelProfile =
+ codebaseIndexModels?.gemini?.[
+ currentSettings.codebaseIndexEmbedderModelId
+ ]
+
+ // Conditionally render the dimension slider only for Gemini models
+ // that explicitly define a min and max dimension in their profile.
+ if (
+ selectedModelProfile?.minDimension &&
+ selectedModelProfile?.maxDimension
+ ) {
+ return (
+