Add LM Studio support to code indexing service and settings

kiwina · MuriloFP · commit 37225088daec · 2025-07-04T16:47:22.000-03:00
diff --git a/packages/types/src/codebase-index.ts b/packages/types/src/codebase-index.ts
@@ -7,7 +7,7 @@ import { z } from "zod"
 export const codebaseIndexConfigSchema = z.object({
 	codebaseIndexEnabled: z.boolean().optional(),
 	codebaseIndexQdrantUrl: z.string().optional(),
-	codebaseIndexEmbedderProvider: z.enum(["openai", "ollama", "openai-compatible", "gemini"]).optional(),
+	codebaseIndexEmbedderProvider: z.enum(["openai", "ollama", "openai-compatible", "gemini", "lmstudio"]).optional(),
 	codebaseIndexEmbedderBaseUrl: z.string().optional(),
 	codebaseIndexEmbedderModelId: z.string().optional(),
 	codebaseIndexSearchMinScore: z.number().min(0).max(1).optional(),
@@ -24,6 +24,7 @@ export const codebaseIndexModelsSchema = z.object({
 	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(),
+	lmstudio: z.record(z.string(), z.object({ dimension: z.number() })).optional(),
 })
 
 export type CodebaseIndexModels = z.infer<typeof codebaseIndexModelsSchema>
@@ -39,6 +40,7 @@ export const codebaseIndexProviderSchema = z.object({
 	codebaseIndexOpenAiCompatibleApiKey: z.string().optional(),
 	codebaseIndexOpenAiCompatibleModelDimension: z.number().optional(),
 	codebaseIndexGeminiApiKey: z.string().optional(),
+	codebaseIndexLmStudioBaseUrl: z.string().optional(),
 })
 
 export type CodebaseIndexProvider = z.infer<typeof codebaseIndexProviderSchema>
diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts
@@ -17,6 +17,7 @@ export class CodeIndexConfigManager {
 	private ollamaOptions?: ApiHandlerOptions
 	private openAiCompatibleOptions?: { baseUrl: string; apiKey: string; modelDimension?: number }
 	private geminiOptions?: { apiKey: string }
+	private lmStudioOptions?: ApiHandlerOptions
 	private qdrantUrl?: string = "http://localhost:6333"
 	private qdrantApiKey?: string
 	private searchMinScore?: number
@@ -73,13 +74,15 @@ export class CodeIndexConfigManager {
 		this.searchMinScore = codebaseIndexSearchMinScore
 		this.openAiOptions = { openAiNativeApiKey: openAiKey }
 
-		// Set embedder provider with support for openai-compatible
+		// Set embedder provider with support for all providers
 		if (codebaseIndexEmbedderProvider === "ollama") {
 			this.embedderProvider = "ollama"
 		} else if (codebaseIndexEmbedderProvider === "openai-compatible") {
 			this.embedderProvider = "openai-compatible"
 		} else if (codebaseIndexEmbedderProvider === "gemini") {
 			this.embedderProvider = "gemini"
+		} else if (codebaseIndexEmbedderProvider === "lmstudio") {
+			this.embedderProvider = "lmstudio"
 		} else {
 			this.embedderProvider = "openai"
 		}
@@ -100,6 +103,10 @@ export class CodeIndexConfigManager {
 				: undefined
 
 		this.geminiOptions = geminiApiKey ? { apiKey: geminiApiKey } : undefined
+
+		this.lmStudioOptions = {
+			lmStudioBaseUrl: codebaseIndexEmbedderBaseUrl,
+		}
 	}
 
 	/**
@@ -116,6 +123,7 @@ export class CodeIndexConfigManager {
 			ollamaOptions?: ApiHandlerOptions
 			openAiCompatibleOptions?: { baseUrl: string; apiKey: string }
 			geminiOptions?: { apiKey: string }
+			lmStudioOptions?: ApiHandlerOptions
 			qdrantUrl?: string
 			qdrantApiKey?: string
 			searchMinScore?: number
@@ -134,6 +142,7 @@ export class CodeIndexConfigManager {
 			openAiCompatibleApiKey: this.openAiCompatibleOptions?.apiKey ?? "",
 			openAiCompatibleModelDimension: this.openAiCompatibleOptions?.modelDimension,
 			geminiApiKey: this.geminiOptions?.apiKey ?? "",
+			lmStudioBaseUrl: this.lmStudioOptions?.lmStudioBaseUrl ?? "",
 			qdrantUrl: this.qdrantUrl ?? "",
 			qdrantApiKey: this.qdrantApiKey ?? "",
 		}
@@ -157,6 +166,7 @@ export class CodeIndexConfigManager {
 				ollamaOptions: this.ollamaOptions,
 				openAiCompatibleOptions: this.openAiCompatibleOptions,
 				geminiOptions: this.geminiOptions,
+				lmStudioOptions: this.lmStudioOptions,
 				qdrantUrl: this.qdrantUrl,
 				qdrantApiKey: this.qdrantApiKey,
 				searchMinScore: this.currentSearchMinScore,
@@ -189,6 +199,12 @@ export class CodeIndexConfigManager {
 			const qdrantUrl = this.qdrantUrl
 			const isConfigured = !!(apiKey && qdrantUrl)
 			return isConfigured
+		} else if (this.embedderProvider === "lmstudio") {
+			// Lm Studio model ID has a default, so only base URL is strictly required for config
+			const lmStudioBaseUrl = this.lmStudioOptions?.lmStudioBaseUrl
+			const qdrantUrl = this.qdrantUrl
+			const isConfigured = !!(lmStudioBaseUrl && qdrantUrl)
+			return isConfigured
 		}
 		return false // Should not happen if embedderProvider is always set correctly
 	}
@@ -222,6 +238,7 @@ export class CodeIndexConfigManager {
 		const prevOpenAiCompatibleApiKey = prev?.openAiCompatibleApiKey ?? ""
 		const prevOpenAiCompatibleModelDimension = prev?.openAiCompatibleModelDimension
 		const prevGeminiApiKey = prev?.geminiApiKey ?? ""
+		const prevLmStudioBaseUrl = prev?.lmStudioBaseUrl ?? ""
 		const prevQdrantUrl = prev?.qdrantUrl ?? ""
 		const prevQdrantApiKey = prev?.qdrantApiKey ?? ""
 
@@ -254,6 +271,7 @@ export class CodeIndexConfigManager {
 			const currentOpenAiCompatibleApiKey = this.openAiCompatibleOptions?.apiKey ?? ""
 			const currentOpenAiCompatibleModelDimension = this.openAiCompatibleOptions?.modelDimension
 			const currentGeminiApiKey = this.geminiOptions?.apiKey ?? ""
+			const currentLmStudioBaseUrl = this.lmStudioOptions?.lmStudioBaseUrl ?? ""
 			const currentQdrantUrl = this.qdrantUrl ?? ""
 			const currentQdrantApiKey = this.qdrantApiKey ?? ""
 
@@ -279,6 +297,14 @@ export class CodeIndexConfigManager {
 				}
 			}
 
+			if (prevGeminiApiKey !== currentGeminiApiKey) {
+				return true
+			}
+
+			if (prevLmStudioBaseUrl !== currentLmStudioBaseUrl) {
+				return true
+			}
+
 			if (prevQdrantUrl !== currentQdrantUrl || prevQdrantApiKey !== currentQdrantApiKey) {
 				return true
 			}
@@ -331,6 +357,7 @@ export class CodeIndexConfigManager {
 			ollamaOptions: this.ollamaOptions,
 			openAiCompatibleOptions: this.openAiCompatibleOptions,
 			geminiOptions: this.geminiOptions,
+			lmStudioOptions: this.lmStudioOptions,
 			qdrantUrl: this.qdrantUrl,
 			qdrantApiKey: this.qdrantApiKey,
 			searchMinScore: this.currentSearchMinScore,
diff --git a/src/services/code-index/embedders/lmstudio.ts b/src/services/code-index/embedders/lmstudio.ts
@@ -0,0 +1,140 @@
+import { OpenAI } from "openai"
+import { ApiHandlerOptions } from "../../../shared/api"
+import { IEmbedder, EmbeddingResponse, EmbedderInfo } from "../interfaces"
+import {
+	MAX_BATCH_TOKENS,
+	MAX_ITEM_TOKENS,
+	MAX_BATCH_RETRIES as MAX_RETRIES,
+	INITIAL_RETRY_DELAY_MS as INITIAL_DELAY_MS,
+} from "../constants"
+
+/**
+ * LM Studio implementation of the embedder interface with batching and rate limiting.
+ * Uses OpenAI-compatible API endpoints with a custom base URL.
+ */
+export class CodeIndexLmStudioEmbedder implements IEmbedder {
+	protected options: ApiHandlerOptions
+	private embeddingsClient: OpenAI
+	private readonly defaultModelId: string
+
+	/**
+	 * Creates a new LM Studio embedder
+	 * @param options API handler options including lmStudioBaseUrl
+	 */
+	constructor(options: ApiHandlerOptions & { embeddingModelId?: string }) {
+		this.options = options
+		this.embeddingsClient = new OpenAI({
+			baseURL: (this.options.lmStudioBaseUrl || "http://localhost:1234") + "/v1",
+			apiKey: "noop", // LM Studio doesn't require a real API key
+		})
+		this.defaultModelId = options.embeddingModelId || "text-embedding-nomic-embed-text-v1.5@f16"
+	}
+
+	/**
+	 * Creates embeddings for the given texts with batching and rate limiting
+	 * @param texts Array of text strings to embed
+	 * @param model Optional model identifier
+	 * @returns Promise resolving to embedding response
+	 */
+	async createEmbeddings(texts: string[], model?: string): Promise<EmbeddingResponse> {
+		const modelToUse = model || this.defaultModelId
+		const allEmbeddings: number[][] = []
+		const usage = { promptTokens: 0, totalTokens: 0 }
+		const remainingTexts = [...texts]
+
+		while (remainingTexts.length > 0) {
+			const currentBatch: string[] = []
+			let currentBatchTokens = 0
+			const processedIndices: number[] = []
+
+			for (let i = 0; i < remainingTexts.length; i++) {
+				const text = remainingTexts[i]
+				const itemTokens = Math.ceil(text.length / 4)
+
+				if (itemTokens > MAX_ITEM_TOKENS) {
+					console.warn(
+						`Text at index ${i} exceeds maximum token limit (${itemTokens} > ${MAX_ITEM_TOKENS}). Skipping.`,
+					)
+					processedIndices.push(i)
+					continue
+				}
+
+				if (currentBatchTokens + itemTokens <= MAX_BATCH_TOKENS) {
+					currentBatch.push(text)
+					currentBatchTokens += itemTokens
+					processedIndices.push(i)
+				} else {
+					break
+				}
+			}
+
+			// Remove processed items from remainingTexts (in reverse order to maintain correct indices)
+			for (let i = processedIndices.length - 1; i >= 0; i--) {
+				remainingTexts.splice(processedIndices[i], 1)
+			}
+
+			if (currentBatch.length > 0) {
+				try {
+					const batchResult = await this._embedBatchWithRetries(currentBatch, modelToUse)
+
+					allEmbeddings.push(...batchResult.embeddings)
+					usage.promptTokens += batchResult.usage.promptTokens
+					usage.totalTokens += batchResult.usage.totalTokens
+				} catch (error) {
+					console.error("Failed to process batch:", error)
+					throw new Error("Failed to create embeddings: batch processing error")
+				}
+			}
+		}
+
+		return { embeddings: allEmbeddings, usage }
+	}
+
+	/**
+	 * Helper method to handle batch embedding with retries and exponential backoff
+	 * @param batchTexts Array of texts to embed in this batch
+	 * @param model Model identifier to use
+	 * @returns Promise resolving to embeddings and usage statistics
+	 */
+	private async _embedBatchWithRetries(
+		batchTexts: string[],
+		model: string,
+	): Promise<{ embeddings: number[][]; usage: { promptTokens: number; totalTokens: number } }> {
+		for (let attempts = 0; attempts < MAX_RETRIES; attempts++) {
+			try {
+				const response = await this.embeddingsClient.embeddings.create({
+					input: batchTexts,
+					model: model,
+					encoding_format: "float",
+				})
+
+				return {
+					embeddings: response.data.map((item) => item.embedding),
+					usage: {
+						promptTokens: response.usage?.prompt_tokens || 0,
+						totalTokens: response.usage?.total_tokens || 0,
+					},
+				}
+			} catch (error: any) {
+				const isRateLimitError = error?.status === 429
+				const hasMoreAttempts = attempts < MAX_RETRIES - 1
+
+				if (isRateLimitError && hasMoreAttempts) {
+					const delayMs = INITIAL_DELAY_MS * Math.pow(2, attempts)
+					await new Promise((resolve) => setTimeout(resolve, delayMs))
+					continue
+				}
+
+				throw error
+			}
+		}
+
+		throw new Error(`Failed to create embeddings after ${MAX_RETRIES} attempts`)
+	}
+
+	get embedderInfo(): EmbedderInfo {
+		return {
+			name: "lmstudio",
+		}
+	}
+}
diff --git a/src/services/code-index/interfaces/config.ts b/src/services/code-index/interfaces/config.ts
@@ -13,6 +13,7 @@ export interface CodeIndexConfig {
 	ollamaOptions?: ApiHandlerOptions
 	openAiCompatibleOptions?: { baseUrl: string; apiKey: string; modelDimension?: number }
 	geminiOptions?: { apiKey: string }
+	lmStudioOptions?: ApiHandlerOptions
 	qdrantUrl?: string
 	qdrantApiKey?: string
 	searchMinScore?: number
@@ -32,6 +33,7 @@ export type PreviousConfigSnapshot = {
 	openAiCompatibleApiKey?: string
 	openAiCompatibleModelDimension?: number
 	geminiApiKey?: string
+	lmStudioBaseUrl?: string
 	qdrantUrl?: string
 	qdrantApiKey?: string
 }
diff --git a/src/services/code-index/interfaces/embedder.ts b/src/services/code-index/interfaces/embedder.ts
@@ -21,7 +21,7 @@ export interface EmbeddingResponse {
 	}
 }
 
-export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini"
+export type AvailableEmbedders = "openai" | "ollama" | "openai-compatible" | "gemini" | "lmstudio"
 
 export interface EmbedderInfo {
 	name: AvailableEmbedders
diff --git a/src/services/code-index/interfaces/manager.ts b/src/services/code-index/interfaces/manager.ts
@@ -70,7 +70,15 @@ export interface ICodeIndexManager {
 }
 
 export type IndexingState = "Standby" | "Indexing" | "Indexed" | "Error"
-export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini"
+
+/**
+ * Supported embedder providers for code indexing.
+ * To add a new provider:
+ * 1. Add the provider name to this union type
+ * 2. Update the switch statements in CodeIndexConfigManager
+ * 3. Add provider-specific configuration options
+ */
+export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "lmstudio"
 
 export interface IndexProgressUpdate {
 	systemStatus: IndexingState
diff --git a/src/services/code-index/service-factory.ts b/src/services/code-index/service-factory.ts
@@ -3,6 +3,7 @@ import { OpenAiEmbedder } from "./embedders/openai"
 import { CodeIndexOllamaEmbedder } from "./embedders/ollama"
 import { OpenAICompatibleEmbedder } from "./embedders/openai-compatible"
 import { GeminiEmbedder } from "./embedders/gemini"
+import { CodeIndexLmStudioEmbedder } from "./embedders/lmstudio"
 import { EmbedderProvider, getDefaultModelId, getModelDimension } from "../../shared/embeddingModels"
 import { QdrantVectorStore } from "./vector-store/qdrant-client"
 import { codeParser, DirectoryScanner, FileWatcher } from "./processors"
@@ -61,6 +62,14 @@ export class CodeIndexServiceFactory {
 				throw new Error("Gemini configuration missing for embedder creation")
 			}
 			return new GeminiEmbedder(config.geminiOptions.apiKey)
+		} else if (provider === "lmstudio") {
+			if (!config.lmStudioOptions?.lmStudioBaseUrl) {
+				throw new Error("LM Studio configuration missing for embedder creation")
+			}
+			return new CodeIndexLmStudioEmbedder({
+				...config.lmStudioOptions,
+				embeddingModelId: config.modelId,
+			})
 		}
 
 		throw new Error(`Invalid embedder type configured: ${config.embedderProvider}`)
diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts
@@ -2,7 +2,7 @@
  * Defines profiles for different embedding models, including their dimensions.
  */
 
-export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" // Add other providers as needed
+export type EmbedderProvider = "openai" | "ollama" | "openai-compatible" | "gemini" | "lmstudio" // Add other providers as needed
 
 export interface EmbeddingModelProfile {
 	dimension: number
@@ -49,6 +49,11 @@ export const EMBEDDING_MODEL_PROFILES: EmbeddingModelProfiles = {
 	gemini: {
 		"text-embedding-004": { dimension: 768 },
 	},
+	lmstudio: {
+		"text-embedding-nomic-embed-text-v1.5@f16": { dimension: 768 },
+		"text-embedding-nomic-embed-text-v1.5@f32": { dimension: 768 },
+		"text-embedding-mxbai-embed-large-v1": { dimension: 1024 },
+	},
 }
 
 /**
@@ -136,6 +141,19 @@ export function getDefaultModelId(provider: EmbedderProvider): string {
 		case "gemini":
 			return "text-embedding-004"
 
+		case "lmstudio": {
+			// Choose a sensible default for LM Studio, e.g., the first one listed or a specific one
+			const lmStudioModels = EMBEDDING_MODEL_PROFILES.lmstudio
+			const defaultLmStudioModel = lmStudioModels && Object.keys(lmStudioModels)[0]
+			if (defaultLmStudioModel) {
+				return defaultLmStudioModel
+			}
+			// Fallback if no LM Studio models are defined (shouldn't happen with the constant)
+			console.warn("No default LM Studio model found in profiles.")
+			// Return a placeholder or throw an error, depending on desired behavior
+			return "unknown-default" // Placeholder specific model ID
+		}
+
 		default:
 			// Fallback for unknown providers
 			console.warn(`Unknown provider for default model ID: ${provider}. Falling back to OpenAI default.`)
diff --git a/webview-ui/src/i18n/locales/en/settings.json b/webview-ui/src/i18n/locales/en/settings.json

Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ export interface EmbeddingResponse {`
`21`	`21`	`}`
`22`	`22`	`}`
`23`	`23`
`24`		`-export type AvailableEmbedders = "openai" \| "ollama" \| "openai-compatible" \| "gemini"`
	`24`	`+export type AvailableEmbedders = "openai" \| "ollama" \| "openai-compatible" \| "gemini" \| "lmstudio"`
`25`	`25`
`26`	`26`	`export interface EmbedderInfo {`
`27`	`27`	`name: AvailableEmbedders`