Skip to content

Commit 7e456cb

Browse files
committed
feat: add query-context smart tool selection
- Add ToolSelectionAnalyzer class for intelligent tool filtering - Analyze query complexity (simple/moderate/complex) - Detect tool mentions and group needs - Score tools based on relevance (0-1 scale) - Select 6-12 most relevant tools per query - Add configuration settings for the feature - Include comprehensive unit tests Fixes #9117
1 parent e98f4b9 commit 7e456cb

File tree

7 files changed

+765
-2
lines changed

7 files changed

+765
-2
lines changed

packages/types/src/provider-settings.ts

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,11 @@ const baseProviderSettingsSchema = z.object({
182182

183183
// Model verbosity.
184184
verbosity: verbosityLevelsSchema.optional(),
185+
186+
// Smart tool selection.
187+
smartToolSelectionEnabled: z.boolean().optional(),
188+
smartToolSelectionMinTools: z.number().min(1).max(20).optional(),
189+
smartToolSelectionMaxTools: z.number().min(1).max(20).optional(),
185190
})
186191

187192
// Several of the providers share common model config properties.

src/core/prompts/system.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ async function generatePrompt(
6262
settings?: SystemPromptSettings,
6363
todoList?: TodoItem[],
6464
modelId?: string,
65+
userQuery?: string,
6566
): Promise<string> {
6667
if (!context) {
6768
throw new Error("Extension context is required for generating system prompt")
@@ -108,6 +109,7 @@ ${getToolDescriptionsForMode(
108109
settings,
109110
enableMcpServerCreation,
110111
modelId,
112+
userQuery,
111113
)}
112114
113115
${getToolUseGuidelinesSection(codeIndexManager)}
@@ -153,6 +155,7 @@ export const SYSTEM_PROMPT = async (
153155
settings?: SystemPromptSettings,
154156
todoList?: TodoItem[],
155157
modelId?: string,
158+
userQuery?: string,
156159
): Promise<string> => {
157160
if (!context) {
158161
throw new Error("Extension context is required for generating system prompt")
@@ -225,5 +228,6 @@ ${customInstructions}`
225228
settings,
226229
todoList,
227230
modelId,
231+
userQuery,
228232
)
229233
}
Lines changed: 308 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,308 @@
1+
import { describe, it, expect, beforeEach } from "vitest"
2+
import { ToolSelectionAnalyzer, SmartToolSelectionConfig } from "./ToolSelectionAnalyzer"
3+
import type { ModeConfig, ToolName, ToolGroup } from "@roo-code/types"
4+
5+
describe("ToolSelectionAnalyzer", () => {
6+
let analyzer: ToolSelectionAnalyzer
7+
8+
beforeEach(() => {
9+
analyzer = new ToolSelectionAnalyzer()
10+
})
11+
12+
const createModeConfig = (groups: ToolGroup[]): ModeConfig => ({
13+
slug: "test",
14+
name: "Test Mode",
15+
roleDefinition: "Test role",
16+
groups: groups as any, // Groups can be string or tuple format
17+
})
18+
19+
const createAvailableTools = (tools: ToolName[]): Set<ToolName> => new Set(tools)
20+
21+
describe("query analysis", () => {
22+
it("should detect simple queries", () => {
23+
const simpleQueries = [
24+
"what does this function do?",
25+
"explain this code",
26+
"what is the purpose of this file?",
27+
"can you summarize this?",
28+
]
29+
30+
const modeConfig = createModeConfig(["read"])
31+
const availableTools = createAvailableTools([
32+
"read_file",
33+
"list_files",
34+
"search_files",
35+
"list_code_definition_names",
36+
"ask_followup_question",
37+
"attempt_completion",
38+
])
39+
40+
simpleQueries.forEach((query) => {
41+
const selected = analyzer.selectTools(query, modeConfig, availableTools)
42+
expect(selected.length).toBeGreaterThanOrEqual(6)
43+
expect(selected.length).toBeLessThanOrEqual(8)
44+
// Should include essential read tools
45+
expect(selected).toContain("read_file")
46+
expect(selected).toContain("list_files")
47+
})
48+
})
49+
50+
it("should detect complex queries", () => {
51+
const complexQueries = [
52+
"refactor the entire authentication system to use JWT tokens",
53+
"implement a comprehensive logging system across all modules",
54+
"redesign the database schema and migrate all existing data",
55+
"build a complete REST API with authentication and rate limiting",
56+
]
57+
58+
const modeConfig = createModeConfig(["read", "edit"])
59+
const availableTools = createAvailableTools([
60+
"read_file",
61+
"write_to_file",
62+
"apply_diff",
63+
"list_files",
64+
"search_files",
65+
"execute_command",
66+
"ask_followup_question",
67+
"attempt_completion",
68+
])
69+
70+
complexQueries.forEach((query) => {
71+
const selected = analyzer.selectTools(query, modeConfig, availableTools)
72+
expect(selected.length).toBeGreaterThanOrEqual(8)
73+
expect(selected.length).toBeLessThanOrEqual(12)
74+
// Should include both read and edit tools
75+
expect(selected).toContain("read_file")
76+
expect(selected).toContain("write_to_file")
77+
expect(selected).toContain("apply_diff")
78+
})
79+
})
80+
})
81+
82+
describe("tool-specific queries", () => {
83+
it("should prioritize mentioned tools", () => {
84+
const query = "run npm test and show me the results"
85+
const modeConfig = createModeConfig(["read", "command"])
86+
const availableTools = createAvailableTools([
87+
"read_file",
88+
"list_files",
89+
"execute_command",
90+
"ask_followup_question",
91+
"attempt_completion",
92+
])
93+
94+
const selected = analyzer.selectTools(query, modeConfig, availableTools)
95+
// execute_command should be included and highly ranked (top 5) due to explicit mention
96+
expect(selected).toContain("execute_command")
97+
const executeIndex = selected.indexOf("execute_command")
98+
expect(executeIndex).toBeGreaterThanOrEqual(0)
99+
expect(executeIndex).toBeLessThan(5) // Within top 5 is still highly ranked
100+
})
101+
102+
it("should detect file operations", () => {
103+
const query = "create a new config file with the settings I provided"
104+
const modeConfig = createModeConfig(["read", "edit"])
105+
const availableTools = createAvailableTools([
106+
"read_file",
107+
"write_to_file",
108+
"apply_diff",
109+
"list_files",
110+
"search_files",
111+
"ask_followup_question",
112+
"attempt_completion",
113+
])
114+
115+
const selected = analyzer.selectTools(query, modeConfig, availableTools)
116+
// write_to_file should be highly ranked
117+
expect(selected.slice(0, 3)).toContain("write_to_file")
118+
})
119+
120+
it("should detect search operations", () => {
121+
const query = "find all places where the API key is used"
122+
const modeConfig = createModeConfig(["read"])
123+
const availableTools = createAvailableTools([
124+
"read_file",
125+
"list_files",
126+
"search_files",
127+
"list_code_definition_names",
128+
"ask_followup_question",
129+
"attempt_completion",
130+
])
131+
132+
const selected = analyzer.selectTools(query, modeConfig, availableTools)
133+
// search_files should be included and highly ranked (top 5) due to explicit search intent
134+
expect(selected).toContain("search_files")
135+
const searchIndex = selected.indexOf("search_files")
136+
expect(searchIndex).toBeGreaterThanOrEqual(0)
137+
expect(searchIndex).toBeLessThan(5) // Within top 5 is still highly ranked
138+
})
139+
})
140+
141+
describe("settings configuration", () => {
142+
it("should respect minimum tools setting", () => {
143+
const query = "simple task"
144+
const modeConfig = createModeConfig(["read"])
145+
const availableTools = createAvailableTools([
146+
"read_file",
147+
"list_files",
148+
"ask_followup_question",
149+
"attempt_completion",
150+
])
151+
152+
const customAnalyzer = new ToolSelectionAnalyzer({
153+
enabled: true,
154+
minTools: 3,
155+
maxTools: 12,
156+
})
157+
158+
const selected = customAnalyzer.selectTools(query, modeConfig, availableTools)
159+
expect(selected.length).toBeGreaterThanOrEqual(3)
160+
})
161+
162+
it("should respect maximum tools setting", () => {
163+
const query = "complex refactoring task involving multiple systems"
164+
const modeConfig = createModeConfig(["read", "edit", "command"])
165+
const availableTools = createAvailableTools([
166+
"read_file",
167+
"write_to_file",
168+
"apply_diff",
169+
"insert_content",
170+
"list_files",
171+
"search_files",
172+
"list_code_definition_names",
173+
"execute_command",
174+
"ask_followup_question",
175+
"attempt_completion",
176+
"switch_mode",
177+
"new_task",
178+
"update_todo_list",
179+
])
180+
181+
const customAnalyzer = new ToolSelectionAnalyzer({
182+
enabled: true,
183+
minTools: 6,
184+
maxTools: 8,
185+
})
186+
187+
const selected = customAnalyzer.selectTools(query, modeConfig, availableTools)
188+
expect(selected.length).toBeLessThanOrEqual(8)
189+
})
190+
191+
it("should return all tools when smart selection is disabled", () => {
192+
const query = "some task"
193+
const modeConfig = createModeConfig(["read"])
194+
const availableTools = createAvailableTools([
195+
"read_file",
196+
"list_files",
197+
"search_files",
198+
"ask_followup_question",
199+
"attempt_completion",
200+
])
201+
202+
const disabledAnalyzer = new ToolSelectionAnalyzer({
203+
enabled: false,
204+
})
205+
206+
const selected = disabledAnalyzer.selectTools(query, modeConfig, availableTools)
207+
expect(selected.length).toBe(availableTools.size)
208+
expect(selected.sort()).toEqual(Array.from(availableTools).sort())
209+
})
210+
})
211+
212+
describe("essential tools", () => {
213+
it("should always include ask_followup_question and attempt_completion", () => {
214+
const queries = ["simple query", "complex refactoring", "create a file", "run tests"]
215+
216+
const modeConfig = createModeConfig(["read", "edit"])
217+
const availableTools = createAvailableTools([
218+
"read_file",
219+
"write_to_file",
220+
"list_files",
221+
"ask_followup_question",
222+
"attempt_completion",
223+
])
224+
225+
queries.forEach((query) => {
226+
const selected = analyzer.selectTools(query, modeConfig, availableTools)
227+
expect(selected).toContain("ask_followup_question")
228+
expect(selected).toContain("attempt_completion")
229+
})
230+
})
231+
})
232+
233+
describe("empty or invalid inputs", () => {
234+
it("should handle empty query gracefully", () => {
235+
const query = ""
236+
const modeConfig = createModeConfig(["read"])
237+
const availableTools = createAvailableTools([
238+
"read_file",
239+
"list_files",
240+
"ask_followup_question",
241+
"attempt_completion",
242+
])
243+
244+
const selected = analyzer.selectTools(query, modeConfig, availableTools)
245+
expect(selected.length).toBeGreaterThan(0)
246+
expect(selected).toContain("ask_followup_question")
247+
expect(selected).toContain("attempt_completion")
248+
})
249+
250+
it("should handle empty available tools", () => {
251+
const query = "do something"
252+
const modeConfig = createModeConfig(["read"])
253+
const availableTools = createAvailableTools([])
254+
255+
const selected = analyzer.selectTools(query, modeConfig, availableTools)
256+
expect(selected).toEqual([])
257+
})
258+
})
259+
260+
describe("tool scoring", () => {
261+
it("should score tools based on query relevance", () => {
262+
const query = "write a test file for the authentication module"
263+
const modeConfig = createModeConfig(["read", "edit"])
264+
const availableTools = createAvailableTools([
265+
"read_file",
266+
"write_to_file",
267+
"apply_diff",
268+
"list_files",
269+
"search_files",
270+
"execute_command",
271+
"ask_followup_question",
272+
"attempt_completion",
273+
])
274+
275+
const selected = analyzer.selectTools(query, modeConfig, availableTools)
276+
277+
// write_to_file should be highly ranked for "write a test file"
278+
const writeIndex = selected.indexOf("write_to_file")
279+
expect(writeIndex).toBeGreaterThanOrEqual(0)
280+
expect(writeIndex).toBeLessThan(4) // Should be in top 4 tools
281+
282+
// read_file should also be included for context
283+
expect(selected).toContain("read_file")
284+
})
285+
286+
it("should handle multiple tool mentions in query", () => {
287+
const query = "read the config file, update it, and run the build command"
288+
const modeConfig = createModeConfig(["read", "edit", "command"])
289+
const availableTools = createAvailableTools([
290+
"read_file",
291+
"write_to_file",
292+
"apply_diff",
293+
"execute_command",
294+
"list_files",
295+
"search_files",
296+
"ask_followup_question",
297+
"attempt_completion",
298+
])
299+
300+
const selected = analyzer.selectTools(query, modeConfig, availableTools)
301+
302+
// Should include all mentioned operations
303+
expect(selected).toContain("read_file")
304+
expect(selected).toContain("apply_diff") // for "update"
305+
expect(selected).toContain("execute_command") // for "run"
306+
})
307+
})
308+
})

0 commit comments

Comments
 (0)