Skip to main content

vtcode_core/prompts/
cache_aware.rs

1use crate::config::constants::tools;
2use crate::llm::provider::ToolDefinition;
3
4/// Priority tools (unified tools + critical controls) that should appear first
5const PRIORITY_TOOLS: &[&str] = &[
6    tools::UNIFIED_SEARCH,
7    tools::UNIFIED_FILE,
8    tools::UNIFIED_EXEC,
9    tools::REQUEST_USER_INPUT,
10    tools::PLAN_TASK_TRACKER,
11    tools::TASK_TRACKER,
12    tools::EXIT_PLAN_MODE,
13];
14
15/// Sort tool definitions with priority grouping to improve LLM attention.
16/// Priority tools appear first (better attention), then alphabetically.
17/// This can save ~50-100 tokens per request via improved LLM focus.
18pub fn sort_tool_definitions(mut tools: Vec<ToolDefinition>) -> Vec<ToolDefinition> {
19    tools.sort_by(|a, b| {
20        let a_name = a
21            .function
22            .as_ref()
23            .map(|func| func.name.as_str())
24            .unwrap_or("");
25        let b_name = b
26            .function
27            .as_ref()
28            .map(|func| func.name.as_str())
29            .unwrap_or("");
30
31        // Find priority positions (None = not priority = end of list)
32        let a_priority = PRIORITY_TOOLS.iter().position(|&p| p == a_name);
33        let b_priority = PRIORITY_TOOLS.iter().position(|&p| p == b_name);
34
35        match (a_priority, b_priority) {
36            // Both priority - sort by priority order
37            (Some(a_pos), Some(b_pos)) => a_pos.cmp(&b_pos),
38            // Only a is priority - a comes first
39            (Some(_), None) => std::cmp::Ordering::Less,
40            // Only b is priority - b comes first
41            (None, Some(_)) => std::cmp::Ordering::Greater,
42            // Neither priority - sort alphabetically
43            (None, None) => {
44                let name_cmp = a_name.cmp(b_name);
45                if name_cmp != std::cmp::Ordering::Equal {
46                    return name_cmp;
47                }
48                a.tool_type.cmp(&b.tool_type)
49            }
50        }
51    });
52    tools
53}
54
55#[cfg(test)]
56mod tests {
57    use hashbrown::HashSet;
58
59    use super::PRIORITY_TOOLS;
60    use super::sort_tool_definitions;
61    use crate::llm::provider::ToolDefinition;
62
63    #[test]
64    fn sort_tool_definitions_orders_by_name() {
65        let tools = vec![
66            ToolDefinition::function("b_tool".to_string(), "b".to_string(), serde_json::json!({})),
67            ToolDefinition::function("a_tool".to_string(), "a".to_string(), serde_json::json!({})),
68        ];
69
70        let sorted = sort_tool_definitions(tools);
71        let names: Vec<&str> = sorted
72            .iter()
73            .filter_map(|tool| tool.function.as_ref().map(|func| func.name.as_str()))
74            .collect();
75
76        assert_eq!(names, vec!["a_tool", "b_tool"]);
77    }
78
79    #[test]
80    fn sort_tool_definitions_prioritizes_unified_tools() {
81        let tools = vec![
82            ToolDefinition::function(
83                "zebra_tool".to_string(),
84                "z".to_string(),
85                serde_json::json!({}),
86            ),
87            ToolDefinition::function(
88                "unified_search".to_string(),
89                "search".to_string(),
90                serde_json::json!({}),
91            ),
92            ToolDefinition::function(
93                "request_user_input".to_string(),
94                "ask".to_string(),
95                serde_json::json!({}),
96            ),
97            ToolDefinition::function(
98                "alpha_tool".to_string(),
99                "a".to_string(),
100                serde_json::json!({}),
101            ),
102            ToolDefinition::function(
103                "unified_file".to_string(),
104                "file".to_string(),
105                serde_json::json!({}),
106            ),
107        ];
108
109        let sorted = sort_tool_definitions(tools);
110        let names: Vec<&str> = sorted
111            .iter()
112            .filter_map(|tool| tool.function.as_ref().map(|func| func.name.as_str()))
113            .collect();
114
115        // Priority tools first (in priority order), then alphabetical
116        assert_eq!(
117            names,
118            vec![
119                "unified_search",
120                "unified_file",
121                "request_user_input",
122                "alpha_tool",
123                "zebra_tool"
124            ]
125        );
126    }
127
128    #[test]
129    fn priority_tools_are_unique() {
130        let unique: HashSet<&str> = PRIORITY_TOOLS.iter().copied().collect();
131        assert_eq!(unique.len(), PRIORITY_TOOLS.len());
132    }
133}