Skip to main content

oxi_agent/tools/
tool_definition_wrapper.rs

1/// Tool definition wrapping utilities
2/// Provides adapters for converting between tool representations.
3use crate::tools::{AgentTool, AgentToolResult, ToolContext, ToolError};
4use crate::types::ToolDefinition;
5use async_trait::async_trait;
6use serde_json::Value;
7use std::sync::Arc;
8use tokio::sync::oneshot;
9
10/// A dynamically-constructed tool from a definition and execution function.
11///
12/// This is useful for wrapping external tool definitions (e.g., from extensions
13/// or plugins) into the `AgentTool` interface without requiring a dedicated struct.
14pub struct DynamicTool {
15    name: String,
16    label: String,
17    description: String,
18    parameters: Value,
19    #[allow(clippy::type_complexity)]
20    execute_fn: Arc<
21        dyn Fn(
22                &str,
23                Value,
24                Option<oneshot::Receiver<()>>,
25            )
26                -> std::pin::Pin<Box<dyn Future<Output = Result<AgentToolResult, ToolError>> + Send>>
27            + Send
28            + Sync,
29    >,
30}
31
32impl DynamicTool {
33    /// Create a new dynamic tool from components.
34    pub fn new(
35        name: impl Into<String>,
36        label: impl Into<String>,
37        description: impl Into<String>,
38        parameters: Value,
39        execute_fn: impl Fn(
40            &str,
41            Value,
42            Option<oneshot::Receiver<()>>,
43        ) -> std::pin::Pin<
44            Box<dyn Future<Output = Result<AgentToolResult, ToolError>> + Send>,
45        > + Send
46        + Sync
47        + 'static,
48    ) -> Self {
49        Self {
50            name: name.into(),
51            label: label.into(),
52            description: description.into(),
53            parameters,
54            execute_fn: Arc::new(execute_fn),
55        }
56    }
57
58    /// Create a `DynamicTool` from a `ToolDefinition` with a simple execution function.
59    pub fn from_definition(
60        def: ToolDefinition,
61        execute_fn: impl Fn(
62            &str,
63            Value,
64            Option<oneshot::Receiver<()>>,
65        ) -> std::pin::Pin<
66            Box<dyn Future<Output = Result<AgentToolResult, ToolError>> + Send>,
67        > + Send
68        + Sync
69        + 'static,
70    ) -> Self {
71        let name_for_label = def.name.clone();
72        let schema =
73            serde_json::to_value(&def.input_schema).unwrap_or(Value::Object(Default::default()));
74        Self {
75            name: def.name,
76            label: name_for_label, // Use name as label fallback
77            description: def.description,
78            parameters: schema,
79            execute_fn: Arc::new(execute_fn),
80        }
81    }
82}
83
84#[async_trait]
85impl AgentTool for DynamicTool {
86    fn name(&self) -> &str {
87        &self.name
88    }
89
90    fn label(&self) -> &str {
91        &self.label
92    }
93
94    fn description(&self) -> &str {
95        &self.description
96    }
97
98    fn parameters_schema(&self) -> Value {
99        self.parameters.clone()
100    }
101
102    async fn execute(
103        &self,
104        tool_call_id: &str,
105        params: Value,
106        signal: Option<oneshot::Receiver<()>>,
107        _ctx: &ToolContext,
108    ) -> Result<AgentToolResult, ToolError> {
109        (self.execute_fn)(tool_call_id, params, signal).await
110    }
111}
112
113/// Trait for tool definitions that can be wrapped into `AgentTool`.
114pub trait ToolDefinitionLike: Send + Sync {
115    /// Returns the tool name used in LLM function-calling.
116    fn tool_name(&self) -> &str;
117    /// Returns a human-readable label for UI display.
118    fn tool_label(&self) -> &str;
119    /// Returns the tool description sent to the LLM.
120    fn tool_description(&self) -> &str;
121    /// Returns the JSON schema for the tool's parameters.
122    fn tool_parameters(&self) -> Value;
123    /// Executes the tool with the given call ID, parameters, and optional cancellation signal.
124    fn tool_execute(
125        &self,
126        tool_call_id: &str,
127        params: Value,
128        signal: Option<oneshot::Receiver<()>>,
129    ) -> std::pin::Pin<Box<dyn Future<Output = Result<AgentToolResult, ToolError>> + Send>>;
130}
131
132/// Wrap a `ToolDefinitionLike` into an `Arc<dyn AgentTool>`.
133pub fn wrap_tool_definition<T: ToolDefinitionLike + 'static>(def: T) -> Arc<dyn AgentTool> {
134    Arc::new(DefinitionWrapper(def))
135}
136
137/// Internal wrapper that adapts `ToolDefinitionLike` to `AgentTool`.
138struct DefinitionWrapper<T>(T);
139
140#[async_trait]
141impl<T: ToolDefinitionLike + 'static> AgentTool for DefinitionWrapper<T> {
142    fn name(&self) -> &str {
143        self.0.tool_name()
144    }
145
146    fn label(&self) -> &str {
147        self.0.tool_label()
148    }
149
150    fn description(&self) -> &str {
151        self.0.tool_description()
152    }
153
154    fn parameters_schema(&self) -> Value {
155        self.0.tool_parameters()
156    }
157
158    async fn execute(
159        &self,
160        tool_call_id: &str,
161        params: Value,
162        signal: Option<oneshot::Receiver<()>>,
163        _ctx: &ToolContext,
164    ) -> Result<AgentToolResult, ToolError> {
165        self.0.tool_execute(tool_call_id, params, signal).await
166    }
167}
168
169/// Create a minimal `ToolDefinition` from an `AgentTool`.
170///
171/// This is useful for internal registries that operate on definitions
172/// even when tools are provided as trait objects.
173pub fn create_tool_definition_from_agent_tool(tool: &dyn AgentTool) -> ToolDefinition {
174    let schema_map: std::collections::HashMap<String, Value> = {
175        let schema = tool.parameters_schema();
176        if let Value::Object(map) = schema {
177            map.into_iter().collect()
178        } else {
179            let mut map = std::collections::HashMap::new();
180            map.insert("type".to_string(), Value::String("object".to_string()));
181            map
182        }
183    };
184
185    ToolDefinition::new(tool.name(), tool.description(), schema_map)
186}
187
188#[cfg(test)]
189mod tests {
190    use super::*;
191
192    #[test]
193    fn test_create_tool_definition_from_agent_tool() {
194        struct TestTool;
195
196        #[async_trait]
197        impl AgentTool for TestTool {
198            fn name(&self) -> &str {
199                "test_tool"
200            }
201            fn label(&self) -> &str {
202                "Test Tool"
203            }
204            fn description(&self) -> &str {
205                "A test tool"
206            }
207            fn parameters_schema(&self) -> Value {
208                serde_json::json!({
209                    "type": "object",
210                    "properties": {
211                        "input": { "type": "string" }
212                    }
213                })
214            }
215            async fn execute(
216                &self,
217                _tool_call_id: &str,
218                _params: Value,
219                _signal: Option<oneshot::Receiver<()>>,
220                _ctx: &ToolContext,
221            ) -> Result<AgentToolResult, ToolError> {
222                Ok(AgentToolResult::success("test"))
223            }
224        }
225
226        let tool = TestTool;
227        let def = create_tool_definition_from_agent_tool(&tool);
228        assert_eq!(def.name, "test_tool");
229        assert_eq!(def.description, "A test tool");
230    }
231
232    #[test]
233    fn test_dynamic_tool_creation() {
234        let tool = DynamicTool::new(
235            "dynamic_test",
236            "Dynamic Test",
237            "A dynamic test tool",
238            serde_json::json!({"type": "object"}),
239            |_id, _params, _signal| {
240                Box::pin(async { Ok(AgentToolResult::success("dynamic result")) })
241            },
242        );
243
244        assert_eq!(tool.name(), "dynamic_test");
245        assert_eq!(tool.label(), "Dynamic Test");
246        assert_eq!(tool.description(), "A dynamic test tool");
247    }
248
249    #[tokio::test]
250    async fn test_dynamic_tool_execution() {
251        let tool = DynamicTool::new(
252            "exec_test",
253            "Exec Test",
254            "Exec test tool",
255            serde_json::json!({"type": "object"}),
256            |_id, params, _signal| {
257                let result = format!("got: {}", params);
258                Box::pin(async move { Ok(AgentToolResult::success(result)) })
259            },
260        );
261
262        let result = tool
263            .execute(
264                "call_1",
265                serde_json::json!({"key": "value"}),
266                None,
267                &ToolContext::default(),
268            )
269            .await;
270        assert!(result.is_ok());
271        assert_eq!(result.unwrap().output, "got: {\"key\":\"value\"}");
272    }
273
274    struct MockDefLike;
275
276    impl ToolDefinitionLike for MockDefLike {
277        fn tool_name(&self) -> &str {
278            "mock_def"
279        }
280        fn tool_label(&self) -> &str {
281            "Mock Definition"
282        }
283        fn tool_description(&self) -> &str {
284            "Mock description"
285        }
286        fn tool_parameters(&self) -> Value {
287            serde_json::json!({"type": "object"})
288        }
289        fn tool_execute(
290            &self,
291            _tool_call_id: &str,
292            _params: Value,
293            _signal: Option<oneshot::Receiver<()>>,
294        ) -> std::pin::Pin<Box<dyn Future<Output = Result<AgentToolResult, ToolError>> + Send>>
295        {
296            Box::pin(async { Ok(AgentToolResult::success("mock result")) })
297        }
298    }
299
300    #[test]
301    fn test_wrap_tool_definition() {
302        let wrapped = wrap_tool_definition(MockDefLike);
303        assert_eq!(wrapped.name(), "mock_def");
304        assert_eq!(wrapped.label(), "Mock Definition");
305    }
306
307    #[tokio::test]
308    async fn test_wrapped_tool_execution() {
309        let wrapped = wrap_tool_definition(MockDefLike);
310        let result = wrapped
311            .execute("call_1", Value::Null, None, &ToolContext::default())
312            .await;
313        assert!(result.is_ok());
314        assert_eq!(result.unwrap().output, "mock result");
315    }
316}