adk_runner/
callbacks.rs

1use adk_core::{CallbackContext, Content, Result};
2use futures::future::BoxFuture;
3use std::sync::Arc;
4
5/// Callback executed before calling the model
6pub type BeforeModelCallback = Box<
7    dyn Fn(Arc<dyn CallbackContext>) -> BoxFuture<'static, Result<Option<Content>>> + Send + Sync,
8>;
9
10/// Callback executed after model response
11pub type AfterModelCallback = Box<
12    dyn Fn(Arc<dyn CallbackContext>) -> BoxFuture<'static, Result<Option<Content>>> + Send + Sync,
13>;
14
15/// Callback executed before tool execution
16pub type BeforeToolCallback = Box<
17    dyn Fn(Arc<dyn CallbackContext>) -> BoxFuture<'static, Result<Option<Content>>> + Send + Sync,
18>;
19
20/// Callback executed after tool execution
21pub type AfterToolCallback = Box<
22    dyn Fn(Arc<dyn CallbackContext>) -> BoxFuture<'static, Result<Option<Content>>> + Send + Sync,
23>;
24
25/// Collection of all callback types
26pub struct Callbacks {
27    pub before_model: Vec<BeforeModelCallback>,
28    pub after_model: Vec<AfterModelCallback>,
29    pub before_tool: Vec<BeforeToolCallback>,
30    pub after_tool: Vec<AfterToolCallback>,
31}
32
33impl Default for Callbacks {
34    fn default() -> Self {
35        Self::new()
36    }
37}
38
39impl Callbacks {
40    pub fn new() -> Self {
41        Self {
42            before_model: Vec::new(),
43            after_model: Vec::new(),
44            before_tool: Vec::new(),
45            after_tool: Vec::new(),
46        }
47    }
48
49    pub fn add_before_model(&mut self, callback: BeforeModelCallback) {
50        self.before_model.push(callback);
51    }
52
53    pub fn add_after_model(&mut self, callback: AfterModelCallback) {
54        self.after_model.push(callback);
55    }
56
57    pub fn add_before_tool(&mut self, callback: BeforeToolCallback) {
58        self.before_tool.push(callback);
59    }
60
61    pub fn add_after_tool(&mut self, callback: AfterToolCallback) {
62        self.after_tool.push(callback);
63    }
64
65    /// Execute all before_model callbacks
66    pub async fn execute_before_model(
67        &self,
68        ctx: Arc<dyn CallbackContext>,
69    ) -> Result<Vec<Content>> {
70        let mut results = Vec::new();
71        for callback in &self.before_model {
72            if let Some(content) = callback(ctx.clone()).await? {
73                results.push(content);
74            }
75        }
76        Ok(results)
77    }
78
79    /// Execute all after_model callbacks
80    pub async fn execute_after_model(&self, ctx: Arc<dyn CallbackContext>) -> Result<Vec<Content>> {
81        let mut results = Vec::new();
82        for callback in &self.after_model {
83            if let Some(content) = callback(ctx.clone()).await? {
84                results.push(content);
85            }
86        }
87        Ok(results)
88    }
89
90    /// Execute all before_tool callbacks
91    pub async fn execute_before_tool(&self, ctx: Arc<dyn CallbackContext>) -> Result<Vec<Content>> {
92        let mut results = Vec::new();
93        for callback in &self.before_tool {
94            if let Some(content) = callback(ctx.clone()).await? {
95                results.push(content);
96            }
97        }
98        Ok(results)
99    }
100
101    /// Execute all after_tool callbacks
102    pub async fn execute_after_tool(&self, ctx: Arc<dyn CallbackContext>) -> Result<Vec<Content>> {
103        let mut results = Vec::new();
104        for callback in &self.after_tool {
105            if let Some(content) = callback(ctx.clone()).await? {
106                results.push(content);
107            }
108        }
109        Ok(results)
110    }
111}
adk_runner/callbacks.rs

adk_runner/
callbacks.rs
