adk_tool/
function_tool.rs

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

adk_tool/
function_tool.rs
