Context 与状态管理

为什么需要 Context

工具函数是独立的 Python 函数，它们之间默认不共享数据。但在实际应用中，工具之间需要共享信息：用户身份、数据库连接、已查询到的数据……Context 提供了一个类型安全的"共享容器"，通过依赖注入传递给每个工具和 Agent。

基础 Context 使用

from agents import Agent, Runner, function_tool, RunContextWrapper
from dataclasses import dataclass, field
from typing import Optional
import asyncio

# ── 定义 Context 数据类 ───────────────────────────────────
@dataclass
class ShoppingContext:
    # 用户信息（来自认证层）
    user_id: str
    user_name: str
    vip_level: int = 0           # 0=普通, 1=VIP, 2=SVIP

    # 当前 Run 的业务状态
    cart_items: list[dict] = field(default_factory=list)
    viewed_products: list[str] = field(default_factory=list)
    discount_applied: bool = False

    # 依赖：数据库会话（注入）
    db_session: Optional[object] = None

# ── 工具中读写 Context ────────────────────────────────────
@function_tool
async def add_to_cart(
    ctx: RunContextWrapper[ShoppingContext],  # 第一个参数为 ctx
    product_id: str,
    quantity: int = 1
) -> str:
    """将商品加入购物车。"""
    # 通过 ctx.context 读写业务状态
    ctx.context.cart_items.append({
        "product_id": product_id,
        "quantity": quantity
    })
    return f"已加入购物车：{product_id} x{quantity}，当前共 {len(ctx.context.cart_items)} 件"

@function_tool
async def apply_discount(
    ctx: RunContextWrapper[ShoppingContext],
    coupon_code: str
) -> str:
    """应用优惠券。每次购物只能使用一次。"""
    if ctx.context.discount_applied:
        return "本次购物已应用过优惠券，无法重复使用。"

    # VIP 用户额外折扣
    discount_rate = 0.9 if ctx.context.vip_level == 0 else 0.85
    ctx.context.discount_applied = True
    return f"优惠券 {coupon_code} 已应用，享受 {discount_rate*100:.0f}% 折扣"

@function_tool
async def checkout(
    ctx: RunContextWrapper[ShoppingContext]
) -> str:
    """结算购物车。"""
    if not ctx.context.cart_items:
        return "购物车为空，请先添加商品。"

    total_items = sum(item["quantity"] for item in ctx.context.cart_items)
    discount_msg = "（含优惠券折扣）" if ctx.context.discount_applied else ""
    return f"订单创建成功！{total_items} 件商品{discount_msg}，用户：{ctx.context.user_name}"

# ── 创建购物助手 Agent ────────────────────────────────────
shopping_agent = Agent(
    name="购物助手",
    instructions="你是购物助手，帮用户管理购物车和结算。",
    model="gpt-4o-mini",
    tools=[add_to_cart, apply_discount, checkout]
)

async def main():
    ctx = ShoppingContext(
        user_id="usr_001",
        user_name="张三",
        vip_level=1  # VIP 用户
    )
    result = await Runner.run(
        shopping_agent,
        "帮我加入商品 P001 两件，用优惠码 SAVE20，然后结算",
        context=ctx
    )
    print(result.final_output)

asyncio.run(main())

数据库集成：Context 注入 DB Session

from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker
from agents import Agent, Runner, function_tool, RunContextWrapper
from dataclasses import dataclass

# ── 数据库设置 ─────────────────────────────────────────────
engine = create_async_engine("postgresql+asyncpg://user:pass@localhost/mydb")
AsyncSessionLocal = async_sessionmaker(engine)

@dataclass
class AppContext:
    user_id: str
    db: AsyncSession  # 数据库会话（每个 Run 独立，线程安全）

# ── 使用 DB Session 的工具 ────────────────────────────────
@function_tool
async def get_user_profile(ctx: RunContextWrapper[AppContext]) -> str:
    """获取当前用户的档案信息。"""
    # 直接使用 Context 中注入的 DB Session
    result = await ctx.context.db.execute(
        "SELECT name, email, created_at FROM users WHERE id = :id",
        {"id": ctx.context.user_id}
    )
    user = result.fetchone()
    if not user:
        return "未找到用户信息"
    return f"用户：{user.name}，邮箱：{user.email}，注册于：{user.created_at}"

# ── 在 FastAPI 路由中使用（每次请求独立 DB Session）────────
from fastapi import FastAPI, Depends, HTTPException
app = FastAPI()

profile_agent = Agent(
    name="用户档案助手",
    instructions="帮助用户查询和更新个人档案信息。",
    model="gpt-4o-mini",
    tools=[get_user_profile]
)

@app.post("/agent/profile")
async def profile_endpoint(user_id: str, message: str):
    # 每次请求创建独立的 DB Session（线程安全）
    async with AsyncSessionLocal() as session:
        ctx = AppContext(user_id=user_id, db=session)
        result = await Runner.run(profile_agent, message, context=ctx)
        return {"reply": result.final_output}

带用户偏好记忆的个性化 Agent

from agents import Agent, Runner, function_tool, RunContextWrapper
from dataclasses import dataclass, field
import json, asyncio

@dataclass
class UserMemory:
    user_id: str
    preferences: dict = field(default_factory=dict)   # 持久化偏好
    session_notes: list[str] = field(default_factory=list)  # 本次会话笔记

@function_tool
async def save_preference(
    ctx: RunContextWrapper[UserMemory],
    key: str,
    value: str
) -> str:
    """保存用户偏好，例如语言、通知设置等。将在下次会话中记住。"""
    ctx.context.preferences[key] = value
    # 同时持久化到数据库
    await persist_preference(ctx.context.user_id, key, value)
    return f"已记住：{key} = {value}"

@function_tool
async def get_preferences(ctx: RunContextWrapper[UserMemory]) -> str:
    """获取当前用户的所有已保存偏好。"""
    if not ctx.context.preferences:
        return "尚未保存任何偏好设置。"
    return "当前偏好：\n" + json.dumps(ctx.context.preferences, ensure_ascii=False, indent=2)

def build_memory_instructions(
    ctx: RunContextWrapper[UserMemory],
    agent: Agent
) -> str:
    prefs = ctx.context.preferences
    pref_text = "\n".join([f"- {k}: {v}" for k, v in prefs.items()]) if prefs else "无"
    return f"""你是个性化 AI 助手，记住用户偏好并据此调整回答。

用户已知偏好：
{pref_text}

当用户提到新的偏好时（如"我喜欢简洁回答"），使用 save_preference 工具保存。
根据用户偏好调整回答风格和内容深度。
"""

memory_agent = Agent(
    name="记忆助手",
    instructions=build_memory_instructions,
    model="gpt-4o-mini",
    tools=[save_preference, get_preferences]
)

async def chat_with_memory(user_id: str):
    # 从数据库加载历史偏好
    saved_prefs = await load_user_preferences(user_id)
    ctx = UserMemory(user_id=user_id, preferences=saved_prefs)

    result = await Runner.run(
        memory_agent,
        "我喜欢简洁的回答，不超过3句话。然后告诉我 Python 3.13 有什么新特性？",
        context=ctx
    )
    print(result.final_output)