Skip to main content

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