adk_tool/
function_tool.rs

1use adk_core::{Result, Tool, ToolContext};
2use async_trait::async_trait;
3use schemars::{JsonSchema, schema::RootSchema};
4use serde::Serialize;
5use serde_json::Value;
6use std::future::Future;
7use std::pin::Pin;
8use std::sync::Arc;
9
10type AsyncHandler = Box<
11    dyn Fn(Arc<dyn ToolContext>, Value) -> Pin<Box<dyn Future<Output = Result<Value>> + Send>>
12        + Send
13        + Sync,
14>;
15
16pub struct FunctionTool {
17    name: String,
18    description: String,
19    handler: AsyncHandler,
20    long_running: bool,
21    read_only: bool,
22    concurrency_safe: bool,
23    parameters_schema: Option<Value>,
24    response_schema: Option<Value>,
25    scopes: Vec<&'static str>,
26}
27
28impl FunctionTool {
29    pub fn new<F, Fut>(name: impl Into<String>, description: impl Into<String>, handler: F) -> Self
30    where
31        F: Fn(Arc<dyn ToolContext>, Value) -> Fut + Send + Sync + 'static,
32        Fut: Future<Output = Result<Value>> + Send + 'static,
33    {
34        Self {
35            name: name.into(),
36            description: description.into(),
37            handler: Box::new(move |ctx, args| Box::pin(handler(ctx, args))),
38            long_running: false,
39            read_only: false,
40            concurrency_safe: false,
41            parameters_schema: None,
42            response_schema: None,
43            scopes: Vec::new(),
44        }
45    }
46
47    pub fn with_long_running(mut self, long_running: bool) -> Self {
48        self.long_running = long_running;
49        self
50    }
51
52    pub fn with_read_only(mut self, read_only: bool) -> Self {
53        self.read_only = read_only;
54        self
55    }
56
57    pub fn with_concurrency_safe(mut self, concurrency_safe: bool) -> Self {
58        self.concurrency_safe = concurrency_safe;
59        self
60    }
61
62    pub fn with_parameters_schema<T>(mut self) -> Self
63    where
64        T: JsonSchema + Serialize,
65    {
66        self.parameters_schema = Some(generate_schema::<T>());
67        self
68    }
69
70    pub fn with_response_schema<T>(mut self) -> Self
71    where
72        T: JsonSchema + Serialize,
73    {
74        self.response_schema = Some(generate_schema::<T>());
75        self
76    }
77
78    /// Declare the scopes required to execute this tool.
79    ///
80    /// When set, the framework will enforce that the calling user possesses
81    /// **all** listed scopes before dispatching `execute()`.
82    ///
83    /// # Example
84    ///
85    /// ```rust,ignore
86    /// let tool = FunctionTool::new("transfer", "Transfer funds", handler)
87    ///     .with_scopes(&["finance:write", "verified"]);
88    /// ```
89    pub fn with_scopes(mut self, scopes: &[&'static str]) -> Self {
90        self.scopes = scopes.to_vec();
91        self
92    }
93
94    pub fn parameters_schema(&self) -> Option<&Value> {
95        self.parameters_schema.as_ref()
96    }
97
98    pub fn response_schema(&self) -> Option<&Value> {
99        self.response_schema.as_ref()
100    }
101}
102
103/// The note appended to long-running tool descriptions to prevent duplicate calls.
104const LONG_RUNNING_NOTE: &str = "NOTE: This is a long-running operation. Do not call this tool again if it has already returned some intermediate or pending status.";
105
106#[async_trait]
107impl Tool for FunctionTool {
108    fn name(&self) -> &str {
109        &self.name
110    }
111
112    fn description(&self) -> &str {
113        &self.description
114    }
115
116    /// Returns an enhanced description for long-running tools that includes
117    /// a note warning the model not to call the tool again if it's already pending.
118    fn enhanced_description(&self) -> String {
119        if self.long_running {
120            if self.description.is_empty() {
121                LONG_RUNNING_NOTE.to_string()
122            } else {
123                format!("{}\n\n{}", self.description, LONG_RUNNING_NOTE)
124            }
125        } else {
126            self.description.clone()
127        }
128    }
129
130    fn is_long_running(&self) -> bool {
131        self.long_running
132    }
133
134    fn is_read_only(&self) -> bool {
135        self.read_only
136    }
137
138    fn is_concurrency_safe(&self) -> bool {
139        self.concurrency_safe
140    }
141
142    fn parameters_schema(&self) -> Option<Value> {
143        self.parameters_schema.clone()
144    }
145
146    fn response_schema(&self) -> Option<Value> {
147        self.response_schema.clone()
148    }
149
150    fn required_scopes(&self) -> &[&str] {
151        &self.scopes
152    }
153
154    #[adk_telemetry::instrument(
155        skip(self, ctx, args),
156        fields(
157            tool.name = %self.name,
158            tool.description = %self.description,
159            tool.long_running = %self.long_running,
160            function_call.id = %ctx.function_call_id()
161        )
162    )]
163    async fn execute(&self, ctx: Arc<dyn ToolContext>, args: Value) -> Result<Value> {
164        adk_telemetry::debug!("Executing tool");
165        (self.handler)(ctx, args).await
166    }
167}
168
169fn generate_schema<T>() -> Value
170where
171    T: JsonSchema + Serialize,
172{
173    let settings = schemars::r#gen::SchemaSettings::openapi3().with(|s| {
174        s.inline_subschemas = true;
175        s.meta_schema = None;
176    });
177    let generator = schemars::r#gen::SchemaGenerator::new(settings);
178    let mut schema: RootSchema = generator.into_root_schema_for::<T>();
179    schema.schema.metadata().title = None;
180    serde_json::to_value(schema).unwrap()
181}
adk_tool/function_tool.rs

adk_tool/
function_tool.rs
