hanzo_agent/
tool.rs

1//! Tool system for agents
2
3use crate::context::RunContext;
4use crate::errors::{AgentError, Result};
5use async_trait::async_trait;
6use serde_json::Value;
7use std::sync::Arc;
8
9/// Tool trait for agent tools
10///
11/// Tools are functions that agents can call to perform actions.
12/// They have a name, description, JSON schema, and an invoke method.
13#[async_trait]
14pub trait Tool: Send + Sync {
15    /// The name of the tool
16    fn name(&self) -> &str;
17
18    /// A description of what the tool does
19    fn description(&self) -> &str;
20
21    /// JSON schema for the tool's parameters
22    fn json_schema(&self) -> Value;
23
24    /// Invoke the tool with the given context and arguments
25    ///
26    /// # Arguments
27    /// * `ctx` - The runtime context
28    /// * `args` - JSON string containing the tool arguments
29    ///
30    /// # Returns
31    /// The tool result as a string, or an error
32    async fn invoke(&self, ctx: &RunContext, args: &str) -> Result<String>;
33}
34
35/// A function-based tool implementation
36pub struct FunctionTool {
37    name: String,
38    description: String,
39    json_schema: Value,
40    handler: Arc<dyn Fn(&RunContext, Value) -> Result<String> + Send + Sync>,
41}
42
43impl FunctionTool {
44    /// Create a new function tool
45    pub fn new(
46        name: impl Into<String>,
47        description: impl Into<String>,
48        json_schema: Value,
49        handler: impl Fn(&RunContext, Value) -> Result<String> + Send + Sync + 'static,
50    ) -> Self {
51        Self {
52            name: name.into(),
53            description: description.into(),
54            json_schema,
55            handler: Arc::new(handler),
56        }
57    }
58
59    /// Builder for creating function tools
60    pub fn builder(name: impl Into<String>) -> FunctionToolBuilder {
61        FunctionToolBuilder {
62            name: name.into(),
63            description: String::new(),
64            json_schema: Value::Null,
65            handler: None,
66        }
67    }
68}
69
70#[async_trait]
71impl Tool for FunctionTool {
72    fn name(&self) -> &str {
73        &self.name
74    }
75
76    fn description(&self) -> &str {
77        &self.description
78    }
79
80    fn json_schema(&self) -> Value {
81        self.json_schema.clone()
82    }
83
84    async fn invoke(&self, ctx: &RunContext, args: &str) -> Result<String> {
85        let value: Value =
86            serde_json::from_str(args).map_err(|e| AgentError::InvalidJson(e.to_string()))?;
87        (self.handler)(ctx, value)
88    }
89}
90
91/// Builder for FunctionTool
92pub struct FunctionToolBuilder {
93    name: String,
94    description: String,
95    json_schema: Value,
96    handler: Option<Arc<dyn Fn(&RunContext, Value) -> Result<String> + Send + Sync>>,
97}
98
99impl FunctionToolBuilder {
100    /// Set the tool description
101    pub fn description(mut self, desc: impl Into<String>) -> Self {
102        self.description = desc.into();
103        self
104    }
105
106    /// Set the JSON schema
107    pub fn schema(mut self, schema: Value) -> Self {
108        self.json_schema = schema;
109        self
110    }
111
112    /// Set the handler function
113    pub fn handler<F>(mut self, f: F) -> Self
114    where
115        F: Fn(&RunContext, Value) -> Result<String> + Send + Sync + 'static,
116    {
117        self.handler = Some(Arc::new(f));
118        self
119    }
120
121    /// Build the function tool
122    pub fn build(self) -> Result<FunctionTool> {
123        let handler = self
124            .handler
125            .ok_or_else(|| AgentError::Configuration("Tool handler not set".to_string()))?;
126
127        Ok(FunctionTool {
128            name: self.name,
129            description: self.description,
130            json_schema: self.json_schema,
131            handler,
132        })
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139    use serde_json::json;
140
141    #[tokio::test]
142    async fn test_function_tool() {
143        let tool = FunctionTool::builder("test_tool")
144            .description("A test tool")
145            .schema(json!({"type": "object"}))
146            .handler(|_ctx, _args| Ok("test result".to_string()))
147            .build()
148            .unwrap();
149
150        assert_eq!(tool.name(), "test_tool");
151        assert_eq!(tool.description(), "A test tool");
152
153        let ctx = RunContext::new();
154        let result = tool.invoke(&ctx, "{}").await.unwrap();
155        assert_eq!(result, "test result");
156    }
157}
hanzo_agent/tool.rs

hanzo_agent/
tool.rs
