LangChain Short-term memory：短期记忆使用完全指南

LangChain 中的短期记忆是AI智能体（Agent）实现单会话/线程内交互记忆的核心能力，让智能体能够记住同一场对话中的历史交互信息，适配用户偏好、学习反馈并完成复杂多步任务。短期记忆的核心载体是对话历史，但受限于大语言模型（LLM）的上下文窗口，需通过裁剪、删除、总结等策略管理，避免上下文溢出或模型性能下降。

本文将从短期记忆的核心概念、基础使用、生产配置、自定义扩展、历史管理策略、记忆读写方式等维度全面讲解。

一、核心概念 📚

1.1 短期记忆的定义与价值

短期记忆让应用在单个线程/对话内记住之前的交互信息，一个线程（Thread） 就是会话中多个交互的集合（类似邮件将同主题消息归为一个对话）。

• ✅ 核心价值：实现上下文连贯的对话、适配用户个性化偏好、完成多步依赖的复杂任务；

• ❌ 核心挑战：长对话的消息历史会超出LLM的上下文窗口，导致上下文丢失/报错，即使模型支持长上下文，也会因冗余信息出现响应变慢、成本升高、注意力分散的问题。

1.2 核心组件

组件	作用	通俗理解
AgentState	管理短期记忆的默认状态类	记忆的「容器」，默认通过messages键存储对话历史
checkpointer	状态持久化工具	记忆的「保存器」，将状态存到内存/数据库，支持会话断点续传
Thread ID	线程唯一标识	对话的「身份证」，区分不同用户/会话的记忆，避免混淆
中间件（Middleware）	处理消息历史的钩子函数	记忆的「管理员」，在模型调用前/后裁剪、删除、总结消息

1.3 消息的组成

LLM的上下文由消息列表组成，交互过程中消息会交替增长：

• 系统消息（System）：给模型的指令（如「你是一个贴心助手」）；

• 人类消息（Human）：用户的输入；

• 智能体消息（AI）：模型的回复；

• 工具消息（Tool）：工具执行的结果。

二、基础使用 ✅

为智能体添加短期记忆的核心要求是：创建Agent时通过checkpointer参数指定状态持久化器，LangChain会将短期记忆作为Agent状态的一部分管理，通过thread_id区分不同会话。

2.1 开发环境配置（内存持久化）

使用InMemorySaver作为checkpointer，数据仅存于内存，重启后丢失，适合开发/测试场景。

代码示例+详解

# 导入核心依赖
from langchain.agents import create_agent
from langgraph.checkpoint.memory import InMemorySaver

# 1. 创建智能体，指定checkpointer开启短期记忆
# model：使用的大模型（如gpt-4o、gpt-5）
# tools：智能体绑定的工具列表（示例为get_user_info）
# checkpointer：内存持久化器，开启短期记忆的核心
agent = create_agent(
    "gpt-5",
    [get_user_info],  # 自定义工具，可替换为任意工具/空列表
    checkpointer=InMemorySaver(),
)

# 2. 调用智能体，通过configurable指定thread_id（会话唯一标识）
# thread_id="1"：标记该交互属于第1个会话，后续同thread_id会复用记忆
agent.invoke(
    {"messages": [{"role": "user", "content": "Hi! My name is Bob."}]},
    {"configurable": {"thread_id": "1"}},
)

• 关键要点：

  1. 必须指定checkpointer，否则智能体无短期记忆，每次调用都是全新会话；

  2. thread_id是会话隔离的核心，不同thread_id的交互不会共享记忆；

  3. 调用时的messages为列表格式，每个元素包含role和content。

2.2 生产环境配置（数据库持久化）

生产环境需使用数据库后端的checkpointer，避免内存持久化的数据丢失，LangChain官方推荐PostgresSaver（PostgreSQL）。

步骤1：安装依赖

pip install langgraph-checkpoint-postgres

步骤2：代码实现+详解

from langchain.agents import create_agent
from langgraph.checkpoint.postgres import PostgresSaver

# 1. 配置数据库连接地址
# 格式：postgresql://用户名:密码@地址:端口/数据库名?sslmode=disable
DB_URI = "postgresql://postgres:postgres@localhost:5442/postgres?sslmode=disable"

