diff --git a/packages/types/src/codebase-index.ts b/packages/types/src/codebase-index.ts index 89d5b168d7..eebd8f73ab 100644 --- a/packages/types/src/codebase-index.ts +++ b/packages/types/src/codebase-index.ts @@ -34,6 +34,8 @@ export const codebaseIndexConfigSchema = z.object({ // OpenAI Compatible specific fields codebaseIndexOpenAiCompatibleBaseUrl: z.string().optional(), codebaseIndexOpenAiCompatibleModelDimension: z.number().optional(), + // Memory storage settings + memoryStorageEnabled: z.boolean().optional(), }) export type CodebaseIndexConfig = z.infer diff --git a/packages/types/src/global-settings.ts b/packages/types/src/global-settings.ts index 6de4d7413f..6d2049189f 100644 --- a/packages/types/src/global-settings.ts +++ b/packages/types/src/global-settings.ts @@ -62,6 +62,10 @@ export const globalSettingsSchema = z.object({ alwaysAllowFollowupQuestions: z.boolean().optional(), followupAutoApproveTimeoutMs: z.number().optional(), alwaysAllowUpdateTodoList: z.boolean().optional(), + + // Memory storage settings + memoryStorageEnabled: z.boolean().optional(), + memoryStorageAutoApprove: z.boolean().optional(), allowedCommands: z.array(z.string()).optional(), deniedCommands: z.array(z.string()).optional(), commandExecutionTimeout: z.number().optional(), @@ -240,6 +244,8 @@ export const EVALS_SETTINGS: RooCodeSettings = { alwaysAllowFollowupQuestions: true, alwaysAllowUpdateTodoList: true, followupAutoApproveTimeoutMs: 0, + memoryStorageEnabled: false, + memoryStorageAutoApprove: false, allowedCommands: ["*"], commandExecutionTimeout: 20, commandTimeoutAllowlist: [], diff --git a/packages/types/src/tool.ts b/packages/types/src/tool.ts index 7a3fd21199..241d5dd7e9 100644 --- a/packages/types/src/tool.ts +++ b/packages/types/src/tool.ts @@ -28,6 +28,8 @@ export const toolNames = [ "use_mcp_tool", "access_mcp_resource", "ask_followup_question", + "ask_memory_aware_followup_question", + "search_memories", "attempt_completion", "switch_mode", "new_task", diff --git a/src/core/assistant-message/presentAssistantMessage.ts b/src/core/assistant-message/presentAssistantMessage.ts index ee3fa148b4..f04c3f177d 100644 --- a/src/core/assistant-message/presentAssistantMessage.ts +++ b/src/core/assistant-message/presentAssistantMessage.ts @@ -21,6 +21,8 @@ import { executeCommandTool } from "../tools/executeCommandTool" import { useMcpToolTool } from "../tools/useMcpToolTool" import { accessMcpResourceTool } from "../tools/accessMcpResourceTool" import { askFollowupQuestionTool } from "../tools/askFollowupQuestionTool" +import { askMemoryAwareFollowupQuestionTool } from "../tools/askMemoryAwareFollowupQuestionTool" +import { searchMemoriesTool } from "../tools/searchMemoriesTool" import { switchModeTool } from "../tools/switchModeTool" import { attemptCompletionTool } from "../tools/attemptCompletionTool" import { newTaskTool } from "../tools/newTaskTool" @@ -151,6 +153,10 @@ export async function presentAssistantMessage(cline: Task) { break } case "tool_use": + // Get customModes early for use in toolDescription + const stateForDescription = await cline.providerRef.deref()?.getState() + const customModesForDescription = stateForDescription?.customModes ?? [] + const toolDescription = (): string => { switch (block.name) { case "execute_command": @@ -200,6 +206,10 @@ export async function presentAssistantMessage(cline: Task) { return `[${block.name} for '${block.params.server_name}']` case "ask_followup_question": return `[${block.name} for '${block.params.question}']` + case "ask_memory_aware_followup_question": + return `[${block.name} for '${block.params.question}']` + case "search_memories": + return `[${block.name} for '${block.params.query}']` case "attempt_completion": return `[${block.name}]` case "switch_mode": @@ -211,7 +221,7 @@ export async function presentAssistantMessage(cline: Task) { case "new_task": { const mode = block.params.mode ?? defaultModeSlug const message = block.params.message ?? "(no message)" - const modeName = getModeBySlug(mode, customModes)?.name ?? mode + const modeName = getModeBySlug(mode, customModesForDescription)?.name ?? mode return `[${block.name} in ${modeName} mode: '${message}']` } } @@ -504,6 +514,19 @@ export async function presentAssistantMessage(cline: Task) { removeClosingTag, ) break + case "ask_memory_aware_followup_question": + await askMemoryAwareFollowupQuestionTool( + cline, + block, + askApproval, + handleError, + pushToolResult, + removeClosingTag, + ) + break + case "search_memories": + await searchMemoriesTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) + break case "switch_mode": await switchModeTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag) break diff --git a/src/core/tools/askFollowupQuestionTool.ts b/src/core/tools/askFollowupQuestionTool.ts index e736936887..d389877b18 100644 --- a/src/core/tools/askFollowupQuestionTool.ts +++ b/src/core/tools/askFollowupQuestionTool.ts @@ -2,6 +2,10 @@ import { Task } from "../task/Task" import { ToolUse, AskApproval, HandleError, PushToolResult, RemoveClosingTag } from "../../shared/tools" import { formatResponse } from "../prompts/responses" import { parseXml } from "../../utils/xml" +import { CodeIndexConfigManager } from "../../services/code-index/config-manager" +import { CodeIndexServiceFactory } from "../../services/code-index/service-factory" +import { MemoryStorageManager } from "../../services/memory-storage/MemoryStorageManager" +import { CacheManager } from "../../services/code-index/cache-manager" export async function askFollowupQuestionTool( cline: Task, @@ -80,6 +84,49 @@ export async function askFollowupQuestionTool( await cline.say("user_feedback", text ?? "", images) pushToolResult(formatResponse.toolResult(`\n${text}\n`, images)) + // Store memory if enabled + try { + const provider = cline.providerRef.deref() + if (provider && text) { + // Get the code index manager from the provider + const codeIndexManager = provider.codeIndexManager + if (codeIndexManager) { + // Create config manager and service factory + const configManager = new CodeIndexConfigManager(provider.contextProxy) + const cacheManager = new CacheManager(provider.context, cline.workspacePath) + const serviceFactory = new CodeIndexServiceFactory( + configManager, + cline.workspacePath, + cacheManager, + ) + + // Get or create the memory storage manager + const memoryManager = MemoryStorageManager.getInstance( + configManager, + serviceFactory, + cline.workspacePath, + ) + + // Store the memory if enabled + if (memoryManager.isEnabled()) { + const memoryService = await memoryManager.getMemoryStorageService() + if (memoryService) { + await memoryService.storeMemory( + question, + text, + follow_up_json.suggest, + cline.taskId, + await cline.getTaskMode(), + ) + } + } + } + } + } catch (error) { + // Log error but don't fail the tool + console.error("Failed to store memory:", error) + } + return } } catch (error) { diff --git a/src/core/tools/askMemoryAwareFollowupQuestionTool.ts b/src/core/tools/askMemoryAwareFollowupQuestionTool.ts new file mode 100644 index 0000000000..d8caf9337b --- /dev/null +++ b/src/core/tools/askMemoryAwareFollowupQuestionTool.ts @@ -0,0 +1,183 @@ +import { Task } from "../task/Task" +import { ToolUse, AskApproval, HandleError, PushToolResult, RemoveClosingTag } from "../../shared/tools" +import { formatResponse } from "../prompts/responses" +import { parseXml } from "../../utils/xml" +import { CodeIndexConfigManager } from "../../services/code-index/config-manager" +import { CodeIndexServiceFactory } from "../../services/code-index/service-factory" +import { MemoryStorageManager } from "../../services/memory-storage/MemoryStorageManager" +import { CacheManager } from "../../services/code-index/cache-manager" + +export async function askMemoryAwareFollowupQuestionTool( + cline: Task, + block: ToolUse, + askApproval: AskApproval, + handleError: HandleError, + pushToolResult: PushToolResult, + removeClosingTag: RemoveClosingTag, +) { + const question: string | undefined = block.params.question + const follow_up: string | undefined = block.params.follow_up + + try { + if (block.partial) { + await cline.ask("followup", removeClosingTag("question", question), block.partial).catch(() => {}) + return + } else { + if (!question) { + cline.consecutiveMistakeCount++ + cline.recordToolError("ask_memory_aware_followup_question") + pushToolResult( + await cline.sayAndCreateMissingParamError("ask_memory_aware_followup_question", "question"), + ) + return + } + + type Suggest = { answer: string; mode?: string } + + let follow_up_json = { + question, + suggest: [] as Suggest[], + } + + if (follow_up) { + // Define the actual structure returned by the XML parser + type ParsedSuggestion = string | { "#text": string; "@_mode"?: string } + + let parsedSuggest: { + suggest: ParsedSuggestion[] | ParsedSuggestion + } + + try { + parsedSuggest = parseXml(follow_up, ["suggest"]) as { + suggest: ParsedSuggestion[] | ParsedSuggestion + } + } catch (error) { + cline.consecutiveMistakeCount++ + cline.recordToolError("ask_memory_aware_followup_question") + await cline.say("error", `Failed to parse operations: ${error.message}`) + pushToolResult(formatResponse.toolError("Invalid operations xml format")) + return + } + + const rawSuggestions = Array.isArray(parsedSuggest?.suggest) + ? parsedSuggest.suggest + : [parsedSuggest?.suggest].filter((sug): sug is ParsedSuggestion => sug !== undefined) + + // Transform parsed XML to our Suggest format + const normalizedSuggest: Suggest[] = rawSuggestions.map((sug) => { + if (typeof sug === "string") { + // Simple string suggestion (no mode attribute) + return { answer: sug } + } else { + // XML object with text content and optional mode attribute + const result: Suggest = { answer: sug["#text"] } + if (sug["@_mode"]) { + result.mode = sug["@_mode"] + } + return result + } + }) + + follow_up_json.suggest = normalizedSuggest + } + + // Get relevant memories before asking the question + let memoryContext = "" + try { + const provider = cline.providerRef.deref() + if (provider) { + const codeIndexManager = provider.codeIndexManager + if (codeIndexManager) { + const configManager = new CodeIndexConfigManager(provider.contextProxy) + const cacheManager = new CacheManager(provider.context, cline.workspacePath) + const serviceFactory = new CodeIndexServiceFactory( + configManager, + cline.workspacePath, + cacheManager, + ) + + const memoryManager = MemoryStorageManager.getInstance( + configManager, + serviceFactory, + cline.workspacePath, + ) + + if (memoryManager.isEnabled()) { + const memoryService = await memoryManager.getMemoryStorageService() + if (memoryService) { + // Search for relevant memories + const relevantMemories = await memoryService.searchMemories(question, 5) + + // Format memories for context + if (relevantMemories.length > 0) { + memoryContext = "\n\nBased on previous interactions:\n" + relevantMemories.forEach((memory, index) => { + memoryContext += `${index + 1}. Q: ${memory.question}\n A: ${memory.answer}\n` + }) + } + } + } + } + } + } catch (error) { + console.error("Failed to retrieve memories:", error) + // Continue without memory context + } + + // Add memory context to the question + const questionWithContext = question + memoryContext + + cline.consecutiveMistakeCount = 0 + const { text, images } = await cline.ask( + "followup", + JSON.stringify({ ...follow_up_json, question: questionWithContext }), + false, + ) + await cline.say("user_feedback", text ?? "", images) + pushToolResult(formatResponse.toolResult(`\n${text}\n`, images)) + + // Store memory if enabled (same as askFollowupQuestionTool) + try { + const provider = cline.providerRef.deref() + if (provider && text) { + const codeIndexManager = provider.codeIndexManager + if (codeIndexManager) { + const configManager = new CodeIndexConfigManager(provider.contextProxy) + const cacheManager = new CacheManager(provider.context, cline.workspacePath) + const serviceFactory = new CodeIndexServiceFactory( + configManager, + cline.workspacePath, + cacheManager, + ) + + const memoryManager = MemoryStorageManager.getInstance( + configManager, + serviceFactory, + cline.workspacePath, + ) + + if (memoryManager.isEnabled()) { + const memoryService = await memoryManager.getMemoryStorageService() + if (memoryService) { + await memoryService.storeMemory( + question, + text, + follow_up_json.suggest, + cline.taskId, + await cline.getTaskMode(), + ) + } + } + } + } + } catch (error) { + console.error("Failed to store memory:", error) + } + + return + } + } catch (error) { + await handleError("asking memory-aware question", error) + return + } +} diff --git a/src/core/tools/searchMemoriesTool.ts b/src/core/tools/searchMemoriesTool.ts new file mode 100644 index 0000000000..5b118c2e03 --- /dev/null +++ b/src/core/tools/searchMemoriesTool.ts @@ -0,0 +1,105 @@ +import { Task } from "../task/Task" +import { ToolUse, AskApproval, HandleError, PushToolResult, RemoveClosingTag } from "../../shared/tools" +import { formatResponse } from "../prompts/responses" +import { CodeIndexConfigManager } from "../../services/code-index/config-manager" +import { CodeIndexServiceFactory } from "../../services/code-index/service-factory" +import { MemoryStorageManager } from "../../services/memory-storage/MemoryStorageManager" +import { CacheManager } from "../../services/code-index/cache-manager" +import { ClineSayTool } from "../../shared/ExtensionMessage" + +export async function searchMemoriesTool( + cline: Task, + block: ToolUse, + askApproval: AskApproval, + handleError: HandleError, + pushToolResult: PushToolResult, + removeClosingTag: RemoveClosingTag, +) { + const query: string | undefined = block.params.query + const limitStr: string | undefined = block.params.limit + const limit: number = limitStr ? parseInt(limitStr, 10) : 10 + + const sharedMessageProps: ClineSayTool = { + tool: "codebaseSearch", + query: removeClosingTag("query", query), + } + + try { + if (block.partial) { + const partialMessage = JSON.stringify({ ...sharedMessageProps, content: "" } satisfies ClineSayTool) + await cline.ask("tool", partialMessage, block.partial).catch(() => {}) + return + } else { + if (!query) { + cline.consecutiveMistakeCount++ + pushToolResult(await cline.sayAndCreateMissingParamError("search_memories", "query")) + return + } + + // Get the provider and check if memory storage is enabled + const provider = cline.providerRef.deref() + if (!provider) { + pushToolResult(formatResponse.toolError("Provider not available")) + return + } + + const codeIndexManager = provider.codeIndexManager + if (!codeIndexManager) { + pushToolResult(formatResponse.toolError("Code index manager not available")) + return + } + + // Create necessary managers + const configManager = new CodeIndexConfigManager(provider.contextProxy) + const cacheManager = new CacheManager(provider.context, cline.workspacePath) + const serviceFactory = new CodeIndexServiceFactory(configManager, cline.workspacePath, cacheManager) + + const memoryManager = MemoryStorageManager.getInstance(configManager, serviceFactory, cline.workspacePath) + + if (!memoryManager.isEnabled()) { + pushToolResult(formatResponse.toolError("Memory storage is not enabled")) + return + } + + const memoryService = await memoryManager.getMemoryStorageService() + if (!memoryService) { + pushToolResult(formatResponse.toolError("Memory storage service not available")) + return + } + + cline.consecutiveMistakeCount = 0 + + // Search for memories + const memories = await memoryService.searchMemories(query, limit) + + // Format the results + let resultText = `Found ${memories.length} relevant memories:\n\n` + + if (memories.length === 0) { + resultText = "No relevant memories found." + } else { + memories.forEach((memory, index) => { + resultText += `${index + 1}. Question: ${memory.question}\n` + resultText += ` Answer: ${memory.answer}\n\n` + }) + } + + const completeMessage = JSON.stringify({ + ...sharedMessageProps, + content: resultText, + } satisfies ClineSayTool) + const didApprove = await askApproval("tool", completeMessage) + + if (!didApprove) { + return + } + + pushToolResult(formatResponse.toolResult(resultText)) + + return + } + } catch (error) { + await handleError("searching memories", error) + return + } +} diff --git a/src/core/webview/webviewMessageHandler.ts b/src/core/webview/webviewMessageHandler.ts index fdb7e90425..98f2d5ed2d 100644 --- a/src/core/webview/webviewMessageHandler.ts +++ b/src/core/webview/webviewMessageHandler.ts @@ -1319,6 +1319,14 @@ export const webviewMessageHandler = async ( await updateGlobalState("includeTaskHistoryInEnhance", message.bool ?? false) await provider.postStateToWebview() break + case "memoryStorageEnabled": + await updateGlobalState("memoryStorageEnabled", message.bool ?? false) + await provider.postStateToWebview() + break + case "memoryStorageAutoApprove": + await updateGlobalState("memoryStorageAutoApprove", message.bool ?? false) + await provider.postStateToWebview() + break case "condensingApiConfigId": await updateGlobalState("condensingApiConfigId", message.text) await provider.postStateToWebview() diff --git a/src/services/code-index/__tests__/config-manager.spec.ts b/src/services/code-index/__tests__/config-manager.spec.ts index 2d6e704d76..dfe4524c0e 100644 --- a/src/services/code-index/__tests__/config-manager.spec.ts +++ b/src/services/code-index/__tests__/config-manager.spec.ts @@ -1292,14 +1292,17 @@ describe("CodeIndexConfigManager", () => { isConfigured: true, embedderProvider: "openai", modelId: "text-embedding-3-large", + modelDimension: undefined, openAiOptions: { openAiNativeApiKey: "test-openai-key" }, ollamaOptions: { ollamaBaseUrl: undefined }, geminiOptions: undefined, + mistralOptions: undefined, openAiCompatibleOptions: undefined, qdrantUrl: "http://qdrant.local", qdrantApiKey: "test-qdrant-key", searchMinScore: 0.4, searchMaxResults: 50, + memoryStorageEnabled: false, }) }) diff --git a/src/services/code-index/__tests__/service-factory.spec.ts b/src/services/code-index/__tests__/service-factory.spec.ts index 1d8f7ba478..c58566a490 100644 --- a/src/services/code-index/__tests__/service-factory.spec.ts +++ b/src/services/code-index/__tests__/service-factory.spec.ts @@ -367,6 +367,7 @@ describe("CodeIndexServiceFactory", () => { "http://localhost:6333", 3072, "test-key", + "codebase", ) }) @@ -392,6 +393,7 @@ describe("CodeIndexServiceFactory", () => { "http://localhost:6333", 768, "test-key", + "codebase", ) }) @@ -417,6 +419,7 @@ describe("CodeIndexServiceFactory", () => { "http://localhost:6333", 3072, "test-key", + "codebase", ) }) @@ -449,6 +452,7 @@ describe("CodeIndexServiceFactory", () => { "http://localhost:6333", modelDimension, // Should use model's built-in dimension, not manual "test-key", + "codebase", ) }) @@ -480,6 +484,7 @@ describe("CodeIndexServiceFactory", () => { "http://localhost:6333", manualDimension, // Should use manual dimension as fallback "test-key", + "codebase", ) }) @@ -509,6 +514,7 @@ describe("CodeIndexServiceFactory", () => { "http://localhost:6333", 768, "test-key", + "codebase", ) }) @@ -578,6 +584,7 @@ describe("CodeIndexServiceFactory", () => { "http://localhost:6333", 3072, "test-key", + "codebase", ) }) @@ -603,6 +610,7 @@ describe("CodeIndexServiceFactory", () => { "http://localhost:6333", 3072, "test-key", + "codebase", ) }) @@ -627,6 +635,7 @@ describe("CodeIndexServiceFactory", () => { "http://localhost:6333", 1536, "test-key", + "codebase", ) }) diff --git a/src/services/code-index/config-manager.ts b/src/services/code-index/config-manager.ts index 1723f1c2a0..cf43b578f2 100644 --- a/src/services/code-index/config-manager.ts +++ b/src/services/code-index/config-manager.ts @@ -23,6 +23,7 @@ export class CodeIndexConfigManager { private qdrantApiKey?: string private searchMinScore?: number private searchMaxResults?: number + private memoryStorageEnabled: boolean = false constructor(private readonly contextProxy: ContextProxy) { // Initialize with current configuration to avoid false restart triggers @@ -50,6 +51,7 @@ export class CodeIndexConfigManager { codebaseIndexEmbedderModelId: "", codebaseIndexSearchMinScore: undefined, codebaseIndexSearchMaxResults: undefined, + memoryStorageEnabled: false, } const { @@ -60,6 +62,7 @@ export class CodeIndexConfigManager { codebaseIndexEmbedderModelId, codebaseIndexSearchMinScore, codebaseIndexSearchMaxResults, + memoryStorageEnabled, } = codebaseIndexConfig const openAiKey = this.contextProxy?.getSecret("codeIndexOpenAiKey") ?? "" @@ -76,6 +79,7 @@ export class CodeIndexConfigManager { this.qdrantApiKey = qdrantApiKey ?? "" this.searchMinScore = codebaseIndexSearchMinScore this.searchMaxResults = codebaseIndexSearchMaxResults + this.memoryStorageEnabled = memoryStorageEnabled ?? false // Validate and set model dimension const rawDimension = codebaseIndexConfig.codebaseIndexEmbedderModelDimension @@ -379,6 +383,7 @@ export class CodeIndexConfigManager { qdrantApiKey: this.qdrantApiKey, searchMinScore: this.currentSearchMinScore, searchMaxResults: this.currentSearchMaxResults, + memoryStorageEnabled: this.memoryStorageEnabled, } } @@ -460,4 +465,11 @@ export class CodeIndexConfigManager { public get currentSearchMaxResults(): number { return this.searchMaxResults ?? DEFAULT_MAX_SEARCH_RESULTS } + + /** + * Gets whether memory storage is enabled for follow-up questions + */ + public get isMemoryStorageEnabled(): boolean { + return this.memoryStorageEnabled + } } diff --git a/src/services/code-index/interfaces/collection-types.ts b/src/services/code-index/interfaces/collection-types.ts new file mode 100644 index 0000000000..bb83d50501 --- /dev/null +++ b/src/services/code-index/interfaces/collection-types.ts @@ -0,0 +1,15 @@ +/** + * Enum for different collection types in Qdrant + */ +export enum QdrantCollectionType { + CODEBASE = "codebase", + MEMORY = "memory", +} + +/** + * Interface for collection configuration + */ +export interface CollectionConfig { + type: QdrantCollectionType + vectorSize: number +} diff --git a/src/services/code-index/interfaces/config.ts b/src/services/code-index/interfaces/config.ts index 9098a60091..62a024b86d 100644 --- a/src/services/code-index/interfaces/config.ts +++ b/src/services/code-index/interfaces/config.ts @@ -18,6 +18,7 @@ export interface CodeIndexConfig { qdrantApiKey?: string searchMinScore?: number searchMaxResults?: number + memoryStorageEnabled?: boolean } /** diff --git a/src/services/code-index/service-factory.ts b/src/services/code-index/service-factory.ts index 68b0f5c0bc..7758db3138 100644 --- a/src/services/code-index/service-factory.ts +++ b/src/services/code-index/service-factory.ts @@ -6,6 +6,7 @@ import { GeminiEmbedder } from "./embedders/gemini" import { MistralEmbedder } from "./embedders/mistral" import { EmbedderProvider, getDefaultModelId, getModelDimension } from "../../shared/embeddingModels" import { QdrantVectorStore } from "./vector-store/qdrant-client" +import { QdrantCollectionType } from "./interfaces/collection-types" import { codeParser, DirectoryScanner, FileWatcher } from "./processors" import { ICodeParser, IEmbedder, IFileWatcher, IVectorStore } from "./interfaces" import { CodeIndexConfigManager } from "./config-manager" @@ -103,8 +104,9 @@ export class CodeIndexServiceFactory { /** * Creates a vector store instance using the current configuration. + * @param collectionType Type of collection to create (defaults to CODEBASE) */ - public createVectorStore(): IVectorStore { + public createVectorStore(collectionType: QdrantCollectionType = QdrantCollectionType.CODEBASE): IVectorStore { const config = this.configManager.getConfig() const provider = config.embedderProvider as EmbedderProvider @@ -136,8 +138,14 @@ export class CodeIndexServiceFactory { throw new Error(t("embeddings:serviceFactory.qdrantUrlMissing")) } - // Assuming constructor is updated: new QdrantVectorStore(workspacePath, url, vectorSize, apiKey?) - return new QdrantVectorStore(this.workspacePath, config.qdrantUrl, vectorSize, config.qdrantApiKey) + // Create vector store with specified collection type + return new QdrantVectorStore( + this.workspacePath, + config.qdrantUrl, + vectorSize, + config.qdrantApiKey, + collectionType, + ) } /** diff --git a/src/services/code-index/vector-store/__tests__/qdrant-client.spec.ts b/src/services/code-index/vector-store/__tests__/qdrant-client.spec.ts index e539c2edde..983ef7dfe1 100644 --- a/src/services/code-index/vector-store/__tests__/qdrant-client.spec.ts +++ b/src/services/code-index/vector-store/__tests__/qdrant-client.spec.ts @@ -4,6 +4,7 @@ import { createHash } from "crypto" import { QdrantVectorStore } from "../qdrant-client" import { getWorkspacePath } from "../../../../utils/path" import { DEFAULT_MAX_SEARCH_RESULTS, DEFAULT_SEARCH_MIN_SCORE } from "../../constants" +import { QdrantCollectionType } from "../../interfaces/collection-types" // Mocks vitest.mock("@qdrant/js-client-rest") @@ -48,7 +49,7 @@ describe("QdrantVectorStore", () => { const mockApiKey = "test-api-key" const mockVectorSize = 1536 const mockHashedPath = "a1b2c3d4e5f6g7h8i9j0k1l2m3n4o5p6" // Needs to be long enough - const expectedCollectionName = `ws-${mockHashedPath.substring(0, 16)}` + const expectedCollectionName = `codebase-${mockHashedPath.substring(0, 16)}` beforeEach(() => { vitest.clearAllMocks() @@ -84,7 +85,23 @@ describe("QdrantVectorStore", () => { // Access private member for testing constructor logic (not ideal, but necessary here) expect((vectorStore as any).collectionName).toBe(expectedCollectionName) expect((vectorStore as any).vectorSize).toBe(mockVectorSize) + expect((vectorStore as any).collectionType).toBe(QdrantCollectionType.CODEBASE) }) + + it("should correctly initialize with MEMORY collection type", () => { + const memoryVectorStore = new QdrantVectorStore( + mockWorkspacePath, + mockQdrantUrl, + mockVectorSize, + mockApiKey, + QdrantCollectionType.MEMORY, + ) + + const expectedMemoryCollectionName = `memory-${mockHashedPath.substring(0, 16)}` + expect((memoryVectorStore as any).collectionName).toBe(expectedMemoryCollectionName) + expect((memoryVectorStore as any).collectionType).toBe(QdrantCollectionType.MEMORY) + }) + it("should handle constructor with default URL when none provided", () => { const vectorStoreWithDefaults = new QdrantVectorStore(mockWorkspacePath, undefined as any, mockVectorSize) diff --git a/src/services/code-index/vector-store/qdrant-client.ts b/src/services/code-index/vector-store/qdrant-client.ts index 0218e37295..e476a124b5 100644 --- a/src/services/code-index/vector-store/qdrant-client.ts +++ b/src/services/code-index/vector-store/qdrant-client.ts @@ -6,6 +6,7 @@ import { IVectorStore } from "../interfaces/vector-store" import { Payload, VectorStoreSearchResult } from "../interfaces" import { DEFAULT_MAX_SEARCH_RESULTS, DEFAULT_SEARCH_MIN_SCORE } from "../constants" import { t } from "../../../i18n" +import { QdrantCollectionType } from "../interfaces/collection-types" /** * Qdrant implementation of the vector store interface @@ -17,13 +18,23 @@ export class QdrantVectorStore implements IVectorStore { private client: QdrantClient private readonly collectionName: string private readonly qdrantUrl: string = "http://localhost:6333" + private readonly collectionType: QdrantCollectionType /** * Creates a new Qdrant vector store * @param workspacePath Path to the workspace * @param url Optional URL to the Qdrant server + * @param vectorSize Size of the vectors + * @param apiKey Optional API key for authentication + * @param collectionType Type of collection (defaults to CODEBASE) */ - constructor(workspacePath: string, url: string, vectorSize: number, apiKey?: string) { + constructor( + workspacePath: string, + url: string, + vectorSize: number, + apiKey?: string, + collectionType: QdrantCollectionType = QdrantCollectionType.CODEBASE, + ) { // Parse the URL to determine the appropriate QdrantClient configuration const parsedUrl = this.parseQdrantUrl(url) @@ -75,10 +86,11 @@ export class QdrantVectorStore implements IVectorStore { }) } - // Generate collection name from workspace path + // Generate collection name from workspace path and collection type const hash = createHash("sha256").update(workspacePath).digest("hex") this.vectorSize = vectorSize - this.collectionName = `ws-${hash.substring(0, 16)}` + this.collectionType = collectionType + this.collectionName = `${collectionType}-${hash.substring(0, 16)}` } /** diff --git a/src/services/memory-storage/MemoryStorageManager.ts b/src/services/memory-storage/MemoryStorageManager.ts new file mode 100644 index 0000000000..c40646c3c5 --- /dev/null +++ b/src/services/memory-storage/MemoryStorageManager.ts @@ -0,0 +1,69 @@ +import { CodeIndexConfigManager } from "../code-index/config-manager" +import { CodeIndexServiceFactory } from "../code-index/service-factory" +import { MemoryStorageService } from "./MemoryStorageService" + +/** + * Singleton manager for the memory storage service. + * Ensures a single instance is shared across the application. + */ +export class MemoryStorageManager { + private static instance: MemoryStorageManager | null = null + private memoryStorageService: MemoryStorageService | null = null + + private constructor( + private configManager: CodeIndexConfigManager, + private serviceFactory: CodeIndexServiceFactory, + private workspacePath: string, + ) {} + + /** + * Get or create the singleton instance + */ + static getInstance( + configManager: CodeIndexConfigManager, + serviceFactory: CodeIndexServiceFactory, + workspacePath: string, + ): MemoryStorageManager { + if (!MemoryStorageManager.instance) { + MemoryStorageManager.instance = new MemoryStorageManager(configManager, serviceFactory, workspacePath) + } + return MemoryStorageManager.instance + } + + /** + * Get the memory storage service, creating it if necessary + */ + async getMemoryStorageService(): Promise { + if (!this.configManager.isMemoryStorageEnabled) { + return null + } + + if (!this.memoryStorageService) { + this.memoryStorageService = new MemoryStorageService( + this.configManager, + this.serviceFactory, + this.workspacePath, + ) + await this.memoryStorageService.initialize() + } + + return this.memoryStorageService + } + + /** + * Check if memory storage is enabled + */ + isEnabled(): boolean { + return this.configManager.isMemoryStorageEnabled + } + + /** + * Reset the singleton instance (mainly for testing) + */ + static reset(): void { + if (MemoryStorageManager.instance?.memoryStorageService) { + MemoryStorageManager.instance.memoryStorageService.dispose() + } + MemoryStorageManager.instance = null + } +} diff --git a/src/services/memory-storage/MemoryStorageService.ts b/src/services/memory-storage/MemoryStorageService.ts new file mode 100644 index 0000000000..971da6f9f0 --- /dev/null +++ b/src/services/memory-storage/MemoryStorageService.ts @@ -0,0 +1,191 @@ +import { IEmbedder } from "../code-index/interfaces/embedder" +import { IVectorStore, PointStruct } from "../code-index/interfaces/vector-store" +import { CodeIndexServiceFactory } from "../code-index/service-factory" +import { CodeIndexConfigManager } from "../code-index/config-manager" +import { QdrantCollectionType } from "../code-index/interfaces/collection-types" +import { v4 as uuidv4 } from "uuid" + +export interface MemoryEntry { + id: string + question: string + answer: string + suggestions?: Array<{ answer: string; mode?: string }> + timestamp: number + taskId: string + mode?: string +} + +export class MemoryStorageService { + private embedder: IEmbedder | null = null + private vectorStore: IVectorStore | null = null + private isInitialized = false + + constructor( + private readonly configManager: CodeIndexConfigManager, + private readonly serviceFactory: CodeIndexServiceFactory, + private readonly workspacePath: string, + ) {} + + /** + * Initialize the memory storage service + */ + async initialize(): Promise { + if (this.isInitialized) { + return + } + + const config = this.configManager.getConfig() + + // Only initialize if memory storage is enabled and the feature is configured + if (!config.memoryStorageEnabled || !config.isConfigured) { + return + } + + try { + // Create embedder and vector store for memory collection + this.embedder = this.serviceFactory.createEmbedder() + this.vectorStore = this.serviceFactory.createVectorStore(QdrantCollectionType.MEMORY) + + // Initialize the vector store + await this.vectorStore.initialize() + + this.isInitialized = true + } catch (error) { + console.error("Failed to initialize memory storage service:", error) + throw error + } + } + + /** + * Store a question-answer pair in memory + */ + async storeMemory( + question: string, + answer: string, + suggestions: Array<{ answer: string; mode?: string }> = [], + taskId: string, + mode?: string, + ): Promise { + if (!this.isInitialized || !this.embedder || !this.vectorStore) { + await this.initialize() + if (!this.isInitialized) { + // Memory storage is disabled or not configured + return + } + } + + if (!this.embedder || !this.vectorStore) { + throw new Error("Memory storage service not properly initialized") + } + + const memoryEntry: MemoryEntry = { + id: uuidv4(), + question, + answer, + suggestions: suggestions.length > 0 ? suggestions : undefined, + timestamp: Date.now(), + taskId, + mode, + } + + try { + // Create embedding for the question + const embeddingResponse = await this.embedder.createEmbeddings([question]) + const vector = embeddingResponse.embeddings[0] + + // Create point for vector store + const point: PointStruct = { + id: memoryEntry.id, + vector, + payload: memoryEntry, + } + + // Store in vector database + await this.vectorStore.upsertPoints([point]) + } catch (error) { + console.error("Failed to store memory:", error) + throw error + } + } + + /** + * Search for similar questions in memory + */ + async searchSimilarQuestions(question: string, limit: number = 5, scoreThreshold?: number): Promise { + if (!this.isInitialized || !this.embedder || !this.vectorStore) { + await this.initialize() + if (!this.isInitialized) { + // Memory storage is disabled or not configured + return [] + } + } + + if (!this.embedder || !this.vectorStore) { + return [] + } + + try { + // Create embedding for the search query + const embeddingResponse = await this.embedder.createEmbeddings([question]) + const vector = embeddingResponse.embeddings[0] + + // Search in vector database + const results = await this.vectorStore.search( + vector, + undefined, // directoryPrefix + scoreThreshold ?? this.configManager.currentSearchMinScore, + limit, + ) + + // Extract and return memory entries + return results.filter((result) => result.payload).map((result) => result.payload as unknown as MemoryEntry) + } catch (error) { + console.error("Failed to search memories:", error) + return [] + } + } + + /** + * Check if memory storage is enabled and configured + */ + isEnabled(): boolean { + const config = this.configManager.getConfig() + return config.memoryStorageEnabled === true && config.isConfigured + } + + /** + * Clear all memories + */ + async clearMemories(): Promise { + if (!this.vectorStore) { + return + } + + try { + await this.vectorStore.clearCollection() + } catch (error) { + console.error("Failed to clear memories:", error) + throw error + } + } + + /** + * Search for memories (alias for searchSimilarQuestions) + */ + async searchMemories(query: string, limit: number = 5): Promise> { + const memories = await this.searchSimilarQuestions(query, limit) + return memories.map((memory) => ({ + question: memory.question, + answer: memory.answer, + })) + } + + /** + * Dispose of resources + */ + dispose(): void { + this.embedder = null + this.vectorStore = null + this.isInitialized = false + } +} diff --git a/src/shared/ExtensionMessage.ts b/src/shared/ExtensionMessage.ts index 930edeac73..d0f608e093 100644 --- a/src/shared/ExtensionMessage.ts +++ b/src/shared/ExtensionMessage.ts @@ -270,6 +270,8 @@ export type ExtensionState = Pick< | "profileThresholds" | "includeDiagnosticMessages" | "maxDiagnosticMessages" + | "memoryStorageEnabled" + | "memoryStorageAutoApprove" > & { version: string clineMessages: ClineMessage[] diff --git a/src/shared/WebviewMessage.ts b/src/shared/WebviewMessage.ts index cb8759d851..9dd8d3fb43 100644 --- a/src/shared/WebviewMessage.ts +++ b/src/shared/WebviewMessage.ts @@ -142,6 +142,8 @@ export interface WebviewMessage { | "systemPrompt" | "enhancementApiConfigId" | "includeTaskHistoryInEnhance" + | "memoryStorageEnabled" + | "memoryStorageAutoApprove" | "updateExperimental" | "autoApprovalEnabled" | "updateCustomMode" diff --git a/src/shared/tools.ts b/src/shared/tools.ts index 67972243fe..235e8930a9 100644 --- a/src/shared/tools.ts +++ b/src/shared/tools.ts @@ -65,6 +65,7 @@ export const toolParamNames = [ "query", "args", "todos", + "limit", ] as const export type ToolParamName = (typeof toolParamNames)[number] @@ -143,6 +144,16 @@ export interface AskFollowupQuestionToolUse extends ToolUse { params: Partial, "question" | "follow_up">> } +export interface AskMemoryAwareFollowupQuestionToolUse extends ToolUse { + name: "ask_memory_aware_followup_question" + params: Partial, "question" | "follow_up">> +} + +export interface SearchMemoriesToolUse extends ToolUse { + name: "search_memories" + params: Partial, "query" | "limit">> +} + export interface AttemptCompletionToolUse extends ToolUse { name: "attempt_completion" params: Partial, "result">> @@ -183,6 +194,8 @@ export const TOOL_DISPLAY_NAMES: Record = { use_mcp_tool: "use mcp tools", access_mcp_resource: "access mcp resources", ask_followup_question: "ask questions", + ask_memory_aware_followup_question: "ask memory-aware questions", + search_memories: "search memories", attempt_completion: "complete tasks", switch_mode: "switch modes", new_task: "create new task", @@ -225,6 +238,8 @@ export const TOOL_GROUPS: Record = { // Tools that are always available to all modes. export const ALWAYS_AVAILABLE_TOOLS: ToolName[] = [ "ask_followup_question", + "ask_memory_aware_followup_question", + "search_memories", "attempt_completion", "switch_mode", "new_task", diff --git a/webview-ui/src/components/settings/ContextManagementSettings.tsx b/webview-ui/src/components/settings/ContextManagementSettings.tsx index 88484e1d63..4dc55aba74 100644 --- a/webview-ui/src/components/settings/ContextManagementSettings.tsx +++ b/webview-ui/src/components/settings/ContextManagementSettings.tsx @@ -2,7 +2,7 @@ import { HTMLAttributes } from "react" import React from "react" import { useAppTranslation } from "@/i18n/TranslationContext" import { VSCodeCheckbox } from "@vscode/webview-ui-toolkit/react" -import { Database, FoldVertical } from "lucide-react" +import { Database, FoldVertical, Brain } from "lucide-react" import { cn } from "@/lib/utils" import { Input, Select, SelectContent, SelectItem, SelectTrigger, SelectValue, Slider, Button } from "@/components/ui" @@ -27,6 +27,8 @@ type ContextManagementSettingsProps = HTMLAttributes & { includeDiagnosticMessages?: boolean maxDiagnosticMessages?: number writeDelayMs: number + memoryStorageEnabled?: boolean + memoryStorageAutoApprove?: boolean setCachedStateField: SetCachedStateField< | "autoCondenseContext" | "autoCondenseContextPercent" @@ -41,6 +43,8 @@ type ContextManagementSettingsProps = HTMLAttributes & { | "includeDiagnosticMessages" | "maxDiagnosticMessages" | "writeDelayMs" + | "memoryStorageEnabled" + | "memoryStorageAutoApprove" > } @@ -60,6 +64,8 @@ export const ContextManagementSettings = ({ includeDiagnosticMessages, maxDiagnosticMessages, writeDelayMs, + memoryStorageEnabled, + memoryStorageAutoApprove, className, ...props }: ContextManagementSettingsProps) => { @@ -438,6 +444,36 @@ export const ContextManagementSettings = ({ )} +
+
+ +
{t("settings:contextManagement.memoryStorage.title")}
+
+ setCachedStateField("memoryStorageEnabled", e.target.checked)} + data-testid="memory-storage-enabled-checkbox"> + {t("settings:contextManagement.memoryStorage.enable.label")} + +
+ {t("settings:contextManagement.memoryStorage.enable.description")} +
+ {memoryStorageEnabled && ( +
+ setCachedStateField("memoryStorageAutoApprove", e.target.checked)} + data-testid="memory-storage-auto-approve-checkbox"> + + {t("settings:contextManagement.memoryStorage.autoApprove.label")} + + +
+ {t("settings:contextManagement.memoryStorage.autoApprove.description")} +
+
+ )} +
) } diff --git a/webview-ui/src/components/settings/SettingsView.tsx b/webview-ui/src/components/settings/SettingsView.tsx index 630b59485d..81879c56aa 100644 --- a/webview-ui/src/components/settings/SettingsView.tsx +++ b/webview-ui/src/components/settings/SettingsView.tsx @@ -183,6 +183,8 @@ const SettingsView = forwardRef(({ onDone, t includeDiagnosticMessages, maxDiagnosticMessages, includeTaskHistoryInEnhance, + memoryStorageEnabled, + memoryStorageAutoApprove, } = cachedState const apiConfiguration = useMemo(() => cachedState.apiConfiguration ?? {}, [cachedState.apiConfiguration]) @@ -342,6 +344,8 @@ const SettingsView = forwardRef(({ onDone, t vscode.postMessage({ type: "updateCondensingPrompt", text: customCondensingPrompt || "" }) vscode.postMessage({ type: "updateSupportPrompt", values: customSupportPrompts || {} }) vscode.postMessage({ type: "includeTaskHistoryInEnhance", bool: includeTaskHistoryInEnhance ?? false }) + vscode.postMessage({ type: "memoryStorageEnabled", bool: memoryStorageEnabled }) + vscode.postMessage({ type: "memoryStorageAutoApprove", bool: memoryStorageAutoApprove }) vscode.postMessage({ type: "upsertApiConfiguration", text: currentApiConfigName, apiConfiguration }) vscode.postMessage({ type: "telemetrySetting", text: telemetrySetting }) vscode.postMessage({ type: "profileThresholds", values: profileThresholds }) @@ -682,6 +686,8 @@ const SettingsView = forwardRef(({ onDone, t includeDiagnosticMessages={includeDiagnosticMessages} maxDiagnosticMessages={maxDiagnosticMessages} writeDelayMs={writeDelayMs} + memoryStorageEnabled={memoryStorageEnabled} + memoryStorageAutoApprove={memoryStorageAutoApprove} setCachedStateField={setCachedStateField} /> )} diff --git a/webview-ui/src/context/ExtensionStateContext.tsx b/webview-ui/src/context/ExtensionStateContext.tsx index da7ab63358..5b1e77da8d 100644 --- a/webview-ui/src/context/ExtensionStateContext.tsx +++ b/webview-ui/src/context/ExtensionStateContext.tsx @@ -147,6 +147,10 @@ export interface ExtensionStateContextType extends ExtensionState { setMaxDiagnosticMessages: (value: number) => void includeTaskHistoryInEnhance?: boolean setIncludeTaskHistoryInEnhance: (value: boolean) => void + memoryStorageEnabled?: boolean + setMemoryStorageEnabled: (value: boolean) => void + memoryStorageAutoApprove?: boolean + setMemoryStorageAutoApprove: (value: boolean) => void } export const ExtensionStateContext = createContext(undefined) @@ -247,6 +251,8 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode alwaysAllowUpdateTodoList: true, includeDiagnosticMessages: true, maxDiagnosticMessages: 50, + memoryStorageEnabled: false, + memoryStorageAutoApprove: false, }) const [didHydrateState, setDidHydrateState] = useState(false) @@ -266,6 +272,8 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode global: {}, }) const [includeTaskHistoryInEnhance, setIncludeTaskHistoryInEnhance] = useState(false) + const [memoryStorageEnabled, setMemoryStorageEnabled] = useState(false) + const [memoryStorageAutoApprove, setMemoryStorageAutoApprove] = useState(false) const setListApiConfigMeta = useCallback( (value: ProviderSettingsEntry[]) => setState((prevState) => ({ ...prevState, listApiConfigMeta: value })), @@ -303,6 +311,14 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode if ((newState as any).includeTaskHistoryInEnhance !== undefined) { setIncludeTaskHistoryInEnhance((newState as any).includeTaskHistoryInEnhance) } + // Update memoryStorageEnabled if present in state message + if ((newState as any).memoryStorageEnabled !== undefined) { + setMemoryStorageEnabled((newState as any).memoryStorageEnabled) + } + // Update memoryStorageAutoApprove if present in state message + if ((newState as any).memoryStorageAutoApprove !== undefined) { + setMemoryStorageAutoApprove((newState as any).memoryStorageAutoApprove) + } // Handle marketplace data if present in state message if (newState.marketplaceItems !== undefined) { setMarketplaceItems(newState.marketplaceItems) @@ -517,6 +533,10 @@ export const ExtensionStateContextProvider: React.FC<{ children: React.ReactNode }, includeTaskHistoryInEnhance, setIncludeTaskHistoryInEnhance, + memoryStorageEnabled, + setMemoryStorageEnabled, + memoryStorageAutoApprove, + setMemoryStorageAutoApprove, } return {children} diff --git a/webview-ui/src/i18n/locales/en/settings.json b/webview-ui/src/i18n/locales/en/settings.json index 46c15556c8..fa7871694a 100644 --- a/webview-ui/src/i18n/locales/en/settings.json +++ b/webview-ui/src/i18n/locales/en/settings.json @@ -559,6 +559,17 @@ "profileDescription": "Custom threshold for this profile only (overrides global default)", "inheritDescription": "This profile inherits the global default threshold ({{threshold}}%)", "usesGlobal": "(uses global {{threshold}}%)" + }, + "memoryStorage": { + "title": "Memory Storage", + "enable": { + "label": "Enable memory storage", + "description": "When enabled, Roo will remember your answers to follow-up questions and use them to provide more personalized suggestions in the future." + }, + "autoApprove": { + "label": "Auto-approve memory storage", + "description": "When enabled, Roo will automatically store question-answer pairs without asking for confirmation each time." + } } }, "terminal": {