agents_toolkit/
builder.rs

1//! Tool builder utilities for creating tools from functions
2//!
3//! This module provides ergonomic helpers for converting regular Rust functions
4//! into Tool implementations that can be registered with agents.
5
6use agents_core::tools::{Tool, ToolBox, ToolContext, ToolParameterSchema, ToolResult, ToolSchema};
7use async_trait::async_trait;
8use serde_json::Value;
9use std::future::Future;
10use std::pin::Pin;
11use std::sync::Arc;
12
13type BoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
14
15/// Type alias for async tool handler functions
16pub type AsyncToolFn =
17    Arc<dyn Fn(Value, ToolContext) -> BoxFuture<'static, anyhow::Result<ToolResult>> + Send + Sync>;
18
19/// Type alias for sync tool handler functions
20pub type SyncToolFn = Arc<dyn Fn(Value, ToolContext) -> anyhow::Result<ToolResult> + Send + Sync>;
21
22/// A tool implementation backed by a function/closure
23pub struct FunctionTool {
24    schema: ToolSchema,
25    handler: AsyncToolFn,
26}
27
28#[async_trait]
29impl Tool for FunctionTool {
30    fn schema(&self) -> ToolSchema {
31        self.schema.clone()
32    }
33
34    async fn execute(&self, args: Value, ctx: ToolContext) -> anyhow::Result<ToolResult> {
35        (self.handler)(args, ctx).await
36    }
37}
38
39/// Builder for creating tools from async functions
40pub struct ToolBuilder {
41    name: String,
42    description: String,
43    parameters: Option<ToolParameterSchema>,
44}
45
46impl ToolBuilder {
47    /// Start building a new tool with the given name and description
48    pub fn new(name: impl Into<String>, description: impl Into<String>) -> Self {
49        Self {
50            name: name.into(),
51            description: description.into(),
52            parameters: None,
53        }
54    }
55
56    /// Set the parameter schema for this tool
57    pub fn with_parameters(mut self, parameters: ToolParameterSchema) -> Self {
58        self.parameters = Some(parameters);
59        self
60    }
61
62    /// Build the tool with an async handler function
63    pub fn build_async<F, Fut>(self, handler: F) -> ToolBox
64    where
65        F: Fn(Value, ToolContext) -> Fut + Send + Sync + 'static,
66        Fut: Future<Output = anyhow::Result<ToolResult>> + Send + 'static,
67    {
68        let schema = ToolSchema::new(
69            self.name,
70            self.description,
71            self.parameters.unwrap_or_else(|| {
72                ToolParameterSchema::object("No parameters", Default::default(), Vec::new())
73            }),
74        );
75
76        let handler: AsyncToolFn = Arc::new(move |args, ctx| Box::pin(handler(args, ctx)));
77
78        Arc::new(FunctionTool { schema, handler })
79    }
80
81    /// Build the tool with a sync handler function
82    pub fn build_sync<F>(self, handler: F) -> ToolBox
83    where
84        F: Fn(Value, ToolContext) -> anyhow::Result<ToolResult> + Send + Sync + 'static,
85    {
86        let handler = Arc::new(handler);
87        self.build_async(move |args, ctx| {
88            let handler = handler.clone();
89            async move { handler(args, ctx) }
90        })
91    }
92}
93
94/// Quick helper to create a simple async tool
95pub fn tool<F, Fut>(
96    name: impl Into<String>,
97    description: impl Into<String>,
98    parameters: ToolParameterSchema,
99    handler: F,
100) -> ToolBox
101where
102    F: Fn(Value, ToolContext) -> Fut + Send + Sync + 'static,
103    Fut: Future<Output = anyhow::Result<ToolResult>> + Send + 'static,
104{
105    ToolBuilder::new(name, description)
106        .with_parameters(parameters)
107        .build_async(handler)
108}
109
110/// Quick helper to create a simple sync tool
111pub fn tool_sync<F>(
112    name: impl Into<String>,
113    description: impl Into<String>,
114    parameters: ToolParameterSchema,
115    handler: F,
116) -> ToolBox
117where
118    F: Fn(Value, ToolContext) -> anyhow::Result<ToolResult> + Send + Sync + 'static,
119{
120    ToolBuilder::new(name, description)
121        .with_parameters(parameters)
122        .build_sync(handler)
123}
124
125/// Simple helper to create a tool from an async closure (backwards compatibility)
126/// This is a simpler API for basic tools that don't need explicit parameter schemas.
127/// The schema will be inferred as accepting any JSON object.
128pub fn create_tool<F, Fut>(
129    name: impl Into<String>,
130    description: impl Into<String>,
131    handler: F,
132) -> ToolBox
133where
134    F: Fn(Value) -> Fut + Send + Sync + 'static,
135    Fut: Future<Output = anyhow::Result<String>> + Send + 'static,
136{
137    let name = name.into();
138    let description_str = description.into();
139
140    // Create a simple schema that accepts any parameters
141    let parameters = ToolParameterSchema::object("Tool parameters", Default::default(), Vec::new());
142
143    let handler = Arc::new(handler);
144
145    ToolBuilder::new(name, description_str)
146        .with_parameters(parameters)
147        .build_async(move |args, ctx| {
148            let handler = handler.clone();
149            async move {
150                let result = handler(args).await?;
151                Ok(ToolResult::text(&ctx, result))
152            }
153        })
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159    use agents_core::state::AgentStateSnapshot;
160    use serde_json::json;
161
162    #[tokio::test]
163    async fn function_tool_executes_handler() {
164        let tool = ToolBuilder::new("echo", "Echoes input")
165            .with_parameters(ToolParameterSchema::object(
166                "Echo parameters",
167                [(
168                    "message".to_string(),
169                    ToolParameterSchema::string("Message to echo"),
170                )]
171                .into_iter()
172                .collect(),
173                vec!["message".to_string()],
174            ))
175            .build_async(|args, ctx| async move {
176                let msg = args["message"].as_str().unwrap_or("empty");
177                Ok(ToolResult::text(&ctx, format!("Echo: {}", msg)))
178            });
179
180        let schema = tool.schema();
181        assert_eq!(schema.name, "echo");
182        assert_eq!(schema.description, "Echoes input");
183
184        let ctx = ToolContext::new(Arc::new(AgentStateSnapshot::default()));
185        let result = tool
186            .execute(json!({"message": "hello"}), ctx)
187            .await
188            .unwrap();
189
190        match result {
191            ToolResult::Message(msg) => {
192                assert_eq!(msg.content.as_text().unwrap(), "Echo: hello");
193            }
194            _ => panic!("Expected message result"),
195        }
196    }
197
198    #[tokio::test]
199    async fn sync_tool_works() {
200        let tool = tool_sync(
201            "add",
202            "Adds two numbers",
203            ToolParameterSchema::object(
204                "Add parameters",
205                [
206                    ("a".to_string(), ToolParameterSchema::number("First number")),
207                    (
208                        "b".to_string(),
209                        ToolParameterSchema::number("Second number"),
210                    ),
211                ]
212                .into_iter()
213                .collect(),
214                vec!["a".to_string(), "b".to_string()],
215            ),
216            |args, ctx| {
217                let a = args["a"].as_f64().unwrap_or(0.0);
218                let b = args["b"].as_f64().unwrap_or(0.0);
219                let sum = a + b;
220                Ok(ToolResult::text(&ctx, format!("Sum: {}", sum)))
221            },
222        );
223
224        let ctx = ToolContext::new(Arc::new(AgentStateSnapshot::default()));
225        let result = tool.execute(json!({"a": 5, "b": 3}), ctx).await.unwrap();
226
227        match result {
228            ToolResult::Message(msg) => {
229                assert_eq!(msg.content.as_text().unwrap(), "Sum: 8");
230            }
231            _ => panic!("Expected message result"),
232        }
233    }
234}
agents_toolkit/builder.rs

agents_toolkit/
builder.rs
