adk_core/
tool.rs

1use crate::{CallbackContext, EventActions, MemoryEntry, Result};
2use async_trait::async_trait;
3use serde_json::Value;
4use std::sync::Arc;
5
6#[async_trait]
7pub trait Tool: Send + Sync {
8    fn name(&self) -> &str;
9    fn description(&self) -> &str;
10
11    /// Returns an enhanced description that may include additional notes.
12    /// For long-running tools, this includes a warning not to call the tool
13    /// again if it has already returned a pending status.
14    /// Default implementation returns the base description.
15    fn enhanced_description(&self) -> String {
16        self.description().to_string()
17    }
18
19    /// Indicates whether the tool is a long-running operation.
20    /// Long-running tools typically return a task ID immediately and
21    /// complete the operation asynchronously.
22    fn is_long_running(&self) -> bool {
23        false
24    }
25    fn parameters_schema(&self) -> Option<Value> {
26        None
27    }
28    fn response_schema(&self) -> Option<Value> {
29        None
30    }
31    async fn execute(&self, ctx: Arc<dyn ToolContext>, args: Value) -> Result<Value>;
32}
33
34#[async_trait]
35pub trait ToolContext: CallbackContext {
36    fn function_call_id(&self) -> &str;
37    fn actions(&self) -> &EventActions;
38    async fn search_memory(&self, query: &str) -> Result<Vec<MemoryEntry>>;
39}
40
41#[async_trait]
42pub trait Toolset: Send + Sync {
43    fn name(&self) -> &str;
44    async fn tools(&self, ctx: Arc<dyn crate::ReadonlyContext>) -> Result<Vec<Arc<dyn Tool>>>;
45}
46
47pub type ToolPredicate = Box<dyn Fn(&dyn Tool) -> bool + Send + Sync>;
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52    use crate::{Content, EventActions, ReadonlyContext, RunConfig};
53
54    struct TestTool {
55        name: String,
56    }
57
58    #[allow(dead_code)]
59    struct TestContext {
60        content: Content,
61        config: RunConfig,
62        actions: EventActions,
63    }
64
65    impl TestContext {
66        fn new() -> Self {
67            Self {
68                content: Content::new("user"),
69                config: RunConfig::default(),
70                actions: EventActions::default(),
71            }
72        }
73    }
74
75    #[async_trait]
76    impl ReadonlyContext for TestContext {
77        fn invocation_id(&self) -> &str {
78            "test"
79        }
80        fn agent_name(&self) -> &str {
81            "test"
82        }
83        fn user_id(&self) -> &str {
84            "user"
85        }
86        fn app_name(&self) -> &str {
87            "app"
88        }
89        fn session_id(&self) -> &str {
90            "session"
91        }
92        fn branch(&self) -> &str {
93            ""
94        }
95        fn user_content(&self) -> &Content {
96            &self.content
97        }
98    }
99
100    #[async_trait]
101    impl CallbackContext for TestContext {
102        fn artifacts(&self) -> Option<Arc<dyn crate::Artifacts>> {
103            None
104        }
105    }
106
107    #[async_trait]
108    impl ToolContext for TestContext {
109        fn function_call_id(&self) -> &str {
110            "call-123"
111        }
112        fn actions(&self) -> &EventActions {
113            &self.actions
114        }
115        async fn search_memory(&self, _query: &str) -> Result<Vec<crate::MemoryEntry>> {
116            Ok(vec![])
117        }
118    }
119
120    #[async_trait]
121    impl Tool for TestTool {
122        fn name(&self) -> &str {
123            &self.name
124        }
125
126        fn description(&self) -> &str {
127            "test tool"
128        }
129
130        async fn execute(&self, _ctx: Arc<dyn ToolContext>, _args: Value) -> Result<Value> {
131            Ok(Value::String("result".to_string()))
132        }
133    }
134
135    #[test]
136    fn test_tool_trait() {
137        let tool = TestTool { name: "test".to_string() };
138        assert_eq!(tool.name(), "test");
139        assert_eq!(tool.description(), "test tool");
140        assert!(!tool.is_long_running());
141    }
142
143    #[tokio::test]
144    async fn test_tool_execute() {
145        let tool = TestTool { name: "test".to_string() };
146        let ctx = Arc::new(TestContext::new()) as Arc<dyn ToolContext>;
147        let result = tool.execute(ctx, Value::Null).await.unwrap();
148        assert_eq!(result, Value::String("result".to_string()));
149    }
150}
adk_core/tool.rs

adk_core/
tool.rs
