mirror of
https://github.com/RooVetGit/Roo-Code.git
synced 2026-09-07 08:26:51 +00:00
Add LM Studio support to code indexing service and settings
This commit is contained in:
parent
50598b22b5
commit
0bfe5c86ee
9 changed files with 214 additions and 6 deletions
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
140
src/services/code-index/embedders/lmstudio.ts
Normal file
140
src/services/code-index/embedders/lmstudio.ts
Normal 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",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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.`)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue