回调

adk-rust 中的回调提供了在关键执行点观察、自定义和控制 Agent 行为的钩子。它们支持日志记录、防护、缓存、响应修改等功能。

概览

adk-rust 支持六种回调类型,它们拦截 Agent 执行的不同阶段:

回调类型执行时机用例
before_agent在 Agent 开始处理之前输入验证、日志记录、提前终止
after_agent在 Agent 完成之后响应修改、日志记录、清理
before_model在 LLM 调用之前请求修改、缓存、速率限制
after_model在 LLM 响应之后响应过滤、日志记录、缓存
before_tool在 Tool 执行之前权限检查、参数验证
after_tool在 Tool 执行之后结果修改、日志记录

回调类型

Agent 回调

Agent 回调包装了整个 Agent 执行周期。

use adk_rust::prelude::*;
use std::sync::Arc;

// BeforeAgentCallback 类型签名
type BeforeAgentCallback = Box<
    dyn Fn(Arc<dyn CallbackContext>) 
        -> Pin<Box<dyn Future<Output = Result<Option<Content>>> + Send>> 
    + Send + Sync
>;

// AfterAgentCallback 类型签名  
type AfterAgentCallback = Box<
    dyn Fn(Arc<dyn CallbackContext>) 
        -> Pin<Box<dyn Future<Output = Result<Option<Content>>> + Send>> 
    + Send + Sync
>;

Model 回调

Model 回调拦截 LLM 请求和响应。

use adk_rust::prelude::*;
use std::sync::Arc;

// BeforeModelResult - 控制回调后的行为
pub enum BeforeModelResult {
    Continue(LlmRequest),  // 继续使用(可能已修改的)请求
    Skip(LlmResponse),     // 跳过模型调用,转而使用此响应
}

// BeforeModelCallback - 可以修改请求或跳过模型调用
type BeforeModelCallback = Box<
    dyn Fn(Arc<dyn CallbackContext>, LlmRequest)
        -> Pin<Box<dyn Future<Output = Result<BeforeModelResult>> + Send>>
    + Send + Sync
>;

// AfterModelCallback - 可以修改响应
type AfterModelCallback = Box<
    dyn Fn(Arc<dyn CallbackContext>, LlmResponse)
        -> Pin<Box<dyn Future<Output = Result<Option<LlmResponse>>> + Send>>
    + Send + Sync
>;

Tool 回调

Tool 回调拦截 Tool 执行。

use adk_rust::prelude::*;
use std::sync::Arc;

// BeforeToolCallback - 可以通过返回 Some(Content) 来跳过 Tool
type BeforeToolCallback = Box<
    dyn Fn(Arc<dyn CallbackContext>) 
        -> Pin<Box<dyn Future<Output = Result<Option<Content>>> + Send>> 
    + Send + Sync
>;

// AfterToolCallback - 可以修改 Tool 结果
type AfterToolCallback = Box<
    dyn Fn(Arc<dyn CallbackContext>) 
        -> Pin<Box<dyn Future<Output = Result<Option<Content>>> + Send>> 
    + Send + Sync
>;

返回值语义

回调使用不同的返回值来控制执行流程:

Agent/Tool 回调

返回值效果
Ok(None)继续正常执行
Ok(Some(content))使用提供的内容覆盖/跳过
Err(e)中止执行并报错

Model 回调

BeforeModelCallback 使用 BeforeModelResult

返回值效果
Ok(BeforeModelResult::Continue(request))继续执行(可能已修改的)请求
Ok(BeforeModelResult::Skip(response))跳过 Model 调用,改用此响应
Err(e)中止执行并报错

AfterModelCallback 使用 Option<LlmResponse>

返回值效果
Ok(None)保留原始响应
Ok(Some(response))替换为修改后的响应
Err(e)中止执行并报错

总结

  • Before agent/tool 回调:返回 None 继续,返回 Some(content) 跳过
  • Before model 回调:返回 Continue(request) 继续,返回 Skip(response) 绕过 Model
  • After 回调:返回 None 保留原始,返回 Some(...) 替换

向 Agent 添加回调

使用 LlmAgentBuilder 向 Agent 添加回调:

use adk_rust::prelude::*;
use std::sync::Arc;

