💾 Checkpointer 持久化系统

"""
InMemorySaver 示例
功能：内存存储，适合开发测试
"""
from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver

# 创建带内存持久化的 Agent
agent = create_agent(
    model="gpt-4o",
    tools=[],
    checkpointer=InMemorySaver(),  # 内存存储
)

# 第一轮对话
response1 = agent.invoke(
    {"messages": [{"role": "user", "content": "你好！我叫张三。"}]},
    {"configurable": {"thread_id": "user_001"}}  # 指定 Thread ID
)
print(response1["messages"][-1].content)

# 第二轮对话（记住之前的内容）
response2 = agent.invoke(
    {"messages": [{"role": "user", "content": "我刚才说我叫什么？"}]},
    {"configurable": {"thread_id": "user_001"}}  # 相同 Thread ID
)
print(response2["messages"][-1].content)  # 输出："你说你叫张三。"

# 安装 PostgreSQL Checkpointer
pip install langgraph-checkpoint-postgres

"""
PostgresSaver 示例
功能：生产级数据库持久化
"""
from langchain.agents import create_agent
from langgraph.checkpoint.postgres import PostgresSaver

# PostgreSQL 连接字符串
DB_URI = "postgresql://username:password@localhost:5432/langchain_db"

# 使用上下文管理器
with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
    # 首次使用需要初始化表结构
    checkpointer.setup()

    # 创建 Agent
    agent = create_agent(
        model="gpt-4o",
        tools=[],
        checkpointer=checkpointer,
    )

    # 用户 A 的对话
    agent.invoke(
        {"messages": [{"role": "user", "content": "我喜欢蓝色"}]},
        {"configurable": {"thread_id": "user_A"}}
    )

    # 用户 B 的对话（完全隔离）
    agent.invoke(
        {"messages": [{"role": "user", "content": "我喜欢红色"}]},
        {"configurable": {"thread_id": "user_B"}}
    )

    # 用户 A 继续对话（记住之前的内容）
    response = agent.invoke(
        {"messages": [{"role": "user", "content": "我喜欢什么颜色？"}]},
        {"configurable": {"thread_id": "user_A"}}
    )
    print(response["messages"][-1].content)  # 输出："你喜欢蓝色。"

# 安装 SQLite Checkpointer
pip install langgraph-checkpoint-sqlite

"""
SqliteSaver 示例
功能：文件持久化，适合小型部署
"""
from langchain.agents import create_agent
from langgraph.checkpoint.sqlite import SqliteSaver

# 创建 SQLite Checkpointer（存储到文件）
with SqliteSaver.from_conn_string("checkpoints.db") as checkpointer:
    checkpointer.setup()

    agent = create_agent(
        model="gpt-4o",
        tools=[],
        checkpointer=checkpointer,
    )

    # 对话会持久化到 checkpoints.db 文件
    agent.invoke(
        {"messages": [{"role": "user", "content": "你好"}]},
        {"configurable": {"thread_id": "session_123"}}
    )

"""
Thread ID 使用示例
功能：演示会话隔离和 Thread ID 设计模式
"""
from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver
import uuid

agent = create_agent("gpt-4o", tools=[], checkpointer=InMemorySaver())

# 模式 1：用户 ID 作为 Thread ID（单用户多会话）
user_id = "user_12345"
session_id = "session_001"
thread_id = f"{user_id}_{session_id}"

agent.invoke(
    {"messages": [{"role": "user", "content": "帮我创建订单"}]},
    {"configurable": {"thread_id": thread_id}}
)

# 模式 2：随机生成 Thread ID（匿名会话）
thread_id = str(uuid.uuid4())  # 生成唯一 ID
agent.invoke(
    {"messages": [{"role": "user", "content": "你好"}]},
    {"configurable": {"thread_id": thread_id}}
)

# 模式 3：业务实体作为 Thread ID（订单对话）
order_id = "ORDER-2024-001"
thread_id = f"order_{order_id}"

