daimon_core/
tool.rs

1//! Tool trait, output/choice types, and retry policy. Implement [`Tool`] for
2//! callable functions an agent can invoke; the registry lives in the `daimon`
3//! facade crate.
4
5use std::future::Future;
6use std::pin::Pin;
7use std::sync::Arc;
8use std::time::Duration;
9
10use serde::Serialize;
11
12use crate::error::Result;
13
14/// The result of executing a tool.
15#[derive(Debug, Clone)]
16pub struct ToolOutput {
17    /// The output content (text or serialized JSON).
18    pub content: String,
19    /// Whether this output represents an error.
20    pub is_error: bool,
21}
22
23impl ToolOutput {
24    /// Create a successful text output.
25    pub fn text(content: impl Into<String>) -> Self {
26        Self {
27            content: content.into(),
28            is_error: false,
29        }
30    }
31
32    /// Create a successful output from a serializable value, encoded as JSON.
33    pub fn json(value: &impl Serialize) -> Result<Self> {
34        Ok(Self {
35            content: serde_json::to_string(value)?,
36            is_error: false,
37        })
38    }
39
40    /// Create an error output.
41    pub fn error(content: impl Into<String>) -> Self {
42        Self {
43            content: content.into(),
44            is_error: true,
45        }
46    }
47}
48
49/// Controls which tools the model is allowed to use.
50#[derive(Debug, Clone, Default)]
51pub enum ToolChoice {
52    /// Let the model decide whether to call tools.
53    #[default]
54    Auto,
55    /// Prevent the model from calling any tools.
56    None,
57    /// Force the model to call at least one tool.
58    Required,
59    /// Force the model to call a specific tool by name.
60    Specific(String),
61}
62
63/// Strategy for computing delay between retries.
64#[derive(Debug, Clone)]
65pub enum BackoffStrategy {
66    /// Fixed delay between retries.
67    Fixed(Duration),
68    /// Exponential backoff: base * 2^attempt, capped at max.
69    Exponential { base: Duration, max: Duration },
70}
71
72impl BackoffStrategy {
73    pub fn delay_for(&self, attempt: usize) -> Duration {
74        match self {
75            BackoffStrategy::Fixed(d) => *d,
76            BackoffStrategy::Exponential { base, max } => {
77                let millis = base.as_millis() as u64 * 2u64.saturating_pow(attempt as u32);
78                Duration::from_millis(millis).min(*max)
79            }
80        }
81    }
82}
83
84/// Policy controlling when and how tool execution is retried on failure.
85#[derive(Debug, Clone)]
86pub struct ToolRetryPolicy {
87    /// Maximum number of retry attempts (0 = no retries).
88    pub max_retries: usize,
89    /// Backoff strategy between attempts.
90    pub backoff: BackoffStrategy,
91    /// If set, only errors whose message contains one of these substrings are retried.
92    pub retryable_patterns: Vec<String>,
93}
94
95impl ToolRetryPolicy {
96    /// Creates a policy with exponential backoff (100ms base, 10s max).
97    pub fn exponential(max_retries: usize) -> Self {
98        Self {
99            max_retries,
100            backoff: BackoffStrategy::Exponential {
101                base: Duration::from_millis(100),
102                max: Duration::from_secs(10),
103            },
104            retryable_patterns: Vec::new(),
105        }
106    }
107
108    /// Creates a policy with fixed delay between retries.
109    pub fn fixed(max_retries: usize, delay: Duration) -> Self {
110        Self {
111            max_retries,
112            backoff: BackoffStrategy::Fixed(delay),
113            retryable_patterns: Vec::new(),
114        }
115    }
116
117    /// Only retry if the error message contains one of these substrings.
118    pub fn retryable_on(mut self, patterns: Vec<String>) -> Self {
119        self.retryable_patterns = patterns;
120        self
121    }
122
123    /// Returns true if the given error message is eligible for retry.
124    pub fn is_retryable(&self, error_msg: &str) -> bool {
125        if self.retryable_patterns.is_empty() {
126            return true;
127        }
128        self.retryable_patterns
129            .iter()
130            .any(|p| error_msg.contains(p.as_str()))
131    }
132}
133
134impl Default for ToolRetryPolicy {
135    fn default() -> Self {
136        Self::exponential(2)
137    }
138}
139
140/// Trait for tools the agent can invoke. Tools must have unique names and declare a JSON Schema for parameters.
141pub trait Tool: Send + Sync {
142    /// Unique identifier for the tool. Used by the model when requesting a call.
143    fn name(&self) -> &str;
144    /// Human-readable description. The model uses this to decide when to call the tool.
145    fn description(&self) -> &str;
146    /// JSON Schema for the tool's parameters. Validates and guides the model's argument generation.
147    fn parameters_schema(&self) -> serde_json::Value;
148
149    /// Executes the tool with the given arguments. Arguments are validated by the model; implementors may still validate.
150    fn execute(&self, input: &serde_json::Value)
151    -> impl Future<Output = Result<ToolOutput>> + Send;
152
153    /// Per-tool retry policy. If `Some`, overrides the agent-level retry policy
154    /// for this tool. Return `None` to use the agent's default.
155    fn retry_policy(&self) -> Option<ToolRetryPolicy> {
156        None
157    }
158}
159
160/// Object-safe wrapper for the `Tool` trait, enabling dynamic dispatch via `Arc<dyn ErasedTool>`.
161pub trait ErasedTool: Send + Sync {
162    fn name(&self) -> &str;
163    fn description(&self) -> &str;
164    fn parameters_schema(&self) -> serde_json::Value;
165
166    fn execute_erased<'a>(
167        &'a self,
168        input: &'a serde_json::Value,
169    ) -> Pin<Box<dyn Future<Output = Result<ToolOutput>> + Send + 'a>>;
170
171    fn retry_policy(&self) -> Option<ToolRetryPolicy>;
172}
173
174impl<T: Tool> ErasedTool for T {
175    fn name(&self) -> &str {
176        Tool::name(self)
177    }
178
179    fn description(&self) -> &str {
180        Tool::description(self)
181    }
182
183    fn parameters_schema(&self) -> serde_json::Value {
184        Tool::parameters_schema(self)
185    }
186
187    fn execute_erased<'a>(
188        &'a self,
189        input: &'a serde_json::Value,
190    ) -> Pin<Box<dyn Future<Output = Result<ToolOutput>> + Send + 'a>> {
191        Box::pin(Tool::execute(self, input))
192    }
193
194    fn retry_policy(&self) -> Option<ToolRetryPolicy> {
195        Tool::retry_policy(self)
196    }
197}
198
199/// Shared ownership of a tool via `Arc<dyn ErasedTool>`. Used by registry and agent.
200pub type SharedTool = Arc<dyn ErasedTool>;
201
202#[cfg(test)]
203mod tests {
204    use super::BackoffStrategy;
205    use super::{SharedTool, Tool, ToolOutput, ToolRetryPolicy};
206    use crate::Result;
207    use std::sync::Arc;
208    use std::time::Duration;
209
210    /// A provider-crate-style Tool impl using only daimon_core items.
211    struct Echo;
212
213    impl Tool for Echo {
214        fn name(&self) -> &str {
215            "echo"
216        }
217
218        fn description(&self) -> &str {
219            "Echoes its input back."
220        }
221
222        fn parameters_schema(&self) -> serde_json::Value {
223            serde_json::json!({
224                "type": "object",
225                "properties": { "text": { "type": "string" } },
226                "required": ["text"]
227            })
228        }
229
230        async fn execute(&self, input: &serde_json::Value) -> Result<ToolOutput> {
231            Ok(ToolOutput::text(input["text"].as_str().unwrap_or_default()))
232        }
233    }
234
235    #[tokio::test]
236    async fn tool_is_implementable_from_core_alone() {
237        let tool = Echo;
238        assert_eq!(tool.name(), "echo");
239        assert!(tool.retry_policy().is_none());
240
241        let out = tool
242            .execute(&serde_json::json!({ "text": "hi" }))
243            .await
244            .unwrap();
245        assert_eq!(out.content, "hi");
246        assert!(!out.is_error);
247
248        let shared: SharedTool = Arc::new(Echo);
249        let out = shared
250            .execute_erased(&serde_json::json!({ "text": "yo" }))
251            .await
252            .unwrap();
253        assert_eq!(out.content, "yo");
254
255        let policy = ToolRetryPolicy::exponential(3).retryable_on(vec!["timeout".into()]);
256        assert!(policy.is_retryable("connection timeout"));
257        assert!(!policy.is_retryable("invalid arguments"));
258    }
259
260    #[test]
261    fn test_exponential_backoff() {
262        let strategy = BackoffStrategy::Exponential {
263            base: Duration::from_millis(100),
264            max: Duration::from_secs(5),
265        };
266        assert_eq!(strategy.delay_for(0), Duration::from_millis(100));
267        assert_eq!(strategy.delay_for(1), Duration::from_millis(200));
268        assert_eq!(strategy.delay_for(2), Duration::from_millis(400));
269        assert_eq!(strategy.delay_for(10), Duration::from_secs(5));
270    }
271
272    #[test]
273    fn test_fixed_backoff() {
274        let strategy = BackoffStrategy::Fixed(Duration::from_secs(1));
275        assert_eq!(strategy.delay_for(0), Duration::from_secs(1));
276        assert_eq!(strategy.delay_for(5), Duration::from_secs(1));
277    }
278
279    #[test]
280    fn test_retryable_patterns() {
281        let policy = ToolRetryPolicy::exponential(3)
282            .retryable_on(vec!["timeout".into(), "rate limit".into()]);
283        assert!(policy.is_retryable("connection timeout"));
284        assert!(policy.is_retryable("rate limit exceeded"));
285        assert!(!policy.is_retryable("invalid arguments"));
286    }
287
288    #[test]
289    fn test_empty_patterns_retries_everything() {
290        let policy = ToolRetryPolicy::exponential(3);
291        assert!(policy.is_retryable("any error at all"));
292    }
293}
daimon_core/tool.rs

daimon_core/
tool.rs