#[tokio::main]
async fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
    let api_key = std::env::var("GOOGLE_API_KEY")?;
    let model = Arc::new(GeminiModel::new(&api_key, "gemini-2.5-flash")?);

    let agent = LlmAgentBuilder::new("my_agent")
        .model(model)
        .instruction("You are a helpful assistant.")
        // 添加 before_agent 回调
        .before_callback(Box::new(|ctx| {
            Box::pin(async move {
                println!("Agent starting: {}", ctx.agent_name());
                Ok(None) // 继续执行
            })
        }))
        // 添加 after_agent 回调
        .after_callback(Box::new(|ctx| {
            Box::pin(async move {
                println!("Agent completed: {}", ctx.agent_name());
                Ok(None) // 保留原始结果
            })
        }))
        .build()?;

    Ok(())
}

CallbackContext 接口

CallbackContext trait 提供对执行上下文的访问:

use adk_rust::prelude::*;

#[async_trait]
pub trait CallbackContext: ReadonlyContext {
    /// 访问 artifact 存储(如果已配置)
    fn artifacts(&self) -> Option<Arc<dyn Artifacts>>;
}

// CallbackContext 扩展 ReadonlyContext
#[async_trait]
pub trait ReadonlyContext: Send + Sync {
    /// 当前调用 ID
    fn invocation_id(&self) -> &str;
    
    /// 当前 Agent 的名称
    fn agent_name(&self) -> &str;
    
    /// 会话中的用户 ID
    fn user_id(&self) -> &str;
    
    /// 应用程序名称
    fn app_name(&self) -> &str;
    
    /// 会话 ID
    fn session_id(&self) -> &str;
    
    /// 当前分支(用于多 Agent)
    fn branch(&self) -> &str;
    
    /// 用户的输入内容
    fn user_content(&self) -> &Content;
}

常见模式

日志回调

记录所有 agent 交互:

use adk_rust::prelude::*;
use std::sync::Arc;

let agent = LlmAgentBuilder::new("logged_agent")
    .model(model)
    .before_callback(Box::new(|ctx| {
        Box::pin(async move {
            println!("[LOG] Agent '{}' starting", ctx.agent_name());
            println!("[LOG] Session: {}", ctx.session_id());
            println!("[LOG] User: {}", ctx.user_id());
            Ok(None)
        })
    }))
    .after_callback(Box::new(|ctx| {
        Box::pin(async move {
            println!("[LOG] Agent '{}' completed", ctx.agent_name());
            Ok(None)
        })
    }))
    .build()?;

输入护栏

在处理之前阻止不适当的内容:

use adk_rust::prelude::*;
use std::sync::Arc;

let agent = LlmAgentBuilder::new("guarded_agent")
    .model(model)
    .before_callback(Box::new(|ctx| {
        Box::pin(async move {
            // Check user input for blocked content
            let user_content = ctx.user_content();
            for part in &user_content.parts {
                if let Part::Text { text } = part {
                    if text.to_lowercase().contains("blocked_word") {
                        // Return early with rejection message
                        return Ok(Some(Content {
                            role: "model".to_string(),
                            parts: vec![Part::Text {
                                text: "I cannot process that request.".to_string(),
                            }],
                        }));
                    }
                }
            }
            Ok(None) // Continue normal execution
        })
    }))
    .build()?;

响应缓存 (在 Model 之前)

缓存 LLM 响应以减少 API 调用:

use adk_rust::prelude::*;
use std::sync::Arc;
use std::collections::HashMap;
use std::sync::Mutex;

// Simple in-memory cache
let cache: Arc<Mutex<HashMap<String, LlmResponse>>> = Arc::new(Mutex::new(HashMap::new()));
let cache_clone = cache.clone();

let agent = LlmAgentBuilder::new("cached_agent")
    .model(model)
    .before_model_callback(Box::new(move |ctx, request| {
        let cache = cache_clone.clone();
        Box::pin(async move {
            // Create cache key from request contents
            let key = format!("{:?}", request.contents);

            // Check cache
            if let Some(cached) = cache.lock().unwrap().get(&key) {
                println!("[CACHE] Hit for request");
                return Ok(BeforeModelResult::Skip(cached.clone()));
            }

            println!("[CACHE] Miss, calling model");
            Ok(BeforeModelResult::Continue(request)) // Continue to model
        })
    }))
    .build()?;

注入多模态内容 (在 Model 之前)

将图像或其他二进制内容注入 LLM 请求中进行多模态分析:

use adk_rust::prelude::*;
use adk_rust::artifact::{ArtifactService, LoadRequest};
use std::sync::Arc;

// Artifact service with pre-loaded image
let artifact_service: Arc<dyn ArtifactService> = /* ... */;
let callback_service = artifact_service.clone();

