feat: Add support for new gemini embedding model. (#5621)

Author: Jopo-JP

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({

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)
+		})
+	})
 })

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<EmbeddingResponse> {
+	async createEmbeddings(
+		texts: string[],
+		model?: string,
+		options?: { dimension?: number },
+	): Promise<EmbeddingResponse> {
 		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),

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<EmbeddingResponse> {
+	async createEmbeddings(
+		texts: string[],
+		model?: string,
+		options?: { dimension?: number },
+	): Promise<EmbeddingResponse> {
 		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<OpenAIEmbeddingResponse> {
+		const body: Record<string, any> = {
+			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
 				}
 

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<EmbeddingResponse>
+	createEmbeddings(texts: string[], model?: string, options?: { dimension?: number }): Promise<EmbeddingResponse>
 
 	/**
 	 * Validates the embedder configuration by testing connectivity and credentials.

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 },
 	},
 }