agent.invoke(
    {"messages": [{"role": "user", "content": "查询订单状态"}]},
    {"configurable": {"thread_id": thread_id}}
)

# 模式 4：多租户隔离
tenant_id = "company_A"
user_id = "user_001"
thread_id = f"{tenant_id}_{user_id}"

agent.invoke(
    {"messages": [{"role": "user", "content": "帮我查询账户余额"}]},
    {"configurable": {"thread_id": thread_id}}
)

"""
自定义状态 Schema 示例
功能：在状态中存储用户信息和偏好设置
"""
from langchain.agents import create_agent, AgentState
from langgraph.checkpoint.memory import InMemorySaver
from langchain.tools import tool, ToolRuntime

# 定义自定义状态类型
class CustomAgentState(AgentState):
    user_id: str
    user_name: str
    preferences: dict
    order_history: list

# 定义访问状态的工具
@tool
def get_user_info(runtime: ToolRuntime[CustomAgentState]) -> str:
    """获取用户信息"""
    state = runtime.state
    user_name = state.get("user_name", "未知用户")
    preferences = state.get("preferences", {})

    return f"""
    用户姓名：{user_name}
    主题偏好：{preferences.get('theme', '未设置')}
    语言偏好：{preferences.get('language', '未设置')}
    """

@tool
def get_order_history(runtime: ToolRuntime[CustomAgentState]) -> str:
    """查询订单历史"""
    orders = runtime.state.get("order_history", [])
    if not orders:
        return "暂无订单记录"

    order_list = "\n".join([f"- {order}" for order in orders])
    return f"订单历史：\n{order_list}"

# 创建带自定义状态的 Agent
agent = create_agent(
    model="gpt-4o",
    tools=[get_user_info, get_order_history],
    state_schema=CustomAgentState,
    checkpointer=InMemorySaver(),
)

# 调用时提供完整状态
result = agent.invoke(
    {
        "messages": [{"role": "user", "content": "查询我的信息"}],
        "user_id": "user_123",
        "user_name": "张三",
        "preferences": {
            "theme": "dark",
            "language": "zh-CN"
        },
        "order_history": [
            "订单#001 - iPhone 15",
            "订单#002 - MacBook Pro"
        ]
    },
    {"configurable": {"thread_id": "user_123"}}
)

print(result["messages"][-1].content)

"""
消息修剪中间件
功能：保留最近的 N 条消息，删除旧消息
"""
from langchain.agents.middleware import before_model
from langchain.messages import RemoveMessage
from langgraph.graph.message import REMOVE_ALL_MESSAGES
from langchain.agents import AgentState, create_agent
from langgraph.checkpoint.memory import InMemorySaver

@before_model
def trim_messages(state: AgentState, runtime) -> dict | None:
    """修剪消息：保留第一条 + 最近 3 条"""
    messages = state["messages"]

    # 消息数量少于阈值，无需修剪
    if len(messages) <= 3:
        return None

    # 保留第一条消息（通常是 SystemMessage）
    first_msg = messages[0]

    # 保留最近 3 条消息
    recent = messages[-3:] if len(messages) % 2 == 0 else messages[-4:]

    # 构建修剪指令
    return {
        "messages": [
            RemoveMessage(id=REMOVE_ALL_MESSAGES),  # 删除所有消息
            first_msg,                              # 恢复第一条
            *recent                                 # 恢复最近的消息
        ]
    }

# 创建带消息修剪的 Agent
agent = create_agent(
    model="gpt-4o",
    tools=[],
    middleware=[trim_messages],
    checkpointer=InMemorySaver(),
)

# 模拟长对话
thread_id = "long_conversation"
for i in range(10):
    agent.invoke(
        {"messages": [{"role": "user", "content": f"消息 {i+1}"}]},
        {"configurable": {"thread_id": thread_id}}
    )

# 历史只会保留最近几条消息

"""
消息删除示例
功能：手动删除特定消息
"""
from langchain.messages import RemoveMessage
from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver

