319 lines
9.3 KiB
Python
319 lines
9.3 KiB
Python
"""
|
||
Google ADK 会话与状态管理完整示例
|
||
展示 Session、State、Memory 的使用方法
|
||
|
||
对应教程：第06章 - 会话与状态管理
|
||
"""
|
||
|
||
# 导入 ADK 核心模块
|
||
from google.adk.agents import Agent # Agent 类
|
||
from google.adk.runners import Runner # 运行器
|
||
from google.adk.sessions import InMemorySessionService # 内存会话服务
|
||
from google.adk.memory import InMemoryMemoryService # 内存记忆服务
|
||
from google.genai import types # 类型定义
|
||
|
||
# 导入异步库
|
||
import asyncio # 异步编程库
|
||
|
||
|
||
# ========================================
|
||
# 示例一：State 管理（购物车）
|
||
# ========================================
|
||
|
||
def add_to_cart(item: str, price: float, ctx) -> dict:
|
||
"""
|
||
将商品添加到购物车
|
||
|
||
Args:
|
||
item (str): 商品名称
|
||
price (float): 商品价格
|
||
ctx: 工具上下文（自动注入）
|
||
|
||
Returns:
|
||
dict: 操作结果
|
||
"""
|
||
# 获取当前购物车
|
||
cart = ctx.state.get("cart", []) # 从 state 读取购物车
|
||
|
||
# 添加新商品
|
||
cart.append({ # 追加商品
|
||
"item": item, # 商品名
|
||
"price": price, # 价格
|
||
})
|
||
|
||
# 更新 state
|
||
ctx.state["cart"] = cart # 写回 state
|
||
|
||
# 计算总价
|
||
total = sum(item["price"] for item in cart) # 求和
|
||
|
||
return { # 返回结果
|
||
"status": "success",
|
||
"message": f"已将 {item} 添加到购物车",
|
||
"cart_items": len(cart),
|
||
"total": total,
|
||
}
|
||
|
||
|
||
def get_cart(ctx) -> dict:
|
||
"""
|
||
获取购物车内容
|
||
|
||
Args:
|
||
ctx: 工具上下文
|
||
|
||
Returns:
|
||
dict: 购物车内容
|
||
"""
|
||
cart = ctx.state.get("cart", []) # 读取购物车
|
||
|
||
if not cart: # 如果为空
|
||
return { # 返回空信息
|
||
"status": "success",
|
||
"message": "购物车是空的",
|
||
"items": [],
|
||
"total": 0,
|
||
}
|
||
|
||
total = sum(item["price"] for item in cart) # 计算总价
|
||
|
||
return { # 返回购物车
|
||
"status": "success",
|
||
"items": cart,
|
||
"total": total,
|
||
}
|
||
|
||
|
||
def clear_cart(ctx) -> dict:
|
||
"""
|
||
清空购物车
|
||
|
||
Args:
|
||
ctx: 工具上下文
|
||
|
||
Returns:
|
||
dict: 操作结果
|
||
"""
|
||
ctx.state["cart"] = [] # 清空购物车
|
||
|
||
return { # 返回结果
|
||
"status": "success",
|
||
"message": "购物车已清空",
|
||
}
|
||
|
||
|
||
# ========================================
|
||
# 示例二：temp: 临时状态
|
||
# ========================================
|
||
|
||
def step1_process(query: str, ctx) -> dict:
|
||
"""
|
||
处理步骤一：预处理数据
|
||
|
||
Args:
|
||
query (str): 用户查询
|
||
ctx: 工具上下文
|
||
|
||
Returns:
|
||
dict: 预处理结果
|
||
"""
|
||
processed = query.strip().lower() # 预处理
|
||
|
||
# 保存到临时状态（当前调用结束后自动清除）
|
||
ctx.state["temp:processed_query"] = processed # 临时存储
|
||
|
||
return { # 返回结果
|
||
"status": "success",
|
||
"processed_query": processed,
|
||
}
|
||
|
||
|
||
def step2_enhance(ctx) -> dict:
|
||
"""
|
||
处理步骤二：增强数据
|
||
|
||
Args:
|
||
ctx: 工具上下文
|
||
|
||
Returns:
|
||
dict: 增强结果
|
||
"""
|
||
# 从临时状态读取
|
||
processed = ctx.state.get("temp:processed_query") # 读取临时状态
|
||
|
||
if not processed: # 如果没有数据
|
||
return { # 返回错误
|
||
"status": "error",
|
||
"error_message": "请先执行预处理步骤。"
|
||
}
|
||
|
||
enhanced = f"[增强] {processed}" # 增强处理
|
||
|
||
return { # 返回结果
|
||
"status": "success",
|
||
"enhanced_query": enhanced,
|
||
}
|
||
|
||
|
||
# ========================================
|
||
# 创建 Agent
|
||
# ========================================
|
||
|
||
shopping_agent = Agent(
|
||
name="shopping_assistant", # Agent 名称
|
||
model="gemini-2.0-flash", # 模型
|
||
instruction=( # 指令
|
||
"你是一个购物助手。\n"
|
||
"帮助用户管理购物车：\n"
|
||
"- 添加商品：使用 add_to_cart\n"
|
||
"- 查看购物车：使用 get_cart\n"
|
||
"- 清空购物车：使用 clear_cart"
|
||
),
|
||
tools=[add_to_cart, get_cart, clear_cart], # 注册工具
|
||
)
|
||
|
||
|
||
# ========================================
|
||
# 示例三：Session 管理
|
||
# ========================================
|
||
|
||
async def session_management_demo():
|
||
"""演示 Session 管理功能"""
|
||
|
||
print("=" * 50) # 分隔线
|
||
print("Session 管理演示") # 标题
|
||
print("=" * 50) # 分隔线
|
||
|
||
# 创建会话服务
|
||
session_service = InMemorySessionService() # 内存会话服务
|
||
|
||
# 创建会话
|
||
session = await session_service.create_session(
|
||
app_name="shopping_app", # 应用名称
|
||
user_id="user_001", # 用户 ID
|
||
session_id="session_001", # 会话 ID
|
||
state={"cart": []}, # 初始化状态
|
||
)
|
||
|
||
print(f"✅ 会话创建成功: {session.id}") # 打印会话 ID
|
||
|
||
# 获取会话
|
||
existing = await session_service.get_session(
|
||
app_name="shopping_app",
|
||
user_id="user_001",
|
||
session_id="session_001",
|
||
)
|
||
print(f"✅ 获取会话: {existing.id}") # 打印会话信息
|
||
|
||
# 删除会话
|
||
await session_service.delete_session(
|
||
app_name="shopping_app",
|
||
user_id="user_001",
|
||
session_id="session_001",
|
||
)
|
||
print("✅ 会话已删除") # 确认删除
|
||
|
||
|
||
# ========================================
|
||
# 示例四：多轮对话
|
||
# ========================================
|
||
|
||
async def multi_turn_demo():
|
||
"""演示多轮对话"""
|
||
|
||
print("\n" + "=" * 50) # 分隔线
|
||
print("多轮对话演示") # 标题
|
||
print("=" * 50) # 分隔线
|
||
|
||
# 初始化
|
||
session_service = InMemorySessionService() # 会话服务
|
||
await session_service.create_session( # 创建会话
|
||
app_name="chat_app", # 应用名称
|
||
user_id="user_001", # 用户 ID
|
||
session_id="chat_001", # 会话 ID
|
||
)
|
||
|
||
# 创建 Runner
|
||
runner = Runner(
|
||
agent=shopping_agent, # Agent
|
||
app_name="chat_app", # 应用名称
|
||
session_service=session_service, # 会话服务
|
||
)
|
||
|
||
# 模拟多轮对话
|
||
conversations = [ # 对话列表
|
||
"帮我添加一个苹果，价格5元", # 第一轮
|
||
"再添加一个香蕉，价格3元", # 第二轮
|
||
"查看我的购物车", # 第三轮
|
||
"清空购物车", # 第四轮
|
||
]
|
||
|
||
for msg in conversations: # 遍历对话
|
||
print(f"\n[用户]: {msg}") # 打印用户消息
|
||
|
||
# 构造消息
|
||
content = types.Content(
|
||
role='user', # 角色
|
||
parts=[types.Part(text=msg)], # 内容
|
||
)
|
||
|
||
# 运行 Agent
|
||
events = runner.run_async(
|
||
user_id="user_001", # 用户 ID
|
||
session_id="chat_001", # 同一会话
|
||
new_message=content, # 消息
|
||
)
|
||
|
||
# 获取响应
|
||
async for event in events: # 遍历事件
|
||
if event.is_final_response(): # 最终响应
|
||
print(f"[Agent]: {event.content.parts[0].text}") # 打印回复
|
||
|
||
|
||
# ========================================
|
||
# 示例五：Memory 管理
|
||
# ========================================
|
||
|
||
async def memory_demo():
|
||
"""演示 Memory 管理功能"""
|
||
|
||
print("\n" + "=" * 50) # 分隔线
|
||
print("Memory 管理演示") # 标题
|
||
print("=" * 50) # 分隔线
|
||
|
||
# 创建记忆服务
|
||
memory_service = InMemoryMemoryService() # 内存记忆服务
|
||
|
||
# 添加记忆（从会话事件中提取）
|
||
await memory_service.add_session_events_to_memory(
|
||
app_name="my_app", # 应用名称
|
||
user_id="user_001", # 用户 ID
|
||
session_id="session_001", # 会话 ID
|
||
)
|
||
print("✅ 记忆已添加") # 确认添加
|
||
|
||
# 搜索记忆
|
||
results = await memory_service.search(
|
||
query="用户偏好", # 搜索查询
|
||
app_name="my_app", # 应用名称
|
||
user_id="user_001", # 用户 ID
|
||
)
|
||
print(f"✅ 搜索完成，结果数: {len(results)}") # 打印结果数
|
||
|
||
|
||
# ========================================
|
||
# 主函数
|
||
# ========================================
|
||
|
||
async def main():
|
||
"""主函数"""
|
||
|
||
# 运行所有演示
|
||
await session_management_demo() # Session 管理
|
||
await multi_turn_demo() # 多轮对话
|
||
await memory_demo() # Memory 管理
|
||
|
||
|
||
if __name__ == "__main__": # 直接运行
|
||
asyncio.run(main()) # 执行主函数
|