claude_agent/tools/search/
index.rs

1//! Tool index for efficient searching.
2
3use crate::mcp::McpToolDefinition;
4
5#[derive(Debug, Clone)]
6pub struct ToolIndexEntry {
7    pub qualified_name: String,
8    pub server_name: String,
9    pub tool_name: String,
10    pub description: String,
11    pub arg_names: Vec<String>,
12    pub arg_descriptions: Vec<String>,
13    pub estimated_tokens: usize,
14}
15
16impl ToolIndexEntry {
17    pub fn from_mcp_tool(server: &str, tool: &McpToolDefinition) -> Self {
18        let (arg_names, arg_descriptions) = Self::extract_arg_info(&tool.input_schema);
19        let estimated_tokens = Self::estimate_tokens(tool);
20
21        Self {
22            qualified_name: crate::mcp::make_mcp_name(server, &tool.name),
23            server_name: server.to_string(),
24            tool_name: tool.name.clone(),
25            description: tool.description.clone(),
26            arg_names,
27            arg_descriptions,
28            estimated_tokens,
29        }
30    }
31
32    fn extract_arg_info(schema: &serde_json::Value) -> (Vec<String>, Vec<String>) {
33        let mut names = Vec::new();
34        let mut descs = Vec::new();
35
36        if let Some(props) = schema.get("properties").and_then(|p| p.as_object()) {
37            for (name, prop) in props {
38                names.push(name.clone());
39                if let Some(desc) = prop.get("description").and_then(|d| d.as_str()) {
40                    descs.push(desc.to_string());
41                }
42            }
43        }
44
45        (names, descs)
46    }
47
48    fn estimate_tokens(tool: &McpToolDefinition) -> usize {
49        let name_tokens = tool.name.len() / 4;
50        let desc_tokens = tool.description.len() / 4;
51        let schema_tokens = tool.input_schema.to_string().len() / 4;
52        name_tokens + desc_tokens + schema_tokens + 20
53    }
54
55    pub fn searchable_text(&self) -> String {
56        format!(
57            "{} {} {} {}",
58            self.tool_name,
59            self.description,
60            self.arg_names.join(" "),
61            self.arg_descriptions.join(" ")
62        )
63    }
64}
65
66#[derive(Debug, Default)]
67pub struct ToolIndex {
68    entries: Vec<ToolIndexEntry>,
69    total_tokens: usize,
70}
71
72impl ToolIndex {
73    pub fn new() -> Self {
74        Self::default()
75    }
76
77    pub fn add(&mut self, entry: ToolIndexEntry) {
78        self.total_tokens += entry.estimated_tokens;
79        self.entries.push(entry);
80    }
81
82    pub fn total_tokens(&self) -> usize {
83        self.total_tokens
84    }
85
86    pub fn len(&self) -> usize {
87        self.entries.len()
88    }
89
90    pub fn is_empty(&self) -> bool {
91        self.entries.is_empty()
92    }
93
94    pub fn entries(&self) -> &[ToolIndexEntry] {
95        &self.entries
96    }
97
98    pub fn get(&self, qualified_name: &str) -> Option<&ToolIndexEntry> {
99        self.entries
100            .iter()
101            .find(|e| e.qualified_name == qualified_name)
102    }
103
104    pub fn clear(&mut self) {
105        self.entries.clear();
106        self.total_tokens = 0;
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    fn make_test_tool(name: &str, desc: &str) -> McpToolDefinition {
115        McpToolDefinition {
116            name: name.to_string(),
117            description: desc.to_string(),
118            input_schema: serde_json::json!({
119                "type": "object",
120                "properties": {
121                    "arg1": { "type": "string", "description": "First argument" }
122                }
123            }),
124        }
125    }
126
127    #[test]
128    fn test_index_entry_creation() {
129        let tool = make_test_tool("read_file", "Read a file from disk");
130        let entry = ToolIndexEntry::from_mcp_tool("filesystem", &tool);
131
132        assert_eq!(entry.qualified_name, "mcp__filesystem_read_file");
133        assert_eq!(entry.server_name, "filesystem");
134        assert_eq!(entry.tool_name, "read_file");
135        assert!(entry.estimated_tokens > 0);
136    }
137
138    #[test]
139    fn test_searchable_text() {
140        let tool = make_test_tool("get_weather", "Get weather for location");
141        let entry = ToolIndexEntry::from_mcp_tool("weather", &tool);
142        let text = entry.searchable_text();
143
144        assert!(text.contains("get_weather"));
145        assert!(text.contains("weather"));
146        assert!(text.contains("location"));
147    }
148
149    #[test]
150    fn test_index_operations() {
151        let mut index = ToolIndex::new();
152        assert!(index.is_empty());
153
154        let tool = make_test_tool("test", "Test tool");
155        let entry = ToolIndexEntry::from_mcp_tool("server", &tool);
156        let tokens = entry.estimated_tokens;
157
158        index.add(entry);
159        assert_eq!(index.len(), 1);
160        assert_eq!(index.total_tokens(), tokens);
161        assert!(index.get("mcp__server_test").is_some());
162    }
163}
claude_agent/tools/search/index.rs

claude_agent/tools/search/
index.rs
