adk-runner 0.5.0

Agent execution runtime for Rust Agent Development Kit (ADK-Rust) agents
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

use adk_core::{CallbackContext, Content, Result};
use futures::future::BoxFuture;
use std::sync::Arc;

/// Callback executed before calling the model
pub type BeforeModelCallback = Box<
    dyn Fn(Arc<dyn CallbackContext>) -> BoxFuture<'static, Result<Option<Content>>> + Send + Sync,
>;

/// Callback executed after model response
pub type AfterModelCallback = Box<
    dyn Fn(Arc<dyn CallbackContext>) -> BoxFuture<'static, Result<Option<Content>>> + Send + Sync,
>;

/// Callback executed before tool execution
pub type BeforeToolCallback = Box<
    dyn Fn(Arc<dyn CallbackContext>) -> BoxFuture<'static, Result<Option<Content>>> + Send + Sync,
>;

/// Callback executed after tool execution
pub type AfterToolCallback = Box<
    dyn Fn(Arc<dyn CallbackContext>) -> BoxFuture<'static, Result<Option<Content>>> + Send + Sync,
>;

/// Collection of all callback types
pub struct Callbacks {
    pub before_model: Vec<BeforeModelCallback>,
    pub after_model: Vec<AfterModelCallback>,
    pub before_tool: Vec<BeforeToolCallback>,
    pub after_tool: Vec<AfterToolCallback>,
}

impl Default for Callbacks {
    fn default() -> Self {
        Self::new()
    }
}

impl Callbacks {
    pub fn new() -> Self {
        Self {
            before_model: Vec::new(),
            after_model: Vec::new(),
            before_tool: Vec::new(),
            after_tool: Vec::new(),
        }
    }

    pub fn add_before_model(&mut self, callback: BeforeModelCallback) {
        self.before_model.push(callback);
    }

    pub fn add_after_model(&mut self, callback: AfterModelCallback) {
        self.after_model.push(callback);
    }

    pub fn add_before_tool(&mut self, callback: BeforeToolCallback) {
        self.before_tool.push(callback);
    }

    pub fn add_after_tool(&mut self, callback: AfterToolCallback) {
        self.after_tool.push(callback);
    }

    /// Execute all before_model callbacks
    pub async fn execute_before_model(
        &self,
        ctx: Arc<dyn CallbackContext>,
    ) -> Result<Vec<Content>> {
        let mut results = Vec::new();
        for callback in &self.before_model {
            if let Some(content) = callback(ctx.clone()).await? {
                results.push(content);
            }
        }
        Ok(results)
    }

    /// Execute all after_model callbacks
    pub async fn execute_after_model(&self, ctx: Arc<dyn CallbackContext>) -> Result<Vec<Content>> {
        let mut results = Vec::new();
        for callback in &self.after_model {
            if let Some(content) = callback(ctx.clone()).await? {
                results.push(content);
            }
        }
        Ok(results)
    }

    /// Execute all before_tool callbacks
    pub async fn execute_before_tool(&self, ctx: Arc<dyn CallbackContext>) -> Result<Vec<Content>> {
        let mut results = Vec::new();
        for callback in &self.before_tool {
            if let Some(content) = callback(ctx.clone()).await? {
                results.push(content);
            }
        }
        Ok(results)
    }

    /// Execute all after_tool callbacks
    pub async fn execute_after_tool(&self, ctx: Arc<dyn CallbackContext>) -> Result<Vec<Content>> {
        let mut results = Vec::new();
        for callback in &self.after_tool {
            if let Some(content) = callback(ctx.clone()).await? {
                results.push(content);
            }
        }
        Ok(results)
    }
}