# 2. 初始化PostgresSaver，自动创建表
# with上下文：自动管理连接的创建/关闭
with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
    checkpointer.setup()  # 自动在PostgreSQL中创建所需表，无需手动建表
    
    # 3. 创建智能体，指定数据库版checkpointer
    agent = create_agent(
        "gpt-5",
        [get_user_info],
        checkpointer=checkpointer,  # 替换为数据库持久化器
    )
    
    # 4. 调用智能体，用法与开发环境一致
    agent.invoke(
        {"messages": [{"role": "user", "content": "Hi! My name is Bob."}]},
        {"configurable": {"thread_id": "1"}},
    )

• 生产建议：

  1. 数据库密码避免硬编码，使用环境变量（如os.getenv("DB_PASSWORD")）；

  2. 调整PostgreSQL的连接池、超时配置，适配高并发场景；

  3. 定期清理过期的会话数据，避免数据库膨胀。

三、自定义智能体记忆 🛠️

默认情况下，智能体通过AgentState管理记忆，仅包含messages键存储对话历史。LangChain支持**扩展AgentState**，添加自定义字段（如用户ID、偏好设置），让短期记忆存储更多个性化信息。

3.1 自定义状态类

通过继承AgentState添加自定义字段，创建Agent时通过state_schema参数指定。

代码示例+详解

from langchain.agents import create_agent, AgentState
from langgraph.checkpoint.memory import InMemorySaver

# 1. 继承AgentState，定义自定义状态类
# 添加user_id（字符串）、preferences（字典）两个自定义字段
class CustomAgentState(AgentState):
    user_id: str
    preferences: dict

# 2. 创建智能体，通过state_schema指定自定义状态
agent = create_agent(
    "gpt-5",
    [get_user_info],
    state_schema=CustomAgentState,  # 启用自定义记忆字段
    checkpointer=InMemorySaver(),
)

# 3. 调用智能体，传入自定义状态字段
# 除了messages，还可以传入user_id、preferences
result = agent.invoke(
    {
        "messages": [{"role": "user", "content": "Hello"}],
        "user_id": "user_123",  # 自定义字段：用户ID
        "preferences": {"theme": "dark"}  # 自定义字段：用户偏好（深色主题）
    },
    {"configurable": {"thread_id": "1"}}
)

• 关键要点：

  1. 自定义字段需指定类型注解（如str/dict），符合Pydantic规范；

  2. 调用时可直接传入自定义字段，会被纳入短期记忆，同thread_id的后续调用可读取；

  3. 自定义状态支持所有可序列化的数据类型（字符串、数字、字典、列表）。

四、对话历史管理的通用策略 📝

长对话的消息列表会持续增长，最终超出LLM的上下文窗口，LangChain提供4种通用解决方案，平衡记忆完整性和上下文窗口限制，其中裁剪、删除、总结是最常用的三种。

4.1 裁剪消息（Trim messages）

核心逻辑：在调用模型前，移除消息列表中最早的N条消息，仅保留最新的部分消息，是最轻量化的管理策略。

实现方式：@before_model中间件

@before_model是模型调用前的钩子函数，会在LLM处理消息前执行，适合做消息预处理，核心工具为RemoveMessage（移除消息）和REMOVE_ALL_MESSAGES（清空所有消息）。

代码示例+详解

# 导入核心依赖
from langchain.messages import RemoveMessage
from langgraph.graph.message import REMOVE_ALL_MESSAGES
from langgraph.checkpoint.memory import InMemorySaver
from langchain.agents import create_agent, AgentState
from langchain.agents.middleware import before_model
from langgraph.runtime import Runtime
from langchain_core.runnables import RunnableConfig
from typing import Any

