vtcode_core/prompts/
cache_aware.rs1use crate::config::constants::tools;
2use crate::llm::provider::ToolDefinition;
3use compact_str::CompactString;
4
5const PRIORITY_TOOLS: &[&str] = &[
7 tools::EXEC_COMMAND,
8 tools::CODE_SEARCH,
9 tools::APPLY_PATCH,
10 tools::WRITE_STDIN,
11 tools::REQUEST_USER_INPUT,
12 tools::TASK_TRACKER,
13];
14
15pub fn sort_tool_definitions(mut tools: Vec<ToolDefinition>) -> Vec<ToolDefinition> {
19 struct SortKey {
20 priority: Option<usize>,
21 name: CompactString,
22 }
23
24 let mut keyed = tools
25 .drain(..)
26 .map(|tool| {
27 let name = tool.function.as_ref().map_or("", |function| function.name.as_str());
28 let priority = PRIORITY_TOOLS.iter().position(|&priority_name| priority_name == name);
29 (SortKey { priority, name: CompactString::from(name) }, tool)
30 })
31 .collect::<Vec<_>>();
32
33 keyed.sort_by(|(left_key, left_tool), (right_key, right_tool)| match (left_key.priority, right_key.priority) {
34 (Some(left_priority), Some(right_priority)) => left_priority.cmp(&right_priority),
35 (Some(_), None) => std::cmp::Ordering::Less,
36 (None, Some(_)) => std::cmp::Ordering::Greater,
37 (None, None) => left_key
38 .name
39 .cmp(&right_key.name)
40 .then_with(|| left_tool.tool_type.cmp(&right_tool.tool_type)),
41 });
42
43 keyed.into_iter().map(|(_, tool)| tool).collect()
44}
45
46#[cfg(test)]
47mod tests {
48 use hashbrown::HashSet;
49
50 use super::PRIORITY_TOOLS;
51 use super::sort_tool_definitions;
52 use crate::llm::provider::ToolDefinition;
53
54 #[test]
55 fn sort_tool_definitions_orders_by_name() {
56 let tools = vec![
57 ToolDefinition::function("b_tool".to_string(), "b".to_string(), serde_json::json!({})),
58 ToolDefinition::function("a_tool".to_string(), "a".to_string(), serde_json::json!({})),
59 ];
60
61 let sorted = sort_tool_definitions(tools);
62 let names: Vec<&str> = sorted
63 .iter()
64 .filter_map(|tool| tool.function.as_ref().map(|func| func.name.as_str()))
65 .collect();
66
67 assert_eq!(names, vec!["a_tool", "b_tool"]);
68 }
69
70 #[test]
71 fn sort_tool_definitions_prioritizes_current_core_tools() {
72 let tools = vec![
73 ToolDefinition::function("zebra_tool".to_string(), "z".to_string(), serde_json::json!({})),
74 ToolDefinition::function("code_search".to_string(), "search".to_string(), serde_json::json!({})),
75 ToolDefinition::function("request_user_input".to_string(), "ask".to_string(), serde_json::json!({})),
76 ToolDefinition::function("alpha_tool".to_string(), "a".to_string(), serde_json::json!({})),
77 ToolDefinition::function("exec_command".to_string(), "shell".to_string(), serde_json::json!({})),
78 ];
79
80 let sorted = sort_tool_definitions(tools);
81 let names: Vec<&str> = sorted
82 .iter()
83 .filter_map(|tool| tool.function.as_ref().map(|func| func.name.as_str()))
84 .collect();
85
86 assert_eq!(
88 names,
89 vec![
90 "exec_command",
91 "code_search",
92 "request_user_input",
93 "alpha_tool",
94 "zebra_tool"
95 ]
96 );
97 }
98
99 #[test]
100 fn priority_tools_are_unique() {
101 let unique: HashSet<&str> = PRIORITY_TOOLS.iter().copied().collect();
102 assert_eq!(unique.len(), PRIORITY_TOOLS.len());
103 }
104
105 #[test]
106 fn duplicate_priority_names_keep_input_order() {
107 let mut first = ToolDefinition::function("code_search".to_string(), "first".to_string(), serde_json::json!({}));
108 first.tool_type = "first_type".to_string();
109 let mut second =
110 ToolDefinition::function("code_search".to_string(), "second".to_string(), serde_json::json!({}));
111 second.tool_type = "second_type".to_string();
112
113 let sorted = sort_tool_definitions(vec![first, second]);
114 let descriptions: Vec<&str> = sorted
115 .iter()
116 .filter_map(|tool| tool.function.as_ref().map(|function| function.description.as_str()))
117 .collect();
118 assert_eq!(descriptions, vec!["first", "second"]);
119 }
120}