Skip to main content

apollo/tools/
tool_search.rs

1//! ToolSearchTool — keyword search over all available tools.
2//!
3//! When many tools are registered, this lets the agent find the right tool
4//! by searching names and descriptions without having to remember them all.
5
6use std::sync::Arc;
7
8use async_trait::async_trait;
9use serde::Deserialize;
10use tokio::sync::RwLock;
11
12use super::traits::*;
13
14pub struct ToolSearchTool {
15    tools: Arc<RwLock<Vec<Arc<dyn Tool>>>>,
16}
17
18impl ToolSearchTool {
19    pub fn new(tools: Arc<RwLock<Vec<Arc<dyn Tool>>>>) -> Self {
20        Self { tools }
21    }
22}
23
24#[derive(Deserialize)]
25struct SearchArgs {
26    /// Keywords to search for in tool names and descriptions
27    query: String,
28    /// Max results to return (default 10)
29    #[serde(default = "default_limit")]
30    limit: usize,
31}
32
33fn default_limit() -> usize {
34    10
35}
36
37#[async_trait]
38impl Tool for ToolSearchTool {
39    fn name(&self) -> &str {
40        "tool_search"
41    }
42
43    fn spec(&self) -> ToolSpec {
44        ToolSpec {
45            name: "tool_search".to_string(),
46            description: "Search for available tools by keyword. \
47                Returns matching tool names and descriptions. \
48                Use when you need to find a tool but aren't sure of its exact name."
49                .to_string(),
50            parameters: serde_json::json!({
51                "type": "object",
52                "properties": {
53                    "query": {
54                        "type": "string",
55                        "description": "Keywords to search (searches tool name and description)"
56                    },
57                    "limit": {
58                        "type": "integer",
59                        "description": "Max results to return (default 10)",
60                        "minimum": 1,
61                        "maximum": 50
62                    }
63                },
64                "required": ["query"]
65            }),
66        }
67    }
68
69    async fn execute(&self, arguments: &str) -> anyhow::Result<ToolResult> {
70        let args: SearchArgs = serde_json::from_str(arguments)?;
71
72        if args.query.trim().is_empty() {
73            return Ok(ToolResult::error("query must not be empty"));
74        }
75
76        let terms: Vec<String> = args
77            .query
78            .to_lowercase()
79            .split_whitespace()
80            .map(String::from)
81            .collect();
82
83        let tools = self.tools.read().await;
84        let limit = args.limit.min(50);
85
86        // Score each tool: +2 for name match, +1 for description match per term
87        let mut scored: Vec<(usize, &Arc<dyn Tool>)> = tools
88            .iter()
89            .filter_map(|t| {
90                let spec = t.spec();
91                let name_lower = spec.name.to_lowercase();
92                let desc_lower = spec.description.to_lowercase();
93                let score: usize = terms
94                    .iter()
95                    .map(|term| {
96                        let in_name = if name_lower.contains(term.as_str()) {
97                            2
98                        } else {
99                            0
100                        };
101                        let in_desc = if desc_lower.contains(term.as_str()) {
102                            1
103                        } else {
104                            0
105                        };
106                        in_name + in_desc
107                    })
108                    .sum();
109                if score > 0 {
110                    Some((score, t))
111                } else {
112                    None
113                }
114            })
115            .collect();
116
117        scored.sort_by_key(|item| std::cmp::Reverse(item.0));
118        scored.truncate(limit);
119
120        if scored.is_empty() {
121            return Ok(ToolResult::success(format!(
122                "No tools found matching '{}'.",
123                args.query
124            )));
125        }
126
127        let lines: Vec<String> = scored
128            .iter()
129            .map(|(score, t)| {
130                let spec = t.spec();
131                let desc_preview: String = spec.description.chars().take(100).collect();
132                format!("- **{}** (score:{}) — {}", spec.name, score, desc_preview)
133            })
134            .collect();
135
136        Ok(ToolResult::success(format!(
137            "Found {} tool(s) matching '{}':\n\n{}",
138            lines.len(),
139            args.query,
140            lines.join("\n")
141        )))
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148    use crate::policy::ExecutionPolicy;
149    use crate::tools::shell::ShellTool;
150    use std::path::PathBuf;
151
152    #[tokio::test]
153    async fn finds_shell_tool() {
154        let shell = Arc::new(ShellTool::new(
155            PathBuf::from("."),
156            Arc::new(ExecutionPolicy::default()),
157        ));
158        let tools: Arc<RwLock<Vec<Arc<dyn Tool>>>> = Arc::new(RwLock::new(vec![shell]));
159        let tool = ToolSearchTool::new(tools);
160
161        let r = tool
162            .execute(r#"{"query":"shell execute command"}"#)
163            .await
164            .unwrap();
165        assert!(!r.is_error);
166        assert!(r.output.contains("exec") || r.output.contains("shell"));
167    }
168
169    #[tokio::test]
170    async fn empty_query_errors() {
171        let tools: Arc<RwLock<Vec<Arc<dyn Tool>>>> = Arc::new(RwLock::new(vec![]));
172        let tool = ToolSearchTool::new(tools);
173        let r = tool.execute(r#"{"query":""}"#).await.unwrap();
174        assert!(r.is_error);
175    }
176
177    #[tokio::test]
178    async fn no_match_returns_message() {
179        let tools: Arc<RwLock<Vec<Arc<dyn Tool>>>> = Arc::new(RwLock::new(vec![]));
180        let tool = ToolSearchTool::new(tools);
181        let r = tool.execute(r#"{"query":"xyzzy"}"#).await.unwrap();
182        assert!(!r.is_error);
183        assert!(r.output.contains("No tools found"));
184    }
185}