let agent = LlmAgentBuilder::new("image_analyst")
    .model(model)
    .instruction("Describe the image provided by the user.")
    .before_model_callback(Box::new(move |_ctx, mut request| {
        let service = callback_service.clone();
        Box::pin(async move {
            // Load image from artifact storage
            if let Ok(response) = service.load(LoadRequest {
                app_name: "my_app".to_string(),
                user_id: "user".to_string(),
                session_id: "session".to_string(),
                file_name: "user:photo.png".to_string(),
                version: None,
            }).await {
                // Inject image into the user's message
                if let Some(last_content) = request.contents.last_mut() {
                    if last_content.role == "user" {
                        last_content.parts.push(response.part);
                    }
                }
            }

            Ok(BeforeModelResult::Continue(request))
        })
    }))
    .build()?;

这种模式对于多模态 AI 至关重要,因为 Tool 响应是 JSON 文本——Model 无法“看到”Tool 返回的图像。通过将图像直接注入请求中,Model 可以接收到实际的图像数据。

响应修改 (在 Model 之后)

修改或过滤 Model 响应:

use adk_rust::prelude::*;
use std::sync::Arc;

let agent = LlmAgentBuilder::new("filtered_agent")
    .model(model)
    .after_model_callback(Box::new(|ctx, mut response| {
        Box::pin(async move {
            // Modify the response content
            if let Some(ref mut content) = response.content {
                for part in &mut content.parts {
                    if let Part::Text { text } = part {
                        // Add disclaimer to all responses
                        *text = format!("{}\n\n[AI-generated response]", text);
                    }
                }
            }
            Ok(Some(response))
        })
    }))
    .build()?;

Tool 权限检查 (在 Tool 之前)

验证 Tool 执行权限:

use adk_rust::prelude::*;
use std::sync::Arc;

let agent = LlmAgentBuilder::new("permission_agent")
    .model(model)
    .tool(Arc::new(GoogleSearchTool::new()))
    .before_tool_callback(Box::new(|ctx| {
        Box::pin(async move {
            // Check if user has permission for tools
            let user_id = ctx.user_id();
            
            // Example: block certain users from using tools
            if user_id == "restricted_user" {
                return Ok(Some(Content {
                    role: "function".to_string(),
                    parts: vec![Part::Text {
                        text: "Tool access denied for this user.".to_string(),
                    }],
                }));
            }
            
            Ok(None) // Allow tool execution
        })
    }))
    .build()?;

Tool 结果日志 (在 Tool 之后)

记录所有 Tool 执行:

use adk_rust::prelude::*;
use std::sync::Arc;

let agent = LlmAgentBuilder::new("tool_logged_agent")
    .model(model)
    .tool(Arc::new(GoogleSearchTool::new()))
    .after_tool_callback(Box::new(|ctx| {
        Box::pin(async move {
            println!("[TOOL LOG] Tool executed for agent: {}", ctx.agent_name());
            println!("[TOOL LOG] Session: {}", ctx.session_id());
            Ok(None) // Keep original result
        })
    }))
    .build()?;

多个回调

您可以添加多个相同类型的回调。它们按顺序执行:

use adk_rust::prelude::*;
use std::sync::Arc;

let agent = LlmAgentBuilder::new("multi_callback_agent")
    .model(model)
    // First before callback - logging
    .before_callback(Box::new(|ctx| {
        Box::pin(async move {
            println!("[1] Logging callback");
            Ok(None)
        })
    }))
    // Second before callback - validation
    .before_callback(Box::new(|ctx| {
        Box::pin(async move {
            println!("[2] Validation callback");
            Ok(None)
        })
    }))
    .build()?;

当回调返回 Some(content) 时,相同类型的后续回调将被跳过。

错误处理

回调可以返回错误以中止执行:

use adk_rust::prelude::*;
use std::sync::Arc;

let agent = LlmAgentBuilder::new("error_handling_agent")
    .model(model)
    .before_callback(Box::new(|ctx| {
        Box::pin(async move {
            // Validate something critical
            if ctx.user_id().is_empty() {
                return Err(AdkError::Agent("User ID is required".to_string()));
            }
            Ok(None)
        })
    }))
    .build()?;

最佳实践

  1. 保持回调轻量: 避免在回调中进行大量计算
  2. 优雅地处理错误: 返回有意义的错误消息
  3. 谨慎使用日志: 过多的日志记录可能会影响性能
  4. 明智地使用缓存: 考虑缓存失效策略
  5. 独立测试回调: 分别对回调逻辑进行单元测试

相关


上一页: ← 状态管理 | 下一页: 工件 →