# 1. 定义裁剪消息的中间件，@before_model装饰器标记为模型前执行
@before_model
def trim_messages(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    """仅保留最新的3/4条消息+第一条消息，适配上下文窗口"""
    messages = state["messages"]  # 从状态中读取消息历史
    if len(messages) <= 3:
        return None  # 消息数≤3时不处理，直接返回
    
    # 保留第一条消息（避免丢失初始上下文）+ 最新的3/4条消息
    first_msg = messages[0]
    # 偶数条保留最后3条，奇数条保留最后4条（自定义规则，可按需修改）
    recent_messages = messages[-3:] if len(messages) % 2 == 0 else messages[-4:]
    new_messages = [first_msg] + recent_messages  # 拼接新的消息列表
    
    # 返回新的消息列表，先清空所有消息，再添加筛选后的消息
    return {
        "messages": [
            RemoveMessage(id=REMOVE_ALL_MESSAGES),  # 清空原有消息
            *new_messages  # 解包添加新消息
        ]
    }

# 2. 创建智能体，将中间件传入middleware参数
agent = create_agent(
    model="gpt-4o",  # 替换为实际使用的模型
    tools=[],  # 无工具，纯对话
    middleware=[trim_messages],  # 启用消息裁剪中间件
    checkpointer=InMemorySaver(),
)

# 3. 多次调用，验证记忆效果
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
agent.invoke({"messages": [{"role": "user", "content": "hi, my name is bob"}]}, config)
agent.invoke({"messages": [{"role": "user", "content": "write a short poem about cats"}]}, config)
agent.invoke({"messages": [{"role": "user", "content": "now do the same but for dogs"}]}, config)

# 第四次调用：查询姓名，验证裁剪后仍保留关键记忆
final_response = agent.invoke({"messages": [{"role": "user", "content": "what's my name?"}]}, config)
final_response["messages"][-1].pretty_print()  # 打印模型回复

输出结果

================================== Ai Message ==================================
Your name is Bob. You told me that earlier. If you'd like me to call you a nickname or use a different name, just say the word.

• 自定义规则建议：

  1. 按令牌数（Token） 裁剪（更精准）：通过tiktoken计算消息的令牌数，达到阈值时裁剪；

  2. 保留系统消息：确保模型不会丢失核心指令；

  3. 避免裁剪工具调用相关消息：防止模型无法匹配工具调用和结果。

4.2 删除消息（Delete messages）

核心逻辑：从LangGraph的状态中永久删除指定消息/所有消息，适合主动清理无用的历史信息，与裁剪的区别是：删除是主动精准清理，裁剪是被动批量保留。

核心工具

• RemoveMessage(id=消息ID)：删除指定消息；

• RemoveMessage(id=REMOVE_ALL_MESSAGES)：删除所有消息；

• 注：需配合add_messages归约器使用，默认AgentState已内置，无需额外配置。

实现方式1：删除指定消息

from langchain.messages import RemoveMessage

def delete_messages(state):
    """删除最早的2条消息"""
    messages = state["messages"]
    if len(messages) > 2:
        # 遍历最早的2条消息，通过ID删除
        return {"messages": [RemoveMessage(id=m.id) for m in messages[:2]]}

实现方式2：删除所有消息

from langgraph.graph.message import REMOVE_ALL_MESSAGES

def delete_messages(state):
    """清空所有消息历史"""
    return {"messages": [RemoveMessage(id=REMOVE_ALL_MESSAGES)]}

实现方式3：模型调用后自动删除（@after_model中间件）

from langchain.messages import RemoveMessage
from langchain.agents import create_agent, AgentState
from langchain.agents.middleware import after_model
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.runtime import Runtime
from langchain_core.runnables import RunnableConfig

# @after_model：模型调用后执行的钩子函数
@after_model
def delete_old_messages(state: AgentState, runtime: Runtime) -> dict | None:
    """模型回复后，自动删除最早的2条消息"""
    messages = state["messages"]
    if len(messages) > 2:
        return {"messages": [RemoveMessage(id=m.id) for m in messages[:2]]}
    return None

# 创建智能体，启用该中间件
agent = create_agent(
    "gpt-5-nano",
    tools=[],
    system_prompt="Please be concise and to the point.",
    middleware=[delete_old_messages],
    checkpointer=InMemorySaver(),
)

# 调用验证
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
# 第一次调用
for event in agent.stream(
    {"messages": [{"role": "user", "content": "hi! I'm bob"}]},
    config,
    stream_mode="values",
):
    print([(message.type, message.content) for message in event["messages"]])
# 第二次调用：查询姓名，验证消息删除后仍保留关键记忆
for event in agent.stream(
    {"messages": [{"role": "user", "content": "what's my name?"}]},
    config,
    stream_mode="values",
):
    print([(message.type, message.content) for message in event["messages"]])

输出结果（消息逐步删除，最终仅保留最新交互）

[('human', "hi! I'm bob")]
[('human', "hi! I'm bob"), ('ai', 'Hi Bob! Nice to meet you. How can I help you today?')]
[('human', "hi! I'm bob"), ('ai', 'Hi Bob! Nice to meet you.'), ('human', "what's my name?")]
[('human', "hi! I'm bob"), ('ai', 'Hi Bob! Nice to meet you.'), ('human', "what's my name?"), ('ai', 'Your name is Bob.')]
[('human', "what's my name?"), ('ai', 'Your name is Bob.')]

⚠️ 注意事项

删除消息后需保证消息历史的有效性，需遵循LLM提供商的规则：

1. 部分提供商要求消息历史以用户消息开头；

2. 大部分提供商要求带工具调用的智能体消息，后续必须紧跟对应的工具结果消息，不可随意删除。

4.3 总结消息（Summarize messages）

核心逻辑：通过轻量LLM对早期的消息历史进行总结，用总结内容替代原始消息，既减少消息数量，又保留核心信息，是最推荐的长对话管理策略（解决裁剪/删除的信息丢失问题）。

实现方式：内置SummarizationMiddleware

LangChain提供开箱即用的总结中间件，无需手动实现总结逻辑，核心参数支持自定义触发条件和保留消息数。

代码示例+详解

from langchain.agents import create_agent
from langchain.agents.middleware import SummarizationMiddleware
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.runnables import RunnableConfig

# 1. 初始化持久化器
checkpointer = InMemorySaver()

# 2. 创建智能体，启用总结中间件
agent = create_agent(
    model="gpt-4o",  # 主模型（处理对话/生成回复）
    tools=[],
    middleware=[
        SummarizationMiddleware(
            model="gpt-4o-mini",  # 总结用的轻量模型（降低成本）
            max_tokens_before_summary=4000,  # 令牌数达到4000时触发总结
            messages_to_keep=20,  # 总结后保留最新的20条消息
        )
    ],
    checkpointer=checkpointer,
)

# 3. 多次调用，验证总结后仍保留记忆
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
agent.invoke({"messages": [{"role": "user", "content": "hi, my name is bob"}]}, config)
agent.invoke({"messages": [{"role": "user", "content": "write a short poem about cats"}]}, config)
agent.invoke({"messages": [{"role": "user", "content": "now do the same but for dogs"}]}, config)

# 第四次调用：查询姓名，验证总结后仍能记住
final_response = agent.invoke({"messages": [{"role": "user", "content": "what's my name?"}]}, config)
final_response["messages"][-1].pretty_print()

输出结果

================================== Ai Message ==================================
Your name is Bob!

• 核心参数优化建议：

  1. 总结模型选择轻量版（如gpt-4o-mini、claude-3-haiku），降低总结成本；

  2. max_tokens_before_summary设置为**主模型上下文窗口的70%**（如主模型为8k令牌，设为5600），预留足够空间；

  3. messages_to_keep根据对话频率调整，高频交互设为10-20，低频设为50-100。

4.4 自定义策略

除了上述三种通用策略，还可根据业务需求实现自定义消息管理逻辑，例如：

• 按关键词过滤：删除包含无意义内容的消息；

• 按时间过滤：删除超过指定时间的历史消息；

• 按消息类型过滤：仅保留人类消息和智能体核心回复，删除工具日志消息；

• 分主题总结：按对话主题对消息分组，每组单独总结。

五、短期记忆的读写方式 📖

LangChain支持多场景读取和修改短期记忆，核心分为工具中读写、中间件中读写（模型前/后），满足不同业务的记忆操作需求。

5.1 在工具中读写记忆

工具是智能体的核心能力，可通过ToolRuntime参数读取/修改短期记忆，该参数对LLM隐藏（模型不会看到），仅工具内部可访问。

5.1.1 读取短期记忆

通过runtime.state["字段名"]读取状态中的信息（包括默认的messages和自定义字段）。

from langchain.agents import create_agent, AgentState
from langchain.tools import tool, ToolRuntime

# 定义自定义状态，包含user_id字段
class CustomState(AgentState):
    user_id: str

# 定义工具，通过runtime.state读取user_id
@tool
def get_user_info(runtime: ToolRuntime) -> str:
    """根据用户ID查询用户信息"""
    user_id = runtime.state["user_id"]  # 读取短期记忆中的user_id
    # 模拟数据库查询，实际可替换为真实逻辑
    return "User is John Smith" if user_id == "user_123" else "Unknown user"

# 创建智能体，绑定工具和自定义状态
agent = create_agent(
    model="gpt-5-nano",
    tools=[get_user_info],
    state_schema=CustomState,
)

# 调用智能体，传入user_id
result = agent.invoke({
    "messages": [{"role": "user", "content": "look up user information"}],
    "user_id": "user_123"
})
print(result["messages"][-1].content)  # 输出：User is John Smith.

5.1.2 写入短期记忆

通过langgraph.types.Command对象修改/新增状态字段，返回Command(update={字段: 新值})即可，支持修改自定义字段和消息历史。

# 导入核心依赖
from langchain.tools import tool, ToolRuntime
from langchain.messages import ToolMessage
from langchain.agents import create_agent, AgentState
from langgraph.types import Command
from pydantic import BaseModel

# 1. 定义自定义状态和上下文
class CustomState(AgentState):
    user_name: str  # 自定义字段：用户名
class CustomContext(BaseModel):
    user_id: str    # 上下文：用户ID（不可变）

# 2. 工具1：更新用户信息到短期记忆
@tool
def update_user_info(runtime: ToolRuntime[CustomContext, CustomState]) -> Command:
    """查询并更新用户信息到短期记忆"""
    user_id = runtime.context.user_id  # 从上下文读取user_id
    name = "John Smith" if user_id == "user_123" else "Unknown user"
    # 返回Command，更新状态字段+消息历史
    return Command(update={
        "user_name": name,  # 新增/修改短期记忆的user_name字段
        "messages": [       # 更新消息历史，添加工具执行结果
            ToolMessage(
                "Successfully looked up user information",
                tool_call_id=runtime.tool_call_id  # 绑定工具调用ID
            )
        ]
    })

# 3. 工具2：从短期记忆读取用户名并问候
@tool
def greet(runtime: ToolRuntime[CustomContext, CustomState]) -> str:
    """读取用户名并问候用户"""
    user_name = runtime.state.get("user_name", None)  # 读取user_name
    if user_name is None:
        # 未获取到用户名时，提示调用update_user_info工具
        return Command(update={
            "messages": [
                ToolMessage(
                    "Please call the 'update_user_info' tool it will get and update the user's name.",
                    tool_call_id=runtime.tool_call_id
                )
            ]
        })
    return f"Hello {user_name}!"  # 获取到则返回问候语

# 4. 创建智能体
agent = create_agent(
    model="gpt-5-nano",
    tools=[update_user_info, greet],
    state_schema=CustomState,
    context_schema=CustomContext,
)

# 5. 调用智能体，传入上下文
agent.invoke(
    {"messages": [{"role": "user", "content": "greet the user"}]},
    context=CustomContext(user_id="user_123"),
)

• 关键要点：

  1. ToolRuntime[CustomContext, CustomState]指定上下文和状态的类型，实现类型校验；

  2. tool_call_id用于绑定工具调用和结果，避免多工具并行调用时的混乱；

  3. 可通过Command同时更新多个字段，包括自定义字段和messages。

5.2 在中间件中读写记忆

通过动态提示中间件、@before_model、@after_model中间件，可在模型调用的不同阶段读写记忆，实现动态提示生成、消息预处理、结果后处理。

5.2.1 动态提示（Dynamic Prompt）

根据记忆/上下文生成个性化的系统提示，例如根据用户名称呼用户，核心工具为@dynamic_prompt装饰器。

from langchain.agents import create_agent
from typing import TypedDict
from langchain.agents.middleware import dynamic_prompt, ModelRequest

# 定义上下文
class CustomContext(TypedDict):
    user_name: str

# 定义工具
def get_weather(city: str) -> str:
    """查询城市天气"""
    return f"The weather in {city} is always sunny!"

# 动态生成系统提示，从上下文读取用户名
@dynamic_prompt
def dynamic_system_prompt(request: ModelRequest) -> str:
    user_name = request.runtime.context["user_name"]  # 读取上下文/记忆
    return f"You are a helpful assistant. Address the user as {user_name}."  # 个性化提示

# 创建智能体
agent = create_agent(
    model="gpt-5-nano",
    tools=[get_weather],
    middleware=[dynamic_system_prompt],  # 启用动态提示
    context_schema=CustomContext,
)

# 调用智能体
result = agent.invoke(
    {"messages": [{"role": "user", "content": "What is the weather in SF?"}]},
    context=CustomContext(user_name="John Smith"),
)

# 打印所有消息
for msg in result["messages"]:
    msg.pretty_print()

输出结果（模型会称呼用户名）

================================ Human Message =================================
What is the weather in SF?
================================== Ai Message ==================================
Tool Calls:
  get_weather (call_WFQlOGn4b2yoJrv7cih342FG)
  Call ID: call_WFQlOGn4b2yoJrv7cih342FG
   Args:
    city: San Francisco
================================ Tool Message =================================
Name: get_weather
The weather in San Francisco is always sunny!
================================== Ai Message ==================================
Hi John Smith, the weather in San Francisco is always sunny!

5.2.2 模型调用前（@before_model）

在模型处理消息前读写记忆，核心用途：裁剪消息、补充上下文、过滤敏感信息，示例见「4.1 裁剪消息」。

5.2.3 模型调用后（@after_model）

在模型生成回复后读写记忆，核心用途：验证回复、删除敏感消息、更新记忆字段。

from langchain.messages import RemoveMessage
from langgraph.checkpoint.memory import InMemorySaver
from langchain.agents import create_agent, AgentState
from langchain.agents.middleware import after_model
from langgraph.runtime import Runtime

# 模型回复后，删除包含敏感词的消息
@after_model
def validate_response(state: AgentState, runtime: Runtime) -> dict | None:
    """移除包含敏感词（password/secret）的消息"""
    STOP_WORDS = ["password", "secret"]
    last_message = state["messages"][-1]  # 读取最新的模型回复
    # 检测是否包含敏感词
    if any(word in last_message.content for word in STOP_WORDS):
        return {"messages": [RemoveMessage(id=last_message.id)]}  # 删除敏感消息
    return None

# 创建智能体
agent = create_agent(
    model="gpt-5-nano",
    tools=[],
    middleware=[validate_response],  # 启用敏感词验证
    checkpointer=InMemorySaver(),
)

六、核心代码片段整合（可直接运行） 🚀

以下代码整合了短期记忆基础使用、自定义状态、消息总结、工具读写记忆、动态提示的核心功能，基于PostgreSQL持久化（开发环境可替换为InMemorySaver），只需替换模型和数据库配置即可运行。

6.1 完整可运行代码

# 安装依赖
# pip install langchain langgraph langgraph-checkpoint-postgres openai tiktoken
import os
from langchain_openai import ChatOpenAI
from langchain.agents import create_agent, AgentState
from langchain.tools import tool, ToolRuntime
from langchain.messages import ToolMessage
from langgraph.checkpoint.postgres import PostgresSaver
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command
from pydantic import BaseModel
from langchain.agents.middleware import SummarizationMiddleware, dynamic_prompt, ModelRequest
from langchain_core.runnables import RunnableConfig

# -------------------------- 1. 配置基础环境 --------------------------
# 替换为自己的OpenAI API Key
os.environ["OPENAI_API_KEY"] = "your-api-key"
# 模型配置（主模型+总结模型）
MAIN_MODEL = ChatOpenAI(model="gpt-4o")
SUMMARY_MODEL = ChatOpenAI(model="gpt-4o-mini")
# 数据库配置（生产环境）/ 内存持久化（开发环境）
USE_PROD = False  # True=生产环境（PostgreSQL），False=开发环境（内存）
DB_URI = "postgresql://postgres:postgres@localhost:5442/postgres?sslmode=disable"

# -------------------------- 2. 自定义状态和上下文 --------------------------
class CustomState(AgentState):
    """自定义短期记忆状态，添加用户ID和偏好"""
    user_id: str
    user_preferences: dict

class CustomContext(BaseModel):
    """自定义上下文，不可变"""
    session_id: str

# -------------------------- 3. 定义工具（读写短期记忆） --------------------------
@tool
def get_user_preference(runtime: ToolRuntime[CustomContext, CustomState]) -> str:
    """读取短期记忆中的用户偏好"""
    pref = runtime.state.get("user_preferences", {})
    user_id = runtime.state["user_id"]
    return f"User {user_id}'s preferences: {pref}"

@tool
def update_user_preference(runtime: ToolRuntime[CustomContext, CustomState]) -> Command:
    """更新用户偏好到短期记忆"""
    # 模拟从上下文获取偏好，实际可替换为用户输入解析
    return Command(update={
        "user_preferences": {"theme": "dark", "language": "zh-CN"},
        "messages": [
            ToolMessage("Updated user preferences successfully", tool_call_id=runtime.tool_call_id)
        ]
    })

# -------------------------- 4. 动态提示中间件 --------------------------
@dynamic_prompt
def custom_dynamic_prompt(request: ModelRequest) -> str:
    """根据会话ID生成个性化提示"""
    session_id = request.runtime.context["session_id"]
    return f"""你是一个贴心的AI助手，会话ID：{session_id}
    请使用中文回复，结合用户的短期记忆信息提供个性化服务。"""

# -------------------------- 5. 初始化Checkpointer --------------------------
if USE_PROD:
    checkpointer = PostgresSaver.from_conn_string(DB_URI)
    checkpointer.setup()
else:
    checkpointer = InMemorySaver()

# -------------------------- 6. 创建智能体（启用短期记忆+总结） --------------------------
agent = create_agent(
    model=MAIN_MODEL,
    tools=[get_user_preference, update_user_preference],
    state_schema=CustomState,  # 自定义状态
    context_schema=CustomContext,  # 自定义上下文
    checkpointer=checkpointer,  # 开启短期记忆
    middleware=[
        # 消息总结中间件
        SummarizationMiddleware(
            model=SUMMARY_MODEL,
            max_tokens_before_summary=4000,
            messages_to_keep=20,
        ),
        # 动态提示中间件
        custom_dynamic_prompt
    ],
    system_prompt="请始终使用中文回复，简洁明了。"
)

# -------------------------- 7. 调用智能体（体验短期记忆） --------------------------
config: RunnableConfig = {"configurable": {"thread_id": "session_001"}}  # 会话ID
# 第一次调用：初始化用户信息
agent.invoke(
    {
        "messages": [{"role": "user", "content": "你好，我的用户ID是user_001"}],
        "user_id": "user_001",
        "user_preferences": {}
    },
    context=CustomContext(session_id="session_001"),
    config=config
)

# 第二次调用：更新用户偏好
agent.invoke(
    {"messages": [{"role": "user", "content": "更新我的用户偏好"}]},
    context=CustomContext(session_id="session_001"),
    config=config
)

# 第三次调用：读取用户偏好（验证短期记忆）
result = agent.invoke(
    {"messages": [{"role": "user", "content": "查询我的用户偏好"}]},
    context=CustomContext(session_id="session_001"),
    config=config
)

# 打印结果
for msg in result["messages"]:
    if msg.type == "ai":
        print(f"AI回复：{msg.content}")

6.2 运行说明

1. 开发环境：将USE_PROD设为False，无需配置PostgreSQL，直接运行；

2. 生产环境：将USE_PROD设为True，配置正确的PostgreSQL连接地址；

3. 替换OPENAI_API_KEY为自己的密钥，支持替换为其他LLM（如Claude、文心一言）；

4. 运行后，智能体将记住user_001的偏好，同thread_id的后续调用可直接读取。

七、核心开发要点总结 📌

1. 开启短期记忆的核心：创建Agent时必须指定checkpointer，开发用InMemorySaver，生产用PostgresSaver；

2. 会话隔离：通过thread_id区分不同会话，避免记忆混淆；

3. 记忆扩展：继承AgentState添加自定义字段，存储用户ID、偏好等个性化信息；

4. 长对话管理：优先使用SummarizationMiddleware总结消息，避免信息丢失，其次使用裁剪/删除；

5. 记忆读写：工具中通过ToolRuntime读写，模型前后通过@before_model/@after_model读写；

6. 类型校验：自定义状态/上下文需指定类型注解，符合Pydantic规范；

7. 生产优化：数据库持久化需配置环境变量、定期清理数据，总结用轻量模型降低成本。