agent = create_agent("gpt-4o", tools=[], checkpointer=InMemorySaver())

# 第一轮对话
response1 = agent.invoke(
    {"messages": [{"role": "user", "content": "第一条消息"}]},
    {"configurable": {"thread_id": "test"}}
)

# 第二轮对话
response2 = agent.invoke(
    {"messages": [{"role": "user", "content": "第二条消息"}]},
    {"configurable": {"thread_id": "test"}}
)

# 获取当前状态
state = agent.get_state({"configurable": {"thread_id": "test"}})
messages = state.values["messages"]

# 删除前两条消息
if len(messages) > 2:
    delete_result = agent.update_state(
        {"configurable": {"thread_id": "test"}},
        {"messages": [RemoveMessage(id=m.id) for m in messages[:2]]}
    )

# 继续对话（前两条消息已被删除）
response3 = agent.invoke(
    {"messages": [{"role": "user", "content": "我之前说了什么？"}]},
    {"configurable": {"thread_id": "test"}}
)
print(response3["messages"][-1].content)  # 模型不会记得前两条

"""
消息摘要中间件
功能：自动总结旧消息，节省 Token
"""
from langchain.agents.middleware import SummarizationMiddleware
from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver

agent = create_agent(
    model="gpt-4o",
    tools=[],
    middleware=[
        SummarizationMiddleware(
            model="gpt-4o-mini",          # 使用轻量模型生成摘要
            trigger=("tokens", 4000),      # 超过 4000 tokens 触发摘要
            keep=("messages", 20),         # 保留最近 20 条消息
        )
    ],
    checkpointer=InMemorySaver(),
    system_prompt="你是一个有用的助手。"
)

# 长对话会自动触发摘要
thread_id = "summarization_test"
for i in range(30):
    agent.invoke(
        {"messages": [{"role": "user", "content": f"这是第 {i+1} 条消息，内容是..."}]},
        {"configurable": {"thread_id": thread_id}}
    )

"""
工具读取状态示例
功能：在工具中访问 Agent 状态
"""
from langchain.tools import tool, ToolRuntime
from langchain.agents import AgentState, create_agent
from langgraph.checkpoint.memory import InMemorySaver

class UserState(AgentState):
    user_id: str
    subscription_level: str

@tool
def get_user_quota(runtime: ToolRuntime[UserState]) -> str:
    """查询用户配额"""
    # 访问状态
    user_id = runtime.state.get("user_id", "unknown")
    level = runtime.state.get("subscription_level", "free")

    # 根据订阅等级返回配额
    quotas = {
        "free": 10,
        "basic": 100,
        "premium": 1000,
    }

    quota = quotas.get(level, 10)
    return f"用户 {user_id} 的配额：{quota} 次/月"

agent = create_agent(
    model="gpt-4o",
    tools=[get_user_quota],
    state_schema=UserState,
    checkpointer=InMemorySaver(),
)

result = agent.invoke(
    {
        "messages": [{"role": "user", "content": "查询我的配额"}],
        "user_id": "user_123",
        "subscription_level": "premium"
    },
    {"configurable": {"thread_id": "user_123"}}
)

print(result["messages"][-1].content)

"""
工具修改状态示例
功能：工具执行后更新状态
"""
from langchain.tools import tool, ToolRuntime
from langgraph.types import Command
from langchain.messages import ToolMessage
from langchain.agents import AgentState, create_agent
from langgraph.checkpoint.memory import InMemorySaver

class OrderState(AgentState):
    current_order_id: str | None
    order_total: float

@tool
def create_order(
    items: list[str],
    runtime: ToolRuntime[OrderState]
) -> Command:
    """创建订单并更新状态"""
    # 模拟订单创建
    order_id = "ORDER-2024-001"
    total = len(items) * 99.0  # 简化计算

    # 返回 Command 更新状态
    return Command(
        update={
            "current_order_id": order_id,
            "order_total": total,
            "messages": [
                ToolMessage(
                    content=f"订单创建成功！订单号：{order_id}，总金额：¥{total}",
                    tool_call_id=runtime.tool_call_id
                )
            ]
        }
    )

