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