spec_ai/spec_ai_core/tools/
mod.rs

1pub mod builtin;
2pub mod mcp;
3pub mod plugin_adapter;
4
5use anyhow::Result;
6use async_trait::async_trait;
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::collections::HashMap;
10use std::sync::{Arc, Mutex};
11use tracing::debug;
12
13use self::builtin::{
14    ActivateSkillTool, AudioTranscriptionTool, BashTool, CodeSearchTool, EchoTool, FileExtractTool,
15    FileReadTool, FileWriteTool, GenerateCodeTool, GraphTool, GrepTool, MathTool, PromptUserTool,
16    RgTool, SearchTool, ShellTool,
17};
18
19#[cfg(feature = "api")]
20use self::builtin::WebSearchTool;
21
22#[cfg(feature = "web-scraping")]
23use self::builtin::WebScraperTool;
24use crate::spec_ai_core::agent::model::ModelProvider;
25use crate::spec_ai_core::agent::safety::RunSafetyBudget;
26use crate::spec_ai_core::embeddings::EmbeddingsClient;
27use crate::spec_ai_core::persistence::Persistence;
28use crate::spec_ai_config::config::McpConfig;
29use std::path::PathBuf;
30
31pub use plugin_adapter::PluginToolAdapter;
32
33#[cfg(feature = "openai")]
34use async_openai::types::chat::ChatCompletionTool;
35
36/// Result of tool execution
37#[derive(Debug, Clone, Serialize, Deserialize)]
38pub struct ToolResult {
39    /// Whether execution succeeded
40    pub success: bool,
41    /// Output from the tool
42    pub output: String,
43    /// Error message if execution failed
44    pub error: Option<String>,
45}
46
47/// Per-invocation context passed to tools that need run-level guardrails.
48#[derive(Clone, Default)]
49pub struct ToolExecutionContext {
50    pub safety: Option<RunSafetyBudget>,
51    pub delegation_depth: usize,
52}
53
54impl ToolResult {
55    /// Create a successful result
56    pub fn success(output: impl Into<String>) -> Self {
57        Self {
58            success: true,
59            output: output.into(),
60            error: None,
61        }
62    }
63
64    /// Create a failure result
65    pub fn failure(error: impl Into<String>) -> Self {
66        Self {
67            success: false,
68            output: String::new(),
69            error: Some(error.into()),
70        }
71    }
72}
73
74/// Trait for all tools that can be executed by the agent
75#[async_trait]
76pub trait Tool: Send + Sync {
77    /// Unique name of the tool
78    fn name(&self) -> &str;
79
80    /// Human-readable description of what the tool does
81    fn description(&self) -> &str;
82
83    /// JSON Schema describing the tool's parameters
84    fn parameters(&self) -> Value;
85
86    /// Execute the tool with the given arguments
87    async fn execute(&self, args: Value) -> Result<ToolResult>;
88
89    /// Execute with run context. Tools that do not need context use `execute`.
90    async fn execute_with_context(
91        &self,
92        args: Value,
93        _context: ToolExecutionContext,
94    ) -> Result<ToolResult> {
95        self.execute(args).await
96    }
97}
98
99/// Registry for managing and executing tools
100pub struct ToolRegistry {
101    tools: Mutex<HashMap<String, Arc<dyn Tool>>>,
102}
103
104impl ToolRegistry {
105    /// Create a new empty tool registry
106    pub fn new() -> Self {
107        Self {
108            tools: Mutex::new(HashMap::new()),
109        }
110    }
111
112    /// Create a registry populated with all built-in tools.
113    ///
114    /// Tools that require persistence (e.g., `graph`) are only registered when
115    /// an [`Arc<Persistence>`] is provided.
116    #[allow(unused_variables)]
117    pub fn with_builtin_tools(
118        persistence: Option<Arc<Persistence>>,
119        embeddings: Option<EmbeddingsClient>,
120        code_model_provider: Option<Arc<dyn ModelProvider>>,
121        skills_dirs: Vec<PathBuf>,
122    ) -> Self {
123        let registry = Self::new();
124
125        // Register all built-in tools
126        registry.register(Arc::new(EchoTool::new()));
127        registry.register(Arc::new(MathTool::new()));
128        registry.register(Arc::new(FileReadTool::new()));
129        registry.register(Arc::new(FileExtractTool::new()));
130        registry.register(Arc::new(FileWriteTool::new()));
131        registry.register(Arc::new(PromptUserTool::new()));
132        registry.register(Arc::new(SearchTool::new()));
133        registry.register(Arc::new(GrepTool::new()));
134        registry.register(Arc::new(RgTool::new()));
135        registry.register(Arc::new(CodeSearchTool::new()));
136        registry.register(Arc::new(BashTool::new()));
137        registry.register(Arc::new(ShellTool::new()));
138        if !skills_dirs.is_empty() {
139            registry.register(Arc::new(ActivateSkillTool::new(skills_dirs)));
140        }
141        if let Some(provider) = code_model_provider {
142            registry.register(Arc::new(GenerateCodeTool::new(provider)));
143        }
144
145        // Register web search if api feature is enabled
146        #[cfg(feature = "api")]
147        registry.register(Arc::new(WebSearchTool::new().with_embeddings(embeddings)));
148
149        // Register web scraper if feature is enabled
150        #[cfg(feature = "web-scraping")]
151        registry.register(Arc::new(WebScraperTool::new()));
152
153        if let Some(persistence) = persistence {
154            registry.register(Arc::new(GraphTool::new(persistence.clone())));
155            registry.register(Arc::new(AudioTranscriptionTool::with_persistence(
156                persistence,
157            )));
158        } else {
159            registry.register(Arc::new(AudioTranscriptionTool::new()));
160        }
161
162        tracing::debug!("ToolRegistry created with {} tools", registry.len());
163        for name in registry.list() {
164            tracing::debug!("  - Tool: {}", name);
165        }
166
167        registry
168    }
169
170    /// Register a tool in the registry
171    pub fn register(&self, tool: Arc<dyn Tool>) {
172        let name = tool.name().to_string();
173        self.tools.lock().unwrap().insert(name, tool);
174    }
175
176    /// Get a tool by name
177    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
178        self.tools.lock().unwrap().get(name).cloned()
179    }
180
181    /// List all registered tool names
182    pub fn list(&self) -> Vec<String> {
183        self.tools
184            .lock()
185            .unwrap()
186            .keys()
187            .map(|s| s.to_string())
188            .collect()
189    }
190
191    /// Check if a tool is registered
192    pub fn has(&self, name: &str) -> bool {
193        self.tools.lock().unwrap().contains_key(name)
194    }
195
196    /// Execute a tool by name with the given arguments
197    pub async fn execute(&self, name: &str, args: Value) -> Result<ToolResult> {
198        self.execute_with_context(name, args, ToolExecutionContext::default())
199            .await
200    }
201
202    /// Execute a tool by name with run context.
203    pub async fn execute_with_context(
204        &self,
205        name: &str,
206        args: Value,
207        context: ToolExecutionContext,
208    ) -> Result<ToolResult> {
209        let tool = self
210            .get(name)
211            .ok_or_else(|| anyhow::anyhow!("Tool not found: {}", name))?;
212
213        debug!("Executing tool '{}'", name);
214        let result = tool.execute_with_context(args, context).await;
215        match &result {
216            Ok(res) => {
217                debug!(
218                    "Tool '{}' completed: success={}, error={:?}",
219                    name, res.success, res.error
220                );
221            }
222            Err(err) => {
223                debug!("Tool '{}' failed to execute: {}", name, err);
224            }
225        }
226        result
227    }
228
229    /// Get the number of registered tools
230    pub fn len(&self) -> usize {
231        self.tools.lock().unwrap().len()
232    }
233
234    /// Check if the registry is empty
235    pub fn is_empty(&self) -> bool {
236        self.tools.lock().unwrap().is_empty()
237    }
238
239    /// Load plugins from a directory and register their tools
240    ///
241    /// # Arguments
242    /// * `dir` - Directory containing plugin libraries
243    /// * `allow_override` - Whether plugins can override built-in tools
244    ///
245    /// # Returns
246    /// Statistics about the loading process
247    pub fn load_plugins(
248        &mut self,
249        dir: &std::path::Path,
250        allow_override: bool,
251    ) -> anyhow::Result<crate::spec_ai_plugin::LoadStats> {
252        use crate::spec_ai_plugin::{PluginLoader, expand_tilde};
253
254        let expanded_dir = expand_tilde(dir);
255
256        let mut loader = PluginLoader::new();
257        let stats = loader.load_directory(&expanded_dir)?;
258
259        // Register tools from plugins
260        for (tool_ref, plugin_name) in loader.all_tools() {
261            let adapter = match PluginToolAdapter::new(tool_ref, plugin_name) {
262                Ok(a) => a,
263                Err(e) => {
264                    tracing::warn!(
265                        "Failed to create adapter for tool from {}: {}",
266                        plugin_name,
267                        e
268                    );
269                    continue;
270                }
271            };
272
273            let tool_name = adapter.name().to_string();
274
275            // Check for conflicts with built-in tools
276            if self.has(&tool_name) {
277                if allow_override {
278                    tracing::info!(
279                        "Plugin tool '{}' from '{}' overriding built-in tool",
280                        tool_name,
281                        plugin_name
282                    );
283                } else {
284                    tracing::warn!(
285                        "Plugin tool '{}' from '{}' would override built-in, skipping (set allow_override_builtin=true to allow)",
286                        tool_name,
287                        plugin_name
288                    );
289                    continue;
290                }
291            }
292
293            tracing::debug!(
294                "Registering plugin tool '{}' from '{}'",
295                tool_name,
296                plugin_name
297            );
298            self.register(Arc::new(adapter));
299        }
300
301        Ok(stats)
302    }
303
304    /// Load MCP servers and register their tools
305    pub async fn load_mcp_servers(&self, config: &McpConfig) -> anyhow::Result<()> {
306        if !config.enabled {
307            return Ok(());
308        }
309
310        let mut manager = crate::spec_ai_core::tools::mcp::McpManager::new();
311        for (name, server_config) in &config.servers {
312            if let Err(e) = manager
313                .connect_stdio(
314                    &server_config.command,
315                    &server_config.args,
316                    &server_config.env,
317                )
318                .await
319            {
320                tracing::error!("Failed to connect to MCP server '{}': {}", name, e);
321            }
322        }
323
324        for adapter in manager.list_tools().await {
325            tracing::info!("Registering MCP tool: {}", adapter.name());
326            self.register(Arc::new(adapter));
327        }
328
329        Ok(())
330    }
331
332    /// Convert all tools in the registry to OpenAI ChatCompletionTool format.
333    ///
334    /// Used by providers that support native function calling (OpenAI-compatible,
335    /// including MLX and LM Studio when enabled).
336    #[cfg(any(feature = "openai", feature = "mlx", feature = "lmstudio"))]
337    pub fn to_openai_tools(&self) -> Vec<ChatCompletionTool> {
338        use crate::spec_ai_core::agent::function_calling::tool_to_openai_function;
339
340        let tools = self.tools.lock().unwrap();
341        tools
342            .values()
343            .map(|tool| {
344                tool_to_openai_function(tool.name(), tool.description(), &tool.parameters())
345            })
346            .collect()
347    }
348}
349
350impl Default for ToolRegistry {
351    fn default() -> Self {
352        Self::new()
353    }
354}
355
356#[cfg(test)]
357mod tests {
358    use super::*;
359
360    struct DummyTool;
361
362    #[async_trait]
363    impl Tool for DummyTool {
364        fn name(&self) -> &str {
365            "dummy"
366        }
367
368        fn description(&self) -> &str {
369            "A dummy tool for testing"
370        }
371
372        fn parameters(&self) -> Value {
373            serde_json::json!({
374                "type": "object",
375                "properties": {}
376            })
377        }
378
379        async fn execute(&self, _args: Value) -> Result<ToolResult> {
380            Ok(ToolResult::success("dummy output"))
381        }
382    }
383
384    #[tokio::test]
385    async fn test_register_and_get_tool() {
386        let registry = ToolRegistry::new();
387        let tool = Arc::new(DummyTool);
388
389        registry.register(tool.clone());
390
391        assert!(registry.has("dummy"));
392        assert!(registry.get("dummy").is_some());
393        assert_eq!(registry.len(), 1);
394    }
395
396    #[tokio::test]
397    async fn test_list_tools() {
398        let registry = ToolRegistry::new();
399        registry.register(Arc::new(DummyTool));
400
401        let tools = registry.list();
402        assert_eq!(tools.len(), 1);
403        assert!(tools.contains(&"dummy".to_string()));
404    }
405
406    #[tokio::test]
407    async fn test_execute_tool() {
408        let registry = ToolRegistry::new();
409        registry.register(Arc::new(DummyTool));
410        let result = registry.execute("dummy", Value::Null).await.unwrap();
411        assert!(result.success);
412        assert_eq!(result.output, "dummy output");
413    }
414
415    #[tokio::test]
416    async fn test_execute_nonexistent_tool() {
417        let registry = ToolRegistry::new();
418        let result = registry.execute("nonexistent", Value::Null).await;
419        assert!(result.is_err());
420    }
421
422    #[tokio::test]
423    async fn test_tool_result_success() {
424        let result = ToolResult::success("test output");
425        assert!(result.success);
426        assert_eq!(result.output, "test output");
427        assert!(result.error.is_none());
428    }
429
430    #[tokio::test]
431    async fn test_tool_result_failure() {
432        let result = ToolResult::failure("test error");
433        assert!(!result.success);
434        assert_eq!(result.error, Some("test error".to_string()));
435    }
436}
spec_ai/spec_ai_core/tools/mod.rs

spec_ai/spec_ai_core/tools/
mod.rs
