Add LM Studio support to code indexing service and settings

This commit is contained in:
kiwina 2025-05-26 23:21:56 +08:00 committed by Daniel Riccio
parent 50598b22b5
commit 0bfe5c86ee
No known key found for this signature in database
GPG key ID: FFD5FD825F8E8209
9 changed files with 214 additions and 6 deletions

View file

@ -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<typeof codebaseIndexModelsSchema>
@ -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<typeof codebaseIndexProviderSchema>

View file

@ -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,

View file

@ -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",
}
}
}

View file

@ -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
}

View file

@ -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

View file

@ -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

View file

@ -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(

View file

@ -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.`)

View file

@ -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",