init

ADKLearning/google-adk-tutorial/code/callback_demo.py

@@ -0,0 +1,255 @@
+"""
+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())                      # 执行主函数