@tool
def get_current_order(runtime: ToolRuntime[OrderState]) -> str:
    """查询当前订单"""
    order_id = runtime.state.get("current_order_id")
    total = runtime.state.get("order_total", 0.0)

    if not order_id:
        return "暂无进行中的订单"

    return f"当前订单：{order_id}，金额：¥{total}"

agent = create_agent(
    model="gpt-4o",
    tools=[create_order, get_current_order],
    state_schema=OrderState,
    checkpointer=InMemorySaver(),
)

# 创建订单
result1 = agent.invoke(
    {
        "messages": [{"role": "user", "content": "帮我创建订单，商品：iPhone, AirPods"}],
        "current_order_id": None,
        "order_total": 0.0
    },
    {"configurable": {"thread_id": "order_session_001"}}
)

# 查询订单（状态已更新）
result2 = agent.invoke(
    {"messages": [{"role": "user", "content": "查询我的订单"}]},
    {"configurable": {"thread_id": "order_session_001"}}
)

print(result2["messages"][-1].content)

"""
多用户客服系统完整示例
功能：演示 Checkpointer 在实际应用中的使用
"""
from langchain.agents import create_agent, AgentState
from langchain.tools import tool, ToolRuntime
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command
from langchain.messages import ToolMessage
from datetime import datetime

# ==================== 自定义状态 ====================

class CustomerServiceState(AgentState):
    user_id: str
    user_name: str
    issue_category: str | None
    ticket_id: str | None
    satisfaction_score: int | None

# ==================== 工具定义 ====================

@tool
def create_ticket(
    category: str,
    description: str,
    runtime: ToolRuntime[CustomerServiceState]
) -> Command:
    """创建客服工单"""
    user_id = runtime.state.get("user_id", "unknown")
    ticket_id = f"TICKET-{datetime.now().strftime('%Y%m%d%H%M%S')}"

    return Command(
        update={
            "ticket_id": ticket_id,
            "issue_category": category,
            "messages": [
                ToolMessage(
                    content=f"工单创建成功！\n工单号：{ticket_id}\n类别：{category}\n问题：{description}",
                    tool_call_id=runtime.tool_call_id
                )
            ]
        }
    )

@tool
def get_ticket_status(runtime: ToolRuntime[CustomerServiceState]) -> str:
    """查询工单状态"""
    ticket_id = runtime.state.get("ticket_id")

    if not ticket_id:
        return "您当前没有进行中的工单"

    category = runtime.state.get("issue_category", "未分类")
    return f"工单号：{ticket_id}\n类别：{category}\n状态：处理中\n预计解决时间：24 小时内"

@tool
def submit_feedback(
    score: int,
    comment: str,
    runtime: ToolRuntime[CustomerServiceState]
) -> Command:
    """提交满意度评分（1-5 分）"""
    if not 1 <= score <= 5:
        return Command(
            update={
                "messages": [
                    ToolMessage(
                        content="评分必须在 1-5 分之间",
                        tool_call_id=runtime.tool_call_id
                    )
                ]
            }
        )

    return Command(
        update={
            "satisfaction_score": score,
            "messages": [
                ToolMessage(
                    content=f"感谢您的反馈！评分：{score} 分\n评价：{comment}",
                    tool_call_id=runtime.tool_call_id
                )
            ]
        }
    )

# ==================== 创建 Agent ====================

customer_service_agent = create_agent(
    model="gpt-4o",
    tools=[create_ticket, get_ticket_status, submit_feedback],
    state_schema=CustomerServiceState,
    checkpointer=InMemorySaver(),  # 生产环境使用 PostgresSaver
    system_prompt="""你是一个专业的客服助手。

职责：
1. 帮助用户创建和查询工单
2. 收集用户反馈和满意度评分
3. 提供友好、专业的服务

工作流程：
- 了解用户问题，使用 create_ticket 创建工单
- 用户查询时，使用 get_ticket_status 查看进度
- 问题解决后，引导用户使用 submit_feedback 评分

注意：
- 始终保持礼貌和耐心
- 记住用户的姓名和问题类别
- 每个用户的对话是独立的"""
)

