Skip to main content

agent_io/tools/
function.rs

1//! Function tool implementation
2
3use async_trait::async_trait;
4use serde::de::DeserializeOwned;
5use serde_json::Value;
6
7use crate::Result;
8use crate::llm::ToolDefinition;
9
10use super::tool::{EphemeralConfig, Tool, ToolResult};
11
12/// A tool implementation using a function
13pub struct FunctionTool<T, F>
14where
15    T: DeserializeOwned + Send + Sync + 'static,
16    F: Fn(T) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String>> + Send>>
17        + Send
18        + Sync,
19{
20    name: String,
21    description: String,
22    parameters_schema: serde_json::Map<String, Value>,
23    func: F,
24    pub(super) ephemeral_config: EphemeralConfig,
25    _marker: std::marker::PhantomData<T>,
26}
27
28impl<T, F> FunctionTool<T, F>
29where
30    T: DeserializeOwned + Send + Sync + 'static,
31    F: Fn(T) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String>> + Send>>
32        + Send
33        + Sync,
34{
35    /// Create a new function tool
36    pub fn new(
37        name: impl Into<String>,
38        description: impl Into<String>,
39        parameters_schema: serde_json::Map<String, Value>,
40        func: F,
41    ) -> Self {
42        Self {
43            name: name.into(),
44            description: description.into(),
45            parameters_schema,
46            func,
47            ephemeral_config: EphemeralConfig::None,
48            _marker: std::marker::PhantomData,
49        }
50    }
51
52    /// Set ephemeral configuration
53    pub fn with_ephemeral(mut self, config: EphemeralConfig) -> Self {
54        self.ephemeral_config = config;
55        self
56    }
57}
58
59#[async_trait]
60impl<T, F> Tool for FunctionTool<T, F>
61where
62    T: DeserializeOwned + Send + Sync + 'static,
63    F: Fn(T) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String>> + Send>>
64        + Send
65        + Sync,
66{
67    fn name(&self) -> &str {
68        &self.name
69    }
70
71    fn description(&self) -> &str {
72        &self.description
73    }
74
75    fn definition(&self) -> ToolDefinition {
76        ToolDefinition::new(
77            &self.name,
78            &self.description,
79            self.parameters_schema.clone(),
80        )
81    }
82
83    async fn execute(&self, args: Value) -> Result<ToolResult> {
84        let parsed: T = serde_json::from_value(args)?;
85        let content = (self.func)(parsed).await?;
86        Ok(ToolResult::new("", content)
87            .with_ephemeral(self.ephemeral_config != EphemeralConfig::None))
88    }
89
90    fn ephemeral(&self) -> EphemeralConfig {
91        self.ephemeral_config
92    }
93}