feat: Add geminiEmbeddingDimension to configuration and update related components

Author: ChuKhaLi

packages/types/src/codebase-index.ts

@@ -11,6 +11,7 @@ export const codebaseIndexConfigSchema = z.object({
 	codebaseIndexEmbedderBaseUrl: z.string().optional(),
 	codebaseIndexEmbedderModelId: z.string().optional(),
 	geminiEmbeddingTaskType: z.string().optional(),
+	geminiEmbeddingDimension: z.number().optional(),
 })
 
 export type CodebaseIndexConfig = z.infer<typeof codebaseIndexConfigSchema>

packages/types/src/provider-settings.ts

@@ -152,6 +152,7 @@ const geminiSchema = apiModelIdProviderModelSchema.extend({
 	geminiApiKey: z.string().optional(),
 	googleGeminiBaseUrl: z.string().optional(),
 	geminiEmbeddingTaskType: z.string().optional(),
+	geminiEmbeddingDimension: z.number().optional(),
 })
 
 const openAiNativeSchema = apiModelIdProviderModelSchema.extend({

src/services/code-index/config-manager.ts

@@ -5,6 +5,18 @@ import { CodeIndexConfig, PreviousConfigSnapshot } from "./interfaces/config"
 import { SEARCH_MIN_SCORE } from "./constants"
 import { getDefaultModelId, getModelDimension } from "../../shared/embeddingModels"
 
+// Define a type for the raw config state from globalState
+interface RawCodebaseIndexConfigState {
+	codebaseIndexEnabled?: boolean
+	codebaseIndexQdrantUrl?: string
+	codebaseIndexSearchMinScore?: number // Assuming this is also from globalState based on default
+	codebaseIndexEmbedderProvider?: "openai" | "ollama" | "gemini"
+	codebaseIndexEmbedderBaseUrl?: string
+	codebaseIndexEmbedderModelId?: string
+	geminiEmbeddingTaskType?: string
+	geminiEmbeddingDimension?: number // Ensure this is part of the raw state type
+}
+
 /**
  * Manages configuration state and validation for the code indexing feature.
  * Handles loading, validating, and providing access to configuration values.
@@ -31,15 +43,16 @@ export class CodeIndexConfigManager {
 	 */
 	private _loadAndSetConfiguration(): void {
 		// Load configuration from storage
-		const codebaseIndexConfig = this.contextProxy?.getGlobalState("codebaseIndexConfig") ?? {
+		const rawConfig = (this.contextProxy?.getGlobalState("codebaseIndexConfig") ?? {
 			codebaseIndexEnabled: false,
 			codebaseIndexQdrantUrl: "http://localhost:6333",
 			codebaseIndexSearchMinScore: 0.4,
 			codebaseIndexEmbedderProvider: "openai",
 			codebaseIndexEmbedderBaseUrl: "",
 			codebaseIndexEmbedderModelId: "",
 			geminiEmbeddingTaskType: "CODE_RETRIEVAL_QUERY",
-		}
+			geminiEmbeddingDimension: undefined,
+		}) as RawCodebaseIndexConfigState // Cast to our defined raw state type
 
 		const {
 			codebaseIndexEnabled,
@@ -48,11 +61,13 @@ export class CodeIndexConfigManager {
 			codebaseIndexEmbedderBaseUrl,
 			codebaseIndexEmbedderModelId,
 			geminiEmbeddingTaskType,
-		} = codebaseIndexConfig
+			geminiEmbeddingDimension,
+		} = rawConfig // Destructure from the typed rawConfig
 
 		const openAiKey = this.contextProxy?.getSecret("codeIndexOpenAiKey") ?? ""
 		const qdrantApiKey = this.contextProxy?.getSecret("codeIndexQdrantApiKey") ?? ""
 		const geminiApiKey = this.contextProxy?.getSecret("geminiApiKey") ?? ""
+		const rateLimitSeconds = this.contextProxy?.getGlobalState("rateLimitSeconds") ?? undefined
 
 		// Update instance variables with configuration
 		this.isEnabled = codebaseIndexEnabled || false
@@ -79,6 +94,8 @@ export class CodeIndexConfigManager {
 			geminiApiKey,
 			geminiEmbeddingTaskType: geminiEmbeddingTaskType || "CODE_RETRIEVAL_QUERY",
 			apiModelId: this.modelId,
+			geminiEmbeddingDimension,
+			rateLimitSeconds,
 		}
 	}
 
@@ -159,6 +176,7 @@ export class CodeIndexConfigManager {
 			// Gemini requires an API key and Qdrant URL
 			const geminiApiKey = this.geminiOptions?.geminiApiKey
 			const geminiEmbeddingTaskType = this.geminiOptions?.geminiEmbeddingTaskType
+
 			const qdrantUrl = this.qdrantUrl
 			const isConfigured = !!(geminiApiKey && geminiEmbeddingTaskType && qdrantUrl)
 			return isConfigured
@@ -180,6 +198,7 @@ export class CodeIndexConfigManager {
 		const prevOpenAiKey = prev?.openAiKey ?? ""
 		const prevOllamaBaseUrl = prev?.ollamaBaseUrl ?? ""
 		const prevGeminiApiKey = prev?.geminiApiKey ?? ""
+		const prevGeminiEmbeddingDimension = prev?.geminiEmbeddingDimension // Access from prev
 		const prevQdrantUrl = prev?.qdrantUrl ?? ""
 		const prevQdrantApiKey = prev?.qdrantApiKey ?? ""
 
@@ -205,7 +224,9 @@ export class CodeIndexConfigManager {
 				return true
 			}
 
-			if (this._hasVectorDimensionChanged(prevProvider, prevModelId)) {
+			// Check for dimension change, including the new geminiEmbeddingDimension
+			if (this._hasVectorDimensionChanged(prevProvider, prevModelId, prev.geminiEmbeddingDimension)) {
+				// Use prev.geminiEmbeddingDimension
 				return true
 			}
 
@@ -229,6 +250,11 @@ export class CodeIndexConfigManager {
 				if (prevGeminiApiKey !== currentGeminiApiKey) {
 					return true
 				}
+
+				const currentGeminiEmbeddingDimension = this.geminiOptions?.geminiEmbeddingDimension
+				if (currentGeminiEmbeddingDimension !== prevGeminiEmbeddingDimension) {
+					return true
+				}
 			}
 
 			// Qdrant configuration changes
@@ -246,19 +272,35 @@ export class CodeIndexConfigManager {
 	/**
 	 * Checks if model changes result in vector dimension changes that require restart.
 	 */
-	private _hasVectorDimensionChanged(prevProvider: EmbedderProvider, prevModelId?: string): boolean {
+	private _hasVectorDimensionChanged(
+		prevProvider: EmbedderProvider,
+		prevModelId?: string,
+		prevGeminiDimension?: number,
+	): boolean {
 		const currentProvider = this.embedderProvider
 		const currentModelId = this.modelId ?? getDefaultModelId(currentProvider)
 		const resolvedPrevModelId = prevModelId ?? getDefaultModelId(prevProvider)
 
 		// If model IDs are the same and provider is the same, no dimension change
 		if (prevProvider === currentProvider && resolvedPrevModelId === currentModelId) {
+			// If provider and model are same, check if gemini dimension changed
+			if (currentProvider === "gemini" && this.geminiOptions?.geminiEmbeddingDimension !== prevGeminiDimension) {
+				return true
+			}
 			return false
 		}
 
 		// Get vector dimensions for both models
-		const prevDimension = getModelDimension(prevProvider, resolvedPrevModelId)
-		const currentDimension = getModelDimension(currentProvider, currentModelId)
+		const prevDimension = getModelDimension(
+			prevProvider,
+			resolvedPrevModelId,
+			prevProvider === "gemini" ? prevGeminiDimension : undefined,
+		)
+		const currentDimension = getModelDimension(
+			currentProvider,
+			currentModelId,
+			currentProvider === "gemini" ? this.geminiOptions?.geminiEmbeddingDimension : undefined,
+		)
 
 		// If we can't determine dimensions, be safe and restart
 		if (prevDimension === undefined || currentDimension === undefined) {
@@ -284,6 +326,7 @@ export class CodeIndexConfigManager {
 			qdrantUrl: this.qdrantUrl,
 			qdrantApiKey: this.qdrantApiKey,
 			searchMinScore: this.searchMinScore,
+			geminiEmbeddingDimension: this.geminiEmbeddingDimension,
 		}
 	}
 

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

@@ -9,6 +9,7 @@ import { GEMINI_RATE_LIMIT_DELAY_MS, MAX_BATCH_RETRIES, INITIAL_RETRY_DELAY_MS }
 export class CodeIndexGeminiEmbedder extends GeminiHandler implements IEmbedder {
 	private readonly defaultModelId: string
 	private readonly defaultTaskType: string
+	private embeddingQueue: Promise<void> = Promise.resolve() // Sequential queue for embedding operations
 
 	/**
 	 * Creates a new Gemini embedder instance.
@@ -21,24 +22,47 @@ export class CodeIndexGeminiEmbedder extends GeminiHandler implements IEmbedder
 	}
 
 	/**
-	 * Creates embeddings for the given texts using the Gemini API.
+	 * Creates embeddings for the given texts using the Gemini API, ensuring sequential processing.
 	 * @param texts - An array of strings to embed.
 	 * @param model - Optional model ID to override the default.
-	 * @returns A promise that resolves to an EmbeddingResponse containing the embeddings and usage data.
+	 * @returns A promise that resolves to an EmbeddingResponse containing the embeddings.
 	 */
-	// Removed async keyword from the method signature as it no longer uses await at the top level.
-	// It constructs and returns a promise.
 	async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
-		try {
-			const modelId = model || this.defaultModelId
-			const result = await this.embedWithTokenLimit(texts, modelId, this.defaultTaskType)
-			return {
-				embeddings: result.embeddings,
+		// This function will be executed when it's this task's turn in the queue.
+		const taskExecution = async (): Promise<EmbeddingResponse> => {
+			try {
+				const modelId = model || this.defaultModelId
+				// embedWithTokenLimit handles batching, internal delays, and retries for API calls.
+				const result = await this.embedWithTokenLimit(texts, modelId, this.defaultTaskType)
+				return {
+					embeddings: result.embeddings,
+					// If EmbeddingResponse is updated to include usage, and result.usage is reliable:
+					// usage: result.usage,
+				}
+			} catch (error: any) {
+				// Errors are logged within embedWithTokenLimit or _embedBatchWithRetries.
+				// This re-throws the error to be caught by the specific caller of createEmbeddings.
+				console.error("Error during Gemini embedding task execution in queue:", error.message)
+				throw error
 			}
-		} catch (error: any) {
-			console.error("Gemini embedding task failed:", error)
-			throw error
 		}
+
+		// Chain this task onto the queue.
+		// The actual execution of taskExecution() is deferred until the previous promise in the queue resolves.
+		const taskPromise = this.embeddingQueue.then(taskExecution)
+
+		// Update the queue to wait for the current task to complete (or fail).
+		// .catch(() => {}) ensures that an error in one task doesn't break the queue for subsequent tasks.
+		// Each task's success/failure is handled by its own promise (taskPromise), which is returned to the caller.
+		this.embeddingQueue = taskPromise
+			.catch(() => {
+				// This task failed, but the queue should proceed for the next one.
+				// The error from taskPromise will be handled by its specific awaiter below.
+			})
+			.then(() => undefined) // Ensure the queue promise resolves to void for the next .then() in the chain.
+
+		// Return the promise for this specific task. The caller will await this.
+		return taskPromise
 	}
 
 	/**
@@ -112,12 +136,15 @@ export class CodeIndexGeminiEmbedder extends GeminiHandler implements IEmbedder
 
 			// Process the current batch if not empty
 			if (currentBatch.length > 0) {
-				const delayMs =
-					this.options.rateLimitSeconds !== undefined
-						? this.options.rateLimitSeconds * 1000
-						: GEMINI_RATE_LIMIT_DELAY_MS
-				console.log(`Adding proactive delay of ${delayMs}ms before Gemini batch`)
-				await new Promise((resolve) => setTimeout(resolve, delayMs))
+				if (!isFirstBatch) {
+					const delayMs =
+						this.options.rateLimitSeconds !== undefined
+							? this.options.rateLimitSeconds * 1000
+							: GEMINI_RATE_LIMIT_DELAY_MS
+					console.log(`Adding proactive delay of ${delayMs}ms before Gemini batch`)
+					await new Promise((resolve) => setTimeout(resolve, delayMs))
+					isFirstBatch = false
+				}
 
 				try {
 					const batchResult = await this._embedBatchWithRetries(currentBatch, model, taskType)

src/services/code-index/interfaces/config.ts

@@ -15,6 +15,7 @@ export interface CodeIndexConfig {
 	qdrantUrl?: string
 	qdrantApiKey?: string
 	searchMinScore?: number
+	geminiEmbeddingDimension?: number
 }
 
 /**
@@ -29,6 +30,7 @@ export type PreviousConfigSnapshot = {
 	ollamaBaseUrl?: string
 	geminiApiKey?: string
 	geminiEmbeddingTaskType?: string
+	geminiEmbeddingDimension?: number // Add here
 	qdrantUrl?: string
 	qdrantApiKey?: string
 }

src/services/code-index/service-factory.ts

@@ -44,7 +44,7 @@ export class CodeIndexServiceFactory {
 				...config.ollamaOptions,
 				ollamaModelId: config.modelId,
 			})
-        } else if (provider === "gemini") {
+		} else if (provider === "gemini") {
 			if (!config.geminiOptions?.geminiApiKey) {
 				throw new Error("Gemini configuration missing for embedder creation")
 			}
@@ -64,8 +64,12 @@ export class CodeIndexServiceFactory {
 		const defaultModel = getDefaultModelId(provider)
 		// Use the embedding model ID from config, not the chat model IDs
 		const modelId = config.modelId ?? defaultModel
+		let requestedDimension: number | undefined
+		if (provider === "gemini") {
+			requestedDimension = config.geminiEmbeddingDimension
+		}
 
-		const vectorSize = getModelDimension(provider, modelId)
+		const vectorSize = getModelDimension(provider, modelId, requestedDimension)
 
 		if (vectorSize === undefined) {
 			throw new Error(

src/shared/embeddingModels.ts

@@ -48,9 +48,14 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
  * Retrieves the embedding dimension for a given provider and model ID.
  * @param provider The embedder provider (e.g., "openai").
  * @param modelId The specific model ID (e.g., "text-embedding-3-small").
+ * @param requestedDimension Optional dimension requested by the user.
  * @returns The dimension size or undefined if the model is not found.
  */
-export function getModelDimension(provider: EmbedderProvider, modelId: string): number | undefined {
+export function getModelDimension(
+	provider: EmbedderProvider,
+	modelId: string,
+	requestedDimension?: number,
+): number | undefined {
 	const providerProfiles = EMBEDDING_MODEL_PROFILES[provider]
 	if (!providerProfiles) {
 		console.warn(`Provider not found in profiles: ${provider}`)
@@ -64,6 +69,14 @@ export function getModelDimension(provider: EmbedderProvider, modelId: string):
 		return undefined // Or potentially return a default/fallback dimension?
 	}
 
+	if (
+		requestedDimension &&
+		modelProfile.supportDimensions &&
+		modelProfile.supportDimensions.includes(requestedDimension)
+	) {
+		return requestedDimension
+	}
+
 	return modelProfile.dimension
 }