feat: Enhance Gemini embedder with configurable dimensions and validation for embedding models

Author: Jopo-JP

src/services/code-index/embedders/__tests__/gemini.spec.ts

@@ -190,4 +190,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

@@ -47,11 +47,15 @@ export class GeminiEmbedder implements IEmbedder {
 	 * @param model Optional model identifier (uses constructor model if not provided)
 	 * @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 {
 			// Use the provided model or fall back to the instance's model
 			const modelToUse = model || this.modelId
-			return await this.openAICompatibleEmbedder.createEmbeddings(texts, modelToUse)
+			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

@@ -52,7 +52,13 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
 	},
 	gemini: {
 		"text-embedding-004": { dimension: 768 },
-		"gemini-embedding-001": { dimension: 3072, scoreThreshold: 0.4 },
+		"gemini-embedding-001": {
+			dimension: 3072, // Fallback, but defaultDimension is preferred
+			minDimension: 128,
+			maxDimension: 3072,
+			defaultDimension: 3072,
+			scoreThreshold: 0.4,
+		},
 	},
 }