256 lines
8.8 KiB
Python
256 lines
8.8 KiB
Python
"""
|
||
Google ADK 回调机制完整示例
|
||
展示四种回调函数的使用方法
|
||
|
||
对应教程：第07章 - 回调机制与事件系统
|
||
"""
|
||
|
||
# 导入 ADK 核心模块
|
||
from google.adk.agents import Agent # Agent 类
|
||
from google.adk.agents.invocation_context import InvocationContext # 调用上下文
|
||
from google.adk.runners import Runner # 运行器
|
||
from google.adk.sessions import InMemorySessionService # 会话服务
|
||
from google.genai import types # 类型定义
|
||
|
||
# 导入异步和时间模块
|
||
import asyncio # 异步编程库
|
||
import time # 时间模块
|
||
|
||
|
||
# ========================================
|
||
# 示例一：日志回调（记录所有操作）
|
||
# ========================================
|
||
|
||
async def log_before_model(cb_ctx, inv_ctx):
|
||
"""
|
||
模型调用前回调
|
||
|
||
Args:
|
||
cb_ctx: 回调上下文
|
||
inv_ctx: 调用上下文
|
||
"""
|
||
# 记录开始时间
|
||
inv_ctx.session.state["temp:model_start_time"] = time.time() # 保存时间戳
|
||
|
||
# 获取调用计数
|
||
count = inv_ctx.session.state.get("model_call_count", 0) # 读取计数
|
||
inv_ctx.session.state["model_call_count"] = count + 1 # 递增
|
||
|
||
print(f"🔍 [模型前] Agent '{inv_ctx.agent.name}' 即将调用 LLM (第{count+1}次)") # 打印日志
|
||
|
||
|
||
async def log_after_model(cb_ctx, inv_ctx):
|
||
"""
|
||
模型调用后回调
|
||
|
||
Args:
|
||
cb_ctx: 回调上下文
|
||
inv_ctx: 调用上下文
|
||
"""
|
||
# 计算耗时
|
||
start = inv_ctx.session.state.get("temp:model_start_time", time.time()) # 开始时间
|
||
elapsed = time.time() - start # 计算耗时
|
||
|
||
response = cb_ctx.response # 获取模型响应
|
||
|
||
if response and response.function_calls: # 如果有工具调用
|
||
tools = [fc.name for fc in response.function_calls] # 工具名列表
|
||
print(f"✅ [模型后] LLM 调用完成 ({elapsed:.2f}s)，决定调用工具: {tools}") # 打印
|
||
else: # 如果直接响应
|
||
print(f"✅ [模型后] LLM 调用完成 ({elapsed:.2f}s)，直接响应") # 打印
|
||
|
||
|
||
async def log_before_tool(cb_ctx, inv_ctx):
|
||
"""
|
||
工具调用前回调
|
||
|
||
Args:
|
||
cb_ctx: 回调上下文
|
||
inv_ctx: 调用上下文
|
||
"""
|
||
fc = cb_ctx.function_call # 获取函数调用
|
||
print(f"🔧 [工具前] 调用工具: {fc.name}({fc.args})") # 打印工具信息
|
||
|
||
|
||
async def log_after_tool(cb_ctx, inv_ctx):
|
||
"""
|
||
工具调用后回调
|
||
|
||
Args:
|
||
cb_ctx: 回调上下文
|
||
inv_ctx: 调用上下文
|
||
"""
|
||
fc = cb_ctx.function_call # 获取函数调用
|
||
result = cb_ctx.tool_result # 获取工具结果
|
||
|
||
status = "成功" # 默认成功
|
||
if result and result.get("status") == "error": # 如果错误
|
||
status = "失败" # 标记失败
|
||
|
||
print(f"📊 [工具后] {fc.name} 执行{status}") # 打印结果
|
||
|
||
|
||
# ========================================
|
||
# 示例二：人工确认回调
|
||
# ========================================
|
||
|
||
async def human_confirmation(cb_ctx, inv_ctx):
|
||
"""
|
||
人工确认回调
|
||
在敏感操作前暂停
|
||
|
||
Args:
|
||
cb_ctx: 回调上下文
|
||
inv_ctx: 调用上下文
|
||
"""
|
||
tool_name = cb_ctx.function_call.name # 获取工具名
|
||
tool_args = cb_ctx.function_call.args # 获取工具参数
|
||
|
||
# 定义敏感操作
|
||
sensitive_ops = { # 敏感操作映射
|
||
"delete_file": "删除文件", # 删除文件
|
||
"send_email": "发送邮件", # 发送邮件
|
||
"make_payment": "发起支付", # 发起支付
|
||
}
|
||
|
||
if tool_name in sensitive_ops: # 如果是敏感操作
|
||
op = sensitive_ops[tool_name] # 获取操作描述
|
||
print(f"\n⚠️ [人工确认] 需要人工确认！") # 打印警告
|
||
print(f" 操作: {op}") # 打印操作
|
||
print(f" 参数: {tool_args}") # 打印参数
|
||
print(f" 状态: 已记录（模拟自动通过）") # 模拟通过
|
||
|
||
|
||
# ========================================
|
||
# 示例三：参数验证回调
|
||
# ========================================
|
||
|
||
async def validate_params(cb_ctx, inv_ctx):
|
||
"""
|
||
参数验证回调
|
||
在工具调用前验证参数
|
||
|
||
Args:
|
||
cb_ctx: 回调上下文
|
||
inv_ctx: 调用上下文
|
||
"""
|
||
tool_name = cb_ctx.function_call.name # 工具名
|
||
tool_args = cb_ctx.function_call.args # 工具参数
|
||
|
||
# 验证搜索查询长度
|
||
if tool_name == "search": # 如果是搜索工具
|
||
query = tool_args.get("query", "") # 获取查询
|
||
if len(query) < 2: # 如果太短
|
||
print(f"⚠️ [验证] 搜索查询过短: '{query}'") # 打印警告
|
||
|
||
# 验证数值范围
|
||
if tool_name == "calculate": # 如果是计算工具
|
||
value = tool_args.get("value", 0) # 获取数值
|
||
if value < 0: # 如果为负数
|
||
print(f"⚠️ [验证] 数值不能为负: {value}") # 打印警告
|
||
|
||
|
||
# ========================================
|
||
# 定义工具函数
|
||
# ========================================
|
||
|
||
def get_weather(city: str) -> dict:
|
||
"""获取天气信息"""
|
||
weather_data = { # 天气数据
|
||
"北京": {"temp": "25°C", "condition": "晴天"},
|
||
"上海": {"temp": "28°C", "condition": "多云"},
|
||
}
|
||
data = weather_data.get(city) # 查找数据
|
||
if not data: # 如果找不到
|
||
return {"status": "error", "error_message": f"未找到'{city}'的天气信息"}
|
||
return {"status": "success", "city": city, **data} # 返回天气
|
||
|
||
|
||
def delete_file(filename: str) -> dict:
|
||
"""删除文件（模拟）"""
|
||
print(f"🗑️ 执行删除: {filename}") # 模拟删除
|
||
return {"status": "success", "message": f"文件 '{filename}' 已删除"}
|
||
|
||
|
||
# ========================================
|
||
# 创建带回调的 Agent
|
||
# ========================================
|
||
|
||
monitored_agent = Agent(
|
||
name="monitored_agent", # Agent 名称
|
||
model="gemini-2.0-flash", # 模型
|
||
instruction=( # 系统指令
|
||
"你是一个天气助手。\n"
|
||
"使用 get_weather 工具查询天气。\n"
|
||
"使用 delete_file 工具删除文件（需要确认）。"
|
||
),
|
||
tools=[get_weather, delete_file], # 注册工具
|
||
before_model_callback=log_before_model, # 模型前回调
|
||
after_model_callback=log_after_model, # 模型后回调
|
||
before_tool_callback=log_before_tool, # 工具前回调
|
||
after_tool_callback=log_after_tool, # 工具后回调
|
||
)
|
||
|
||
|
||
# ========================================
|
||
# 运行演示
|
||
# ========================================
|
||
|
||
APP_NAME = "callback_demo" # 应用名称
|
||
USER_ID = "user_001" # 用户 ID
|
||
SESSION_ID = "session_001" # 会话 ID
|
||
|
||
|
||
async def main():
|
||
"""主函数"""
|
||
|
||
print("=" * 60) # 分隔线
|
||
print("Google ADK 回调机制演示") # 标题
|
||
print("=" * 60) # 分隔线
|
||
|
||
# 初始化
|
||
session_service = InMemorySessionService() # 会话服务
|
||
await session_service.create_session( # 创建会话
|
||
app_name=APP_NAME, # 应用名称
|
||
user_id=USER_ID, # 用户 ID
|
||
session_id=SESSION_ID, # 会话 ID
|
||
)
|
||
|
||
# 创建 Runner
|
||
runner = Runner(
|
||
agent=monitored_agent, # Agent
|
||
app_name=APP_NAME, # 应用名称
|
||
session_service=session_service, # 会话服务
|
||
)
|
||
|
||
# 测试问题
|
||
queries = [ # 测试列表
|
||
"北京天气怎么样？", # 天气查询
|
||
]
|
||
|
||
for query in queries: # 遍历测试
|
||
print(f"\n{'='*60}") # 分隔线
|
||
print(f"[用户]: {query}") # 打印用户输入
|
||
|
||
# 构造消息
|
||
content = types.Content(
|
||
role='user', # 角色
|
||
parts=[types.Part(text=query)], # 内容
|
||
)
|
||
|
||
# 运行 Agent
|
||
events = runner.run_async(
|
||
user_id=USER_ID, # 用户 ID
|
||
session_id=SESSION_ID, # 会话 ID
|
||
new_message=content, # 消息
|
||
)
|
||
|
||
# 处理事件
|
||
async for event in events: # 遍历事件
|
||
if event.is_final_response(): # 最终响应
|
||
print(f"[Agent]: {event.content.parts[0].text}") # 打印回复
|
||
|
||
|
||
if __name__ == "__main__": # 直接运行
|
||
asyncio.run(main()) # 执行主函数
|