# ==================== 模拟多用户对话 ====================

def simulate_customer_conversation(user_id: str, user_name: str):
    """模拟单个用户的完整对话流程"""
    thread_id = f"customer_{user_id}"
    print(f"\n{'='*60}")
    print(f"用户：{user_name} (ID: {user_id})")
    print(f"{'='*60}\n")

    # 初始状态
    initial_state = {
        "user_id": user_id,
        "user_name": user_name,
        "issue_category": None,
        "ticket_id": None,
        "satisfaction_score": None
    }

    # 第 1 轮：报告问题
    print("👤 用户：我的账号登录不了，一直显示密码错误")
    response1 = customer_service_agent.invoke(
        {**initial_state, "messages": [{"role": "user", "content": "我的账号登录不了，一直显示密码错误"}]},
        {"configurable": {"thread_id": thread_id}}
    )
    print(f"🤖 客服：{response1['messages'][-1].content}\n")

    # 第 2 轮：查询进度
    print("👤 用户：我的工单处理得怎么样了？")
    response2 = customer_service_agent.invoke(
        {"messages": [{"role": "user", "content": "我的工单处理得怎么样了？"}]},
        {"configurable": {"thread_id": thread_id}}
    )
    print(f"🤖 客服：{response2['messages'][-1].content}\n")

    # 第 3 轮：提交反馈
    print("👤 用户：问题已解决，我给 5 分好评！")
    response3 = customer_service_agent.invoke(
        {"messages": [{"role": "user", "content": "问题已解决，我给 5 分好评，服务态度很棒！"}]},
        {"configurable": {"thread_id": thread_id}}
    )
    print(f"🤖 客服：{response3['messages'][-1].content}\n")

# ==================== 测试 ====================

if __name__ == "__main__":
    # 模拟多个用户的独立对话（完全隔离）
    simulate_customer_conversation("user_001", "张三")
    simulate_customer_conversation("user_002", "李四")
    simulate_customer_conversation("user_003", "王五")

    print("\n" + "="*60)
    print("多用户对话完成！每个用户的状态完全隔离。")
    print("="*60)

# 从 InMemorySaver 迁移到 PostgresSaver
old_checkpointer = InMemorySaver()
new_checkpointer = PostgresSaver.from_conn_string(DB_URI)

# 读取旧状态
state = old_checkpointer.get({"configurable": {"thread_id": "thread_001"}})

# 写入新 Checkpointer
new_checkpointer.put(
    {"configurable": {"thread_id": "thread_001"}},
    state.checkpoint,
    state.metadata
)

# 使用 PostgresSaver 定期清理旧数据
import psycopg2
from datetime import datetime, timedelta

# 删除 30 天前的检查点
conn = psycopg2.connect(DB_URI)
cursor = conn.cursor()

cutoff_date = datetime.now() - timedelta(days=30)
cursor.execute("""
    DELETE FROM checkpoints
    WHERE created_at < %s
""", (cutoff_date,))

conn.commit()
cursor.close()
conn.close()

# 监控 Checkpointer 使用情况
from langchain.agents.middleware import after_agent
import logging

logger = logging.getLogger(__name__)

@after_agent
def log_checkpoint_stats(state, runtime):
    """记录检查点统计信息"""
    message_count = len(state["messages"])
    thread_id = runtime.config.get("configurable", {}).get("thread_id")

    logger.info(f"Thread {thread_id}: {message_count} messages saved")

    # 警告：消息数量过多
    if message_count > 100:
        logger.warning(f"Thread {thread_id} has {message_count} messages - consider cleanup")

    return None