Skip to content

Commit fc60262

Browse files
committed
feat: implement memory storage for follow-up questions (#6553)
- Add QdrantCollectionType enum to separate CODEBASE and MEMORY collections - Update QdrantVectorStore to support different collection types - Add memory storage settings to global configuration - Create MemoryStorageService for storing and retrieving Q&A pairs - Create MemoryStorageManager singleton for service management - Update askFollowupQuestionTool to store memories when enabled - Add askMemoryAwareFollowupQuestionTool for memory-aware questions - Add searchMemoriesTool for searching stored memories - Add UI controls in settings for enabling/disabling memory storage - Update tool registration to conditionally include memory tools - Add comprehensive test coverage for all changes This implementation allows Roo to learn from user decisions when answering follow-up questions, providing more personalized suggestions over time.
1 parent 305a5da commit fc60262

25 files changed

+814
-9
lines changed

packages/types/src/codebase-index.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,8 @@ export const codebaseIndexConfigSchema = z.object({
3434
// OpenAI Compatible specific fields
3535
codebaseIndexOpenAiCompatibleBaseUrl: z.string().optional(),
3636
codebaseIndexOpenAiCompatibleModelDimension: z.number().optional(),
37+
// Memory storage settings
38+
memoryStorageEnabled: z.boolean().optional(),
3739
})
3840

3941
export type CodebaseIndexConfig = z.infer<typeof codebaseIndexConfigSchema>

packages/types/src/global-settings.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,10 @@ export const globalSettingsSchema = z.object({
6262
alwaysAllowFollowupQuestions: z.boolean().optional(),
6363
followupAutoApproveTimeoutMs: z.number().optional(),
6464
alwaysAllowUpdateTodoList: z.boolean().optional(),
65+
66+
// Memory storage settings
67+
memoryStorageEnabled: z.boolean().optional(),
68+
memoryStorageAutoApprove: z.boolean().optional(),
6569
allowedCommands: z.array(z.string()).optional(),
6670
deniedCommands: z.array(z.string()).optional(),
6771
commandExecutionTimeout: z.number().optional(),
@@ -240,6 +244,8 @@ export const EVALS_SETTINGS: RooCodeSettings = {
240244
alwaysAllowFollowupQuestions: true,
241245
alwaysAllowUpdateTodoList: true,
242246
followupAutoApproveTimeoutMs: 0,
247+
memoryStorageEnabled: false,
248+
memoryStorageAutoApprove: false,
243249
allowedCommands: ["*"],
244250
commandExecutionTimeout: 20,
245251
commandTimeoutAllowlist: [],

packages/types/src/tool.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,8 @@ export const toolNames = [
2828
"use_mcp_tool",
2929
"access_mcp_resource",
3030
"ask_followup_question",
31+
"ask_memory_aware_followup_question",
32+
"search_memories",
3133
"attempt_completion",
3234
"switch_mode",
3335
"new_task",

src/core/assistant-message/presentAssistantMessage.ts

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@ import { executeCommandTool } from "../tools/executeCommandTool"
2121
import { useMcpToolTool } from "../tools/useMcpToolTool"
2222
import { accessMcpResourceTool } from "../tools/accessMcpResourceTool"
2323
import { askFollowupQuestionTool } from "../tools/askFollowupQuestionTool"
24+
import { askMemoryAwareFollowupQuestionTool } from "../tools/askMemoryAwareFollowupQuestionTool"
25+
import { searchMemoriesTool } from "../tools/searchMemoriesTool"
2426
import { switchModeTool } from "../tools/switchModeTool"
2527
import { attemptCompletionTool } from "../tools/attemptCompletionTool"
2628
import { newTaskTool } from "../tools/newTaskTool"
@@ -151,6 +153,10 @@ export async function presentAssistantMessage(cline: Task) {
151153
break
152154
}
153155
case "tool_use":
156+
// Get customModes early for use in toolDescription
157+
const stateForDescription = await cline.providerRef.deref()?.getState()
158+
const customModesForDescription = stateForDescription?.customModes ?? []
159+
154160
const toolDescription = (): string => {
155161
switch (block.name) {
156162
case "execute_command":
@@ -200,6 +206,10 @@ export async function presentAssistantMessage(cline: Task) {
200206
return `[${block.name} for '${block.params.server_name}']`
201207
case "ask_followup_question":
202208
return `[${block.name} for '${block.params.question}']`
209+
case "ask_memory_aware_followup_question":
210+
return `[${block.name} for '${block.params.question}']`
211+
case "search_memories":
212+
return `[${block.name} for '${block.params.query}']`
203213
case "attempt_completion":
204214
return `[${block.name}]`
205215
case "switch_mode":
@@ -211,7 +221,7 @@ export async function presentAssistantMessage(cline: Task) {
211221
case "new_task": {
212222
const mode = block.params.mode ?? defaultModeSlug
213223
const message = block.params.message ?? "(no message)"
214-
const modeName = getModeBySlug(mode, customModes)?.name ?? mode
224+
const modeName = getModeBySlug(mode, customModesForDescription)?.name ?? mode
215225
return `[${block.name} in ${modeName} mode: '${message}']`
216226
}
217227
}
@@ -504,6 +514,19 @@ export async function presentAssistantMessage(cline: Task) {
504514
removeClosingTag,
505515
)
506516
break
517+
case "ask_memory_aware_followup_question":
518+
await askMemoryAwareFollowupQuestionTool(
519+
cline,
520+
block,
521+
askApproval,
522+
handleError,
523+
pushToolResult,
524+
removeClosingTag,
525+
)
526+
break
527+
case "search_memories":
528+
await searchMemoriesTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
529+
break
507530
case "switch_mode":
508531
await switchModeTool(cline, block, askApproval, handleError, pushToolResult, removeClosingTag)
509532
break

src/core/tools/askFollowupQuestionTool.ts

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,10 @@ import { Task } from "../task/Task"
22
import { ToolUse, AskApproval, HandleError, PushToolResult, RemoveClosingTag } from "../../shared/tools"
33
import { formatResponse } from "../prompts/responses"
44
import { parseXml } from "../../utils/xml"
5+
import { CodeIndexConfigManager } from "../../services/code-index/config-manager"
6+
import { CodeIndexServiceFactory } from "../../services/code-index/service-factory"
7+
import { MemoryStorageManager } from "../../services/memory-storage/MemoryStorageManager"
8+
import { CacheManager } from "../../services/code-index/cache-manager"
59

610
export async function askFollowupQuestionTool(
711
cline: Task,
@@ -80,6 +84,49 @@ export async function askFollowupQuestionTool(
8084
await cline.say("user_feedback", text ?? "", images)
8185
pushToolResult(formatResponse.toolResult(`<answer>\n${text}\n</answer>`, images))
8286

87+
// Store memory if enabled
88+
try {
89+
const provider = cline.providerRef.deref()
90+
if (provider && text) {
91+
// Get the code index manager from the provider
92+
const codeIndexManager = provider.codeIndexManager
93+
if (codeIndexManager) {
94+
// Create config manager and service factory
95+
const configManager = new CodeIndexConfigManager(provider.contextProxy)
96+
const cacheManager = new CacheManager(provider.context, cline.workspacePath)
97+
const serviceFactory = new CodeIndexServiceFactory(
98+
configManager,
99+
cline.workspacePath,
100+
cacheManager,
101+
)
102+
103+
// Get or create the memory storage manager
104+
const memoryManager = MemoryStorageManager.getInstance(
105+
configManager,
106+
serviceFactory,
107+
cline.workspacePath,
108+
)
109+
110+
// Store the memory if enabled
111+
if (memoryManager.isEnabled()) {
112+
const memoryService = await memoryManager.getMemoryStorageService()
113+
if (memoryService) {
114+
await memoryService.storeMemory(
115+
question,
116+
text,
117+
follow_up_json.suggest,
118+
cline.taskId,
119+
await cline.getTaskMode(),
120+
)
121+
}
122+
}
123+
}
124+
}
125+
} catch (error) {
126+
// Log error but don't fail the tool
127+
console.error("Failed to store memory:", error)
128+
}
129+
83130
return
84131
}
85132
} catch (error) {
Lines changed: 183 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,183 @@
1+
import { Task } from "../task/Task"
2+
import { ToolUse, AskApproval, HandleError, PushToolResult, RemoveClosingTag } from "../../shared/tools"
3+
import { formatResponse } from "../prompts/responses"
4+
import { parseXml } from "../../utils/xml"
5+
import { CodeIndexConfigManager } from "../../services/code-index/config-manager"
6+
import { CodeIndexServiceFactory } from "../../services/code-index/service-factory"
7+
import { MemoryStorageManager } from "../../services/memory-storage/MemoryStorageManager"
8+
import { CacheManager } from "../../services/code-index/cache-manager"
9+
10+
export async function askMemoryAwareFollowupQuestionTool(
11+
cline: Task,
12+
block: ToolUse,
13+
askApproval: AskApproval,
14+
handleError: HandleError,
15+
pushToolResult: PushToolResult,
16+
removeClosingTag: RemoveClosingTag,
17+
) {
18+
const question: string | undefined = block.params.question
19+
const follow_up: string | undefined = block.params.follow_up
20+
21+
try {
22+
if (block.partial) {
23+
await cline.ask("followup", removeClosingTag("question", question), block.partial).catch(() => {})
24+
return
25+
} else {
26+
if (!question) {
27+
cline.consecutiveMistakeCount++
28+
cline.recordToolError("ask_memory_aware_followup_question")
29+
pushToolResult(
30+
await cline.sayAndCreateMissingParamError("ask_memory_aware_followup_question", "question"),
31+
)
32+
return
33+
}
34+
35+
type Suggest = { answer: string; mode?: string }
36+
37+
let follow_up_json = {
38+
question,
39+
suggest: [] as Suggest[],
40+
}
41+
42+
if (follow_up) {
43+
// Define the actual structure returned by the XML parser
44+
type ParsedSuggestion = string | { "#text": string; "@_mode"?: string }
45+
46+
let parsedSuggest: {
47+
suggest: ParsedSuggestion[] | ParsedSuggestion
48+
}
49+
50+
try {
51+
parsedSuggest = parseXml(follow_up, ["suggest"]) as {
52+
suggest: ParsedSuggestion[] | ParsedSuggestion
53+
}
54+
} catch (error) {
55+
cline.consecutiveMistakeCount++
56+
cline.recordToolError("ask_memory_aware_followup_question")
57+
await cline.say("error", `Failed to parse operations: ${error.message}`)
58+
pushToolResult(formatResponse.toolError("Invalid operations xml format"))
59+
return
60+
}
61+
62+
const rawSuggestions = Array.isArray(parsedSuggest?.suggest)
63+
? parsedSuggest.suggest
64+
: [parsedSuggest?.suggest].filter((sug): sug is ParsedSuggestion => sug !== undefined)
65+
66+
// Transform parsed XML to our Suggest format
67+
const normalizedSuggest: Suggest[] = rawSuggestions.map((sug) => {
68+
if (typeof sug === "string") {
69+
// Simple string suggestion (no mode attribute)
70+
return { answer: sug }
71+
} else {
72+
// XML object with text content and optional mode attribute
73+
const result: Suggest = { answer: sug["#text"] }
74+
if (sug["@_mode"]) {
75+
result.mode = sug["@_mode"]
76+
}
77+
return result
78+
}
79+
})
80+
81+
follow_up_json.suggest = normalizedSuggest
82+
}
83+
84+
// Get relevant memories before asking the question
85+
let memoryContext = ""
86+
try {
87+
const provider = cline.providerRef.deref()
88+
if (provider) {
89+
const codeIndexManager = provider.codeIndexManager
90+
if (codeIndexManager) {
91+
const configManager = new CodeIndexConfigManager(provider.contextProxy)
92+
const cacheManager = new CacheManager(provider.context, cline.workspacePath)
93+
const serviceFactory = new CodeIndexServiceFactory(
94+
configManager,
95+
cline.workspacePath,
96+
cacheManager,
97+
)
98+
99+
const memoryManager = MemoryStorageManager.getInstance(
100+
configManager,
101+
serviceFactory,
102+
cline.workspacePath,
103+
)
104+
105+
if (memoryManager.isEnabled()) {
106+
const memoryService = await memoryManager.getMemoryStorageService()
107+
if (memoryService) {
108+
// Search for relevant memories
109+
const relevantMemories = await memoryService.searchMemories(question, 5)
110+
111+
// Format memories for context
112+
if (relevantMemories.length > 0) {
113+
memoryContext = "\n\nBased on previous interactions:\n"
114+
relevantMemories.forEach((memory, index) => {
115+
memoryContext += `${index + 1}. Q: ${memory.question}\n A: ${memory.answer}\n`
116+
})
117+
}
118+
}
119+
}
120+
}
121+
}
122+
} catch (error) {
123+
console.error("Failed to retrieve memories:", error)
124+
// Continue without memory context
125+
}
126+
127+
// Add memory context to the question
128+
const questionWithContext = question + memoryContext
129+
130+
cline.consecutiveMistakeCount = 0
131+
const { text, images } = await cline.ask(
132+
"followup",
133+
JSON.stringify({ ...follow_up_json, question: questionWithContext }),
134+
false,
135+
)
136+
await cline.say("user_feedback", text ?? "", images)
137+
pushToolResult(formatResponse.toolResult(`<answer>\n${text}\n</answer>`, images))
138+
139+
// Store memory if enabled (same as askFollowupQuestionTool)
140+
try {
141+
const provider = cline.providerRef.deref()
142+
if (provider && text) {
143+
const codeIndexManager = provider.codeIndexManager
144+
if (codeIndexManager) {
145+
const configManager = new CodeIndexConfigManager(provider.contextProxy)
146+
const cacheManager = new CacheManager(provider.context, cline.workspacePath)
147+
const serviceFactory = new CodeIndexServiceFactory(
148+
configManager,
149+
cline.workspacePath,
150+
cacheManager,
151+
)
152+
153+
const memoryManager = MemoryStorageManager.getInstance(
154+
configManager,
155+
serviceFactory,
156+
cline.workspacePath,
157+
)
158+
159+
if (memoryManager.isEnabled()) {
160+
const memoryService = await memoryManager.getMemoryStorageService()
161+
if (memoryService) {
162+
await memoryService.storeMemory(
163+
question,
164+
text,
165+
follow_up_json.suggest,
166+
cline.taskId,
167+
await cline.getTaskMode(),
168+
)
169+
}
170+
}
171+
}
172+
}
173+
} catch (error) {
174+
console.error("Failed to store memory:", error)
175+
}
176+
177+
return
178+
}
179+
} catch (error) {
180+
await handleError("asking memory-aware question", error)
181+
return
182+
}
183+
}

0 commit comments

Comments
 (0)