diff --git a/packages/types/src/codebase-index.ts b/packages/types/src/codebase-index.ts index 0ad19d8676..ffe9774609 100644 --- a/packages/types/src/codebase-index.ts +++ b/packages/types/src/codebase-index.ts @@ -21,7 +21,7 @@ export const CODEBASE_INDEX_DEFAULTS = { 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(), codebaseIndexEmbedderModelDimension: z.number().optional(), @@ -47,6 +47,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 @@ -62,6 +63,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 diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts index 245621a1bd..adee27156f 100644 --- 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 } private geminiOptions?: { apiKey: string } + private lmStudioOptions?: ApiHandlerOptions private qdrantUrl?: string = "http://localhost:6333" private qdrantApiKey?: string private searchMinScore?: number @@ -92,13 +93,15 @@ export class CodeIndexConfigManager { 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" } @@ -118,6 +121,10 @@ export class CodeIndexConfigManager { : undefined this.geminiOptions = geminiApiKey ? { apiKey: geminiApiKey } : undefined + + this.lmStudioOptions = { + lmStudioBaseUrl: codebaseIndexEmbedderBaseUrl, + } } /** @@ -134,6 +141,7 @@ export class CodeIndexConfigManager { ollamaOptions?: ApiHandlerOptions openAiCompatibleOptions?: { baseUrl: string; apiKey: string } geminiOptions?: { apiKey: string } + lmStudioOptions?: ApiHandlerOptions qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number @@ -152,6 +160,7 @@ export class CodeIndexConfigManager { openAiCompatibleBaseUrl: this.openAiCompatibleOptions?.baseUrl ?? "", openAiCompatibleApiKey: this.openAiCompatibleOptions?.apiKey ?? "", geminiApiKey: this.geminiOptions?.apiKey ?? "", + lmStudioBaseUrl: this.lmStudioOptions?.lmStudioBaseUrl ?? "", qdrantUrl: this.qdrantUrl ?? "", qdrantApiKey: this.qdrantApiKey ?? "", } @@ -175,6 +184,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, @@ -207,6 +217,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 } @@ -240,6 +256,7 @@ export class CodeIndexConfigManager { const prevOpenAiCompatibleApiKey = prev?.openAiCompatibleApiKey ?? "" const prevModelDimension = prev?.modelDimension const prevGeminiApiKey = prev?.geminiApiKey ?? "" + const prevLmStudioBaseUrl = prev?.lmStudioBaseUrl ?? "" const prevQdrantUrl = prev?.qdrantUrl ?? "" const prevQdrantApiKey = prev?.qdrantApiKey ?? "" @@ -269,6 +286,7 @@ export class CodeIndexConfigManager { const currentOpenAiCompatibleApiKey = this.openAiCompatibleOptions?.apiKey ?? "" const currentModelDimension = this.modelDimension const currentGeminiApiKey = this.geminiOptions?.apiKey ?? "" + const currentLmStudioBaseUrl = this.lmStudioOptions?.lmStudioBaseUrl ?? "" const currentQdrantUrl = this.qdrantUrl ?? "" const currentQdrantApiKey = this.qdrantApiKey ?? "" @@ -292,6 +310,14 @@ export class CodeIndexConfigManager { return true } + if (prevGeminiApiKey !== currentGeminiApiKey) { + return true + } + + if (prevLmStudioBaseUrl !== currentLmStudioBaseUrl) { + return true + } + if (prevQdrantUrl !== currentQdrantUrl || prevQdrantApiKey !== currentQdrantApiKey) { return true } @@ -343,6 +369,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 new file mode 100644 index 0000000000..25d6ef07fa --- /dev/null +++ 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 { + 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 index 190a23e2a3..ea44e404de 100644 --- 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 } geminiOptions?: { apiKey: string } + lmStudioOptions?: ApiHandlerOptions qdrantUrl?: string qdrantApiKey?: string searchMinScore?: number @@ -33,6 +34,7 @@ export type PreviousConfigSnapshot = { openAiCompatibleBaseUrl?: string openAiCompatibleApiKey?: string 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 index 0a74446d5e..58532ff478 100644 --- a/src/services/code-index/interfaces/embedder.ts +++ b/src/services/code-index/interfaces/embedder.ts @@ -28,7 +28,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 index 70e3fd9765..4e96aef30b 100644 --- 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 index 818dafb497..1fe6a9b4a4 100644 --- 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" @@ -62,6 +63,14 @@ export class CodeIndexServiceFactory { throw new Error(t("embeddings:serviceFactory.geminiConfigMissing")) } 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( diff --git a/src/shared/embeddingModels.ts b/src/shared/embeddingModels.ts index 4c6bc24319..09570765a8 100644 --- 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 index da40058b00..e91b9ddcf2 100644 --- a/webview-ui/src/i18n/locales/en/settings.json +++ b/webview-ui/src/i18n/locales/en/settings.json @@ -50,9 +50,10 @@ "openaiProvider": "OpenAI", "ollamaProvider": "Ollama", "geminiProvider": "Gemini", + "openaiCompatibleProvider": "OpenAI Compatible", + "lmstudioProvider": "LM Studio", "geminiApiKeyLabel": "API Key:", "geminiApiKeyPlaceholder": "Enter your Gemini API key", - "openaiCompatibleProvider": "OpenAI Compatible", "openAiKeyLabel": "OpenAI API Key", "openAiKeyPlaceholder": "Enter your OpenAI API key", "openAiCompatibleBaseUrlLabel": "Base URL", @@ -62,6 +63,7 @@ "modelDimensionLabel": "Model Dimension", "openAiCompatibleModelDimensionPlaceholder": "e.g., 1536", "openAiCompatibleModelDimensionDescription": "The embedding dimension (output size) for your model. Check your provider's documentation for this value. Common values: 384, 768, 1536, 3072.", + "lmstudioUrlLabel": "LM Studio URL:", "modelLabel": "Model", "modelPlaceholder": "Enter model name", "selectModel": "Select a model",