Skip to main content

vtcode_core/prompts/
cache_aware.rs

1use crate::config::constants::tools;
2use crate::llm::provider::ToolDefinition;
3use compact_str::CompactString;
4
5/// Priority tools that should appear first in model-visible definitions.
6const 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
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    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        // Priority tools first (in priority order), then alphabetical
87        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}