Feat/gemini embedding model

packages/types/src/codebase-index.ts

@@ -42,12 +42,25 @@ export type CodebaseIndexConfig = z.infer<typeof codebaseIndexConfigSchema>
  * 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(),
-	mistral: 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(),
+	mistral: z.record(z.string(), modelProfileSchema).optional(),
 })
 
 export type CodebaseIndexModels = z.infer<typeof codebaseIndexModelsSchema>

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

@@ -104,7 +104,7 @@ describe("GeminiEmbedder", () => {
 				const result = await embedder.createEmbeddings(texts)
 
 				// Assert
-				expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "gemini-embedding-001")
+				expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "gemini-embedding-001", undefined)
 				expect(result).toEqual(mockResponse)
 			})
 
@@ -124,7 +124,7 @@ describe("GeminiEmbedder", () => {
 				const result = await embedder.createEmbeddings(texts, "gemini-embedding-001")
 
 				// Assert
-				expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "gemini-embedding-001")
+				expect(mockCreateEmbeddings).toHaveBeenCalledWith(texts, "gemini-embedding-001", undefined)
 				expect(result).toEqual(mockResponse)
 			})
 
@@ -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, "gemini-embedding-001", 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

@@ -82,7 +82,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
@@ -150,7 +154,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
@@ -192,7 +196,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: {
@@ -202,11 +217,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) {
@@ -245,6 +256,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
@@ -258,7 +270,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({
@@ -268,6 +280,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

@@ -5,9 +5,13 @@
 export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "mistral" // 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
 }
 
@@ -48,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,
+		},
 	},
 	mistral: {
 		"codestral-embed-2505": { dimension: 1536, scoreThreshold: 0.4 },