magi_tool/
tool_kit.rs

1use dashmap::{DashMap, mapref::entry::Entry};
2use serde_json::Value;
3use std::{future::Future, pin::Pin, sync::Arc};
4use thiserror::Error;
5
6use crate::Tool as ToolDefinition;
7
8type ToolCallFuture = Pin<Box<dyn Future<Output = Result<String, String>> + Send>>;
9type ToolHandler = Arc<dyn Fn(Value) -> ToolCallFuture + Send + Sync>;
10
11#[derive(Clone)]
12struct Tool {
13    definition: ToolDefinition,
14    handler: ToolHandler,
15}
16
17/// Error type for local tool registration and execution.
18#[derive(Debug, Error, Clone, PartialEq, Eq)]
19pub enum ToolCallError {
20    #[error("Tool '{0}' already exists")]
21    ToolAlreadyExists(String),
22    #[error("Tool '{0}' not found")]
23    ToolNotFound(String),
24    #[error("Tool '{tool_name}' call failed: {message}")]
25    CallFailed { tool_name: String, message: String },
26}
27
28/// Local tool registry backed by tool-name-to-closure mappings.
29#[derive(Clone)]
30pub struct ToolKit {
31    tools: Arc<DashMap<String, Tool>>,
32}
33
34impl ToolKit {
35    /// Creates a new empty ToolKit.
36    pub fn new() -> Self {
37        Self {
38            tools: Arc::new(DashMap::new()),
39        }
40    }
41
42    /// Creates a new ToolKit containing one registered local tool.
43    pub fn with_tool<F, Fut>(tool: ToolDefinition, handler: F) -> Result<Self, ToolCallError>
44    where
45        F: Fn(Value) -> Fut + Send + Sync + 'static,
46        Fut: Future<Output = Result<String, String>> + Send + 'static,
47    {
48        let tool_kit = Self::new();
49        tool_kit.register_tool(tool, handler)?;
50        Ok(tool_kit)
51    }
52
53    /// Registers a local tool and the closure that executes it.
54    pub fn register_tool<F, Fut>(
55        &self,
56        tool: ToolDefinition,
57        handler: F,
58    ) -> Result<(), ToolCallError>
59    where
60        F: Fn(Value) -> Fut + Send + Sync + 'static,
61        Fut: Future<Output = Result<String, String>> + Send + 'static,
62    {
63        let tool_name = tool.name.to_string();
64        let handler = Arc::new(move |arguments| Box::pin(handler(arguments)) as ToolCallFuture);
65
66        match self.tools.entry(tool_name) {
67            Entry::Occupied(entry) => Err(ToolCallError::ToolAlreadyExists(entry.key().clone())),
68            Entry::Vacant(entry) => {
69                entry.insert(Tool {
70                    definition: tool,
71                    handler,
72                });
73                Ok(())
74            }
75        }
76    }
77
78    /// Gets a tool by its local name.
79    pub fn get_tool(&self, tool_name: &str) -> Option<ToolDefinition> {
80        self.tools
81            .get(tool_name)
82            .map(|tool| tool.value().definition.clone())
83    }
84
85    /// Looks up multiple tools by local name, preserving missing entries as `None`.
86    pub fn lookup_tools(&self, tool_names: &[&str]) -> Vec<Option<ToolDefinition>> {
87        tool_names
88            .iter()
89            .map(|tool_name| self.get_tool(tool_name))
90            .collect()
91    }
92
93    /// Gets all registered local tool names.
94    pub fn tool_names(&self) -> Vec<String> {
95        self.tools.iter().map(|entry| entry.key().clone()).collect()
96    }
97
98    /// Gets all registered local tool definitions.
99    pub fn tools(&self) -> Vec<ToolDefinition> {
100        self.tools
101            .iter()
102            .map(|entry| entry.value().definition.clone())
103            .collect()
104    }
105
106    /// Removes a local tool by name.
107    pub fn unregister_tool(&self, tool_name: &str) -> Result<(), ToolCallError> {
108        if self.tools.remove(tool_name).is_none() {
109            return Err(ToolCallError::ToolNotFound(tool_name.to_string()));
110        }
111
112        Ok(())
113    }
114
115    /// Returns the number of registered local tools.
116    pub fn tool_count(&self) -> usize {
117        self.tools.len()
118    }
119
120    /// Calls a local tool by name and returns the text result for the LLM.
121    pub async fn call_tool(
122        &self,
123        tool_name: &str,
124        arguments: Value,
125    ) -> Result<String, ToolCallError> {
126        let handler = {
127            self.tools
128                .get(tool_name)
129                .map(|tool| Arc::clone(&tool.value().handler))
130                .ok_or_else(|| ToolCallError::ToolNotFound(tool_name.to_string()))?
131        };
132
133        handler(arguments)
134            .await
135            .map_err(|message| ToolCallError::CallFailed {
136                tool_name: tool_name.to_string(),
137                message,
138            })
139    }
140}
141
142impl Default for ToolKit {
143    fn default() -> Self {
144        Self::new()
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151    use schemars::JsonSchema;
152    use serde_json::json;
153
154    #[derive(JsonSchema)]
155    #[allow(dead_code)]
156    struct EchoParams {
157        value: String,
158    }
159
160    fn echo_tool() -> ToolDefinition {
161        crate::get_tool::<EchoParams, _, _>("echo", Some("Echo a value")).unwrap()
162    }
163
164    #[test]
165    fn test_with_tool_registers_one_tool() {
166        let tool_kit =
167            ToolKit::with_tool(echo_tool(), |_arguments| async { Ok(String::new()) }).unwrap();
168
169        assert_eq!(tool_kit.tool_count(), 1);
170        assert!(tool_kit.get_tool("echo").is_some());
171    }
172
173    #[tokio::test]
174    async fn test_register_and_call_tool() {
175        let tool_kit = ToolKit::new();
176        tool_kit
177            .register_tool(echo_tool(), |arguments| async move {
178                let value = arguments
179                    .get("value")
180                    .and_then(Value::as_str)
181                    .unwrap_or_default();
182                Ok(value.to_string())
183            })
184            .unwrap();
185
186        assert_eq!(tool_kit.tool_count(), 1);
187        assert!(tool_kit.get_tool("echo").is_some());
188
189        let result = tool_kit.call_tool("echo", json!({"value": "hello"})).await;
190        assert_eq!(result.unwrap(), "hello");
191    }
192
193    #[test]
194    fn test_duplicate_tool_fails() {
195        let tool_kit = ToolKit::new();
196        tool_kit
197            .register_tool(echo_tool(), |_arguments| async { Ok(String::new()) })
198            .unwrap();
199
200        let error = tool_kit
201            .register_tool(echo_tool(), |_arguments| async { Ok(String::new()) })
202            .unwrap_err();
203
204        assert_eq!(error, ToolCallError::ToolAlreadyExists("echo".to_string()));
205    }
206
207    #[tokio::test]
208    async fn test_missing_tool_call_fails() {
209        let tool_kit = ToolKit::new();
210        let error = tool_kit
211            .call_tool("missing", Value::Null)
212            .await
213            .unwrap_err();
214
215        assert_eq!(error, ToolCallError::ToolNotFound("missing".to_string()));
216    }
217
218    #[tokio::test]
219    async fn test_handler_error_includes_called_tool_name() {
220        let tool_kit = ToolKit::new();
221        tool_kit
222            .register_tool(echo_tool(), |_arguments| async {
223                Err("missing value".to_string())
224            })
225            .unwrap();
226
227        let error = tool_kit.call_tool("echo", Value::Null).await.unwrap_err();
228
229        assert_eq!(
230            error,
231            ToolCallError::CallFailed {
232                tool_name: "echo".to_string(),
233                message: "missing value".to_string(),
234            }
235        );
236    }
237
238    #[test]
239    fn test_lookup_tools_preserves_missing_entries() {
240        let tool_kit = ToolKit::new();
241        tool_kit
242            .register_tool(echo_tool(), |_arguments| async { Ok(String::new()) })
243            .unwrap();
244
245        let tools = tool_kit.lookup_tools(&["echo", "missing"]);
246
247        assert_eq!(tools.len(), 2);
248        assert!(tools[0].is_some());
249        assert!(tools[1].is_none());
250    }
251
252    #[test]
253    fn test_unregister_tool() {
254        let tool_kit = ToolKit::new();
255        tool_kit
256            .register_tool(echo_tool(), |_arguments| async { Ok(String::new()) })
257            .unwrap();
258
259        assert_eq!(tool_kit.tool_names(), vec!["echo".to_string()]);
260        assert_eq!(tool_kit.tools().len(), 1);
261
262        tool_kit.unregister_tool("echo").unwrap();
263
264        assert_eq!(tool_kit.tool_count(), 0);
265        assert!(tool_kit.get_tool("echo").is_none());
266    }
267}
magi_tool/tool_kit.rs

magi_tool/
tool_kit.rs
