init
This commit is contained in:
Nelson
318
ADKLearning/google-adk-tutorial/code/session_state_demo.py
Normal file
318
ADKLearning/google-adk-tutorial/code/session_state_demo.py
Normal file
@@ -0,0 +1,318 @@
|
||||
"""
|
||||
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()) # 执行主函数
|
||||
Reference in New Issue
Block a user