Skip to main content

cascade_agent/tools/
mod.rs

1//! Tool system: trait definition, registry, and built-in tools.
2
3pub mod builtin;
4pub mod knowledge_tool;
5pub mod planning_tools;
6pub mod search;
7
8use async_trait::async_trait;
9use serde::{Deserialize, Serialize};
10use std::collections::HashMap;
11use std::sync::Arc;
12
13use crate::error::{AgentError, Result};
14
15// ---------------------------------------------------------------------------
16// ToolStatus / ToolResult
17// ---------------------------------------------------------------------------
18
19#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
20#[serde(rename_all = "lowercase")]
21pub enum ToolStatus {
22    Success,
23    Error,
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct ToolResult {
28    pub status: ToolStatus,
29    pub data: serde_json::Value,
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub error: Option<String>,
32}
33
34impl ToolResult {
35    /// Convenience constructor for a successful result with arbitrary data.
36    pub fn ok(data: serde_json::Value) -> Self {
37        Self {
38            status: ToolStatus::Success,
39            data,
40            error: None,
41        }
42    }
43
44    /// Convenience constructor for a successful result whose data is a string.
45    pub fn ok_string(s: impl Into<String>) -> Self {
46        Self::ok(serde_json::Value::String(s.into()))
47    }
48
49    /// Convenience constructor for an error result.
50    pub fn err(msg: impl Into<String>) -> Self {
51        Self {
52            status: ToolStatus::Error,
53            data: serde_json::Value::Null,
54            error: Some(msg.into()),
55        }
56    }
57
58    /// Serialize to a compact JSON string suitable for feeding back to the LLM.
59    pub fn to_json_string(&self) -> String {
60        serde_json::to_string(self).unwrap_or_else(|_| format!("{:?}", self))
61    }
62
63    /// Return the status as a lowercase string ("success" or "error").
64    pub fn status_str(&self) -> &str {
65        match self.status {
66            ToolStatus::Success => "success",
67            ToolStatus::Error => "error",
68        }
69    }
70}
71
72// ---------------------------------------------------------------------------
73// Tool trait
74// ---------------------------------------------------------------------------
75
76#[async_trait]
77pub trait Tool: Send + Sync {
78    /// Machine-readable tool name (e.g. "echo").
79    fn name(&self) -> &str;
80
81    /// Human-readable description the LLM uses to decide when to call this tool.
82    fn description(&self) -> &str;
83
84    /// JSON Schema describing the parameters this tool accepts.
85    fn parameters_schema(&self) -> serde_json::Value;
86
87    /// Execute the tool with the given JSON arguments.
88    async fn execute(&self, args: serde_json::Value) -> ToolResult;
89
90    /// Build a `llm_cascade::ToolDefinition` that can be passed to the model.
91    fn to_definition(&self) -> llm_cascade::ToolDefinition {
92        llm_cascade::ToolDefinition {
93            name: self.name().to_owned(),
94            description: self.description().to_owned(),
95            parameters: self.parameters_schema(),
96        }
97    }
98}
99
100// ---------------------------------------------------------------------------
101// ToolRegistry
102// ---------------------------------------------------------------------------
103
104/// A thread-safe registry that maps tool names to their implementations.
105pub struct ToolRegistry {
106    tools: Arc<std::sync::Mutex<HashMap<String, Arc<dyn Tool>>>>,
107}
108
109impl Default for ToolRegistry {
110    fn default() -> Self {
111        Self::new()
112    }
113}
114
115impl ToolRegistry {
116    /// Create an empty registry.
117    pub fn new() -> Self {
118        Self {
119            tools: Arc::new(std::sync::Mutex::new(HashMap::new())),
120        }
121    }
122
123    /// Register a tool. Overwrites any previously registered tool with the same name.
124    pub fn register<T: Tool + 'static>(&self, tool: T) {
125        self.tools
126            .lock()
127            .unwrap()
128            .insert(tool.name().to_owned(), Arc::new(tool));
129    }
130
131    /// Register a tool that is already boxed as `Arc<dyn Tool>`.
132    pub fn register_arc(&self, tool: Arc<dyn Tool>) {
133        self.tools
134            .lock()
135            .unwrap()
136            .insert(tool.name().to_owned(), tool);
137    }
138
139    /// Look up a tool by name.
140    pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
141        self.tools.lock().unwrap().get(name).cloned()
142    }
143
144    /// Return all registered tool names.
145    pub fn tool_names(&self) -> Vec<String> {
146        self.tools.lock().unwrap().keys().cloned().collect()
147    }
148
149    /// Return all registered tools as `Arc<dyn Tool>`.
150    pub fn all_tools(&self) -> Vec<Arc<dyn Tool>> {
151        self.tools.lock().unwrap().values().cloned().collect()
152    }
153
154    /// Build a list of `ToolDefinition` for every registered tool.
155    pub fn tool_definitions(&self) -> Vec<llm_cascade::ToolDefinition> {
156        self.tools
157            .lock()
158            .unwrap()
159            .values()
160            .map(|t| t.to_definition())
161            .collect()
162    }
163
164    /// Convenience: execute a tool by name with the given JSON args.
165    pub async fn execute(&self, name: &str, args: serde_json::Value) -> Result<ToolResult> {
166        let tool = self.get(name).ok_or_else(|| AgentError::ToolFailed {
167            tool: name.to_owned(),
168            reason: format!("Tool '{}' is not registered", name),
169        })?;
170        Ok(tool.execute(args).await)
171    }
172
173    /// Create a `ListToolsTool` that references this registry and register it.
174    pub fn register_list_tools(&self) {
175        let names = self.tool_names();
176        let tools = self.all_tools();
177        self.register(builtin::ListToolsTool::new(names, tools));
178    }
179}