Skip to main content

agent_io/tools/
builder.rs

1//! Tool builder
2
3use serde::de::DeserializeOwned;
4use serde_json::{Value, json};
5
6use crate::Result;
7
8use super::function::FunctionTool;
9use super::tool::{EphemeralConfig, Tool};
10
11/// Builder for creating tools
12pub struct ToolBuilder {
13    name: String,
14    description: String,
15    parameters_schema: serde_json::Map<String, Value>,
16    ephemeral: EphemeralConfig,
17}
18
19impl ToolBuilder {
20    pub fn new(name: impl Into<String>) -> Self {
21        Self {
22            name: name.into(),
23            description: String::new(),
24            parameters_schema: serde_json::Map::new(),
25            ephemeral: EphemeralConfig::None,
26        }
27    }
28
29    pub fn description(mut self, desc: impl Into<String>) -> Self {
30        self.description = desc.into();
31        self
32    }
33
34    pub fn parameter(mut self, name: &str, schema: Value) -> Self {
35        self.parameters_schema.insert(name.to_string(), schema);
36        self
37    }
38
39    pub fn string_param(self, name: &str, description: &str) -> Self {
40        self.parameter(
41            name,
42            json!({
43                "type": "string",
44                "description": description
45            }),
46        )
47    }
48
49    pub fn number_param(self, name: &str, description: &str) -> Self {
50        self.parameter(
51            name,
52            json!({
53                "type": "number",
54                "description": description
55            }),
56        )
57    }
58
59    pub fn boolean_param(self, name: &str, description: &str) -> Self {
60        self.parameter(
61            name,
62            json!({
63                "type": "boolean",
64                "description": description
65            }),
66        )
67    }
68
69    pub fn ephemeral(mut self, config: EphemeralConfig) -> Self {
70        self.ephemeral = config;
71        self
72    }
73
74    pub fn build<F, T>(self, func: F) -> Box<dyn Tool>
75    where
76        T: DeserializeOwned + Send + Sync + 'static,
77        F: Fn(T) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String>> + Send>>
78            + Send
79            + Sync
80            + 'static,
81    {
82        let mut tool = FunctionTool::new(self.name, self.description, self.parameters_schema, func);
83        tool.ephemeral_config = self.ephemeral;
84        Box::new(tool)
85    }
86}