adk_tool/
function_tool.rs1use 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 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
89const 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 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}