05 短期记忆

记忆模块在一个AI Agent系统是非常重要的，LLM本身是没有记忆的，因此需要我们自己来管理记忆信息，而记忆信息又分为短期记忆和长期记忆。

短期记忆：一般是指保存最近的几轮对话信息，在请求LLM时，一起将这些对话填充到填充系统消息、人类消息中，一起发送给LLM，这样Agent就拥有短期记忆。

长期记忆：当在同一个会话中，对话内容过长时，如果将所有的对话内容当成短期记忆传递给LLM，就可能超出上下文大小范围、分不清提问的重点、答案出现幻觉等情况。这个时候我们就需要把过往的对话信息总结成一段长期记忆，在请求LLM时，同样将长期记忆信息填充系统消息、人类消息中，一起发送给LLM，这样Agent就拥有了长期记忆，并且长期记忆支持跨会话读取和使用。

一、短期记忆基础用法

在LangChainv1中，通过create_agent创建的agent底层是通过LangGraph实现的，在LangGraph中，实现短期记忆功能的组件是State，因此学习短期记忆组件主要就是学习如何使用State组件。

在LangGraph中，想要使用短期记忆组件，就必须先创建一个checkpointer，并将这个checkpointer传给Agent，这里使用的是一个基于内存的checkpointer：InMemorySaver。

短期记忆使用示例如下：

from langchain.agents import create_agent
from langchain.messages import HumanMessage
from langgraph.checkpoint.memory import InMemorySaver

import dotenv
dotenv.load_dotenv()

# 1.定义Agent
checkpointer = InMemorySaver()
agent = create_agent(model="deepseek-v4-flash",
                     checkpointer=checkpointer)

# 2.第一次调用Agent
first_human_message = HumanMessage("你好，我是大志，你是？")
state = agent.invoke({"messages": [first_human_message]},
                     config={"configurable": {"thread_id": "1000"}},
                     )
first_content = state.get("messages")[-1].content
print(f"Human：{first_human_message.content}")
print(f"AI:{first_content}")
print("--------------------------------------------")

# 3.第二次调用Agent
second_human_message = HumanMessage("你还记得我是谁吗？")
state = agent.invoke({"messages": ["你还记得我是谁吗？"]},
                     config={"configurable": {"thread_id": "1000"}},
                     )
second_content = state.get("messages")[-1].content
print(f"Human：{second_human_message.content}")
print(f"AI:{second_content}")

执行结果：

在生产环境使用推荐使用PostgresSaver作为Agent的checkpointer，在使用之前需要安装如下依赖，具体用法与类似，这里就不过多演示了。

uv add langgraph-checkpoint-postgres

二、自定义短期记忆内容

在使用create_agent创建的Agent中，默认使用AgentState来管理短期记忆，默认的AgentState类中，核心功能是管理了消息列表相关的messages信息，如果我们想自定义短期记忆的内容，可以直接继承AgentState类，在定义类中添加想要的短期记忆内容。

代码示例如下：

from langchain.agents import create_agent
from langchain.messages import HumanMessage
from langchain.agents import AgentState
from langgraph.checkpoint.memory import InMemorySaver

import dotenv

dotenv.load_dotenv()


# 1.定义短期记忆内容
class CustomState(AgentState):
    username: str


# 2.定义Agent
checkpointer = InMemorySaver()
agent = create_agent(model="deepseek-v4-flash",
                     state_schema=CustomState,
                     checkpointer=checkpointer)

# 3.调用Agent
state = agent.invoke({"messages": [HumanMessage("1+1等于几")], "username": "大志说编程"},
                     config={"configurable": {"thread_id": "1000"}},
                     )

print(state.get("username"))

执行结果如下，在agent执行之后，打印state中的username属性，已经保存成功

三、消息处理

3.1 修剪消息

修剪消息功能可以通过中间件的方式去实现，使用@before_model装饰器定义一个中间件，被该装饰器标记的方法，内部逻辑会在执行model调用之前执行，具体中间件的知识会在后续的文章中进行讲解，这里只是了解用法即可。

在下面的示例中，定义了一个中间件trim_messages，当短期记忆中的消息超过五条就截取最近的5条消息，采用的方式是使用RemoveMessage(id=REMOVE_ALL_MESSAGES)删除所有消息，再把新的消息加入进去。

from langchain.messages import HumanMessage
from langchain.messages import RemoveMessage
from langgraph.graph.message import REMOVE_ALL_MESSAGES
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

from langgraph.checkpoint.memory import InMemorySaver

import dotenv

dotenv.load_dotenv()


# 1、定义消息修剪中间件
@before_model
def trim_messages(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    messages = state["messages"]
    if len(messages) <= 5:
        return None

    return {
        "messages": [
            RemoveMessage(id=REMOVE_ALL_MESSAGES),
            *messages[-5:]
        ]
    }


# 2.定义Agent
checkpointer = InMemorySaver()
agent = create_agent(model="deepseek-v4-flash",
                     system_prompt="你是一个专门为程序员服务的AI助手",
                     middleware=[trim_messages],
                     checkpointer=checkpointer)

# 3.调用Agent
config: RunnableConfig = {"configurable": {"thread_id": "1001"}}

agent.invoke({"messages": [HumanMessage("我是大志你是")]}, config=config)
agent.invoke({"messages": [HumanMessage("1+1等于几？")]}, config=config)
agent.invoke({"messages": [HumanMessage("1+2等于几？")]}, config=config)
agent.invoke({"messages": [HumanMessage("1+3等于几？")]}, config=config)

result = agent.invoke({"messages": [HumanMessage("你知道我是谁吗？")]}, config=config)
print(result.get("messages")[-1].content)

执行结果如下，当提问三轮对话后，第一轮的对话已经被剪裁到，所以Agent无法得知用户的身份。

3.2 删除消息

下面的示例，使用了@after_model装饰器，被该装饰器标记的方法，内部逻辑会在执行model调用之后执行，依然使用RemoveMessage进行消息删除，删除短期记忆的前两条消息。

from langchain.messages import HumanMessage
from langchain.messages import RemoveMessage
from langchain.agents import create_agent, AgentState
from langchain.agents.middleware import after_model
from langgraph.runtime import Runtime
from langchain_core.runnables import RunnableConfig
from typing import Any

from langgraph.checkpoint.memory import InMemorySaver

import dotenv

dotenv.load_dotenv()


# 1、定义消息删除中间件
@after_model
def delete_messages(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    messages = state["messages"]
    if len(messages) <= 2:
        return None

    return {
        "messages": [
            RemoveMessage(id=m.id) for m in messages[:2]
        ]
    }


# 2.定义Agent
checkpointer = InMemorySaver()
agent = create_agent(model="deepseek-v4-flash",
                     system_prompt="你是一个专门为程序员服务的AI助手",
                     middleware=[delete_messages],
                     checkpointer=checkpointer)

# 3.调用Agent
config: RunnableConfig = {"configurable": {"thread_id": "1001"}}

agent.invoke({"messages": [HumanMessage("我是大志你是")]}, config=config)
agent.invoke({"messages": [HumanMessage("1+1等于几？")]}, config=config)
result = agent.invoke({"messages": [HumanMessage("你知道我是谁吗？")]}, config=config)
print(result.get("messages")[-1].content)

执行结果如下，因为前两条消息被删除了，所Agent同样不知道我们的身份信息。

3.3 总结消息

当消息达到一定数量之后，为了减少短期记忆的上下文长度，除了对消息进行修剪和删除之外，LangChain还提供了一个压缩短期记忆的中间件SummarizationMiddleware，SummarizationMiddleware可以配置当短期记忆token达到多少时进行压缩，以及保留最近几条消息，到达指定token数量后，自动触发压缩。

from langchain.messages import HumanMessage
from langchain.agents import create_agent
from langchain.agents.middleware import SummarizationMiddleware
from langchain_core.runnables import RunnableConfig

from langgraph.checkpoint.memory import InMemorySaver

import dotenv

dotenv.load_dotenv()


# 1.定义Agent
checkpointer = InMemorySaver()
agent = create_agent(model="deepseek-v4-flash",
                     system_prompt="你是一个专门为程序员服务的AI助手",
                     middleware=[SummarizationMiddleware(
                         model="deepseek-v4-flash",
                         trigger=("tokens", 100),
                         keep=("messages", 2)
                     )],
                     checkpointer=checkpointer)

# 2.调用Agent
config: RunnableConfig = {"configurable": {"thread_id": "1001"}}

agent.invoke({"messages": [HumanMessage("我是大志你是")]}, config=config)
agent.invoke({"messages": [HumanMessage("1+1等于几？")]}, config=config)
agent.invoke({"messages": [HumanMessage("1+2等于几？")]}, config=config)
agent.invoke({"messages": [HumanMessage("1+3等于几？")]}, config=config)

result = agent.invoke({"messages": [HumanMessage("你知道我是谁吗？")]}, config=config)
print(result.get("messages")[-1].content)
print(result.get("messages"))

执行结果如下，SummarizationMiddleware中间件将除了最近的两条消息之前的消息进行总结，并且以HumanMessage的形式插入到短期记忆的消息列表中，并且会把多余的消息从短期记忆中剔除。

四、读取记忆信息

4.1 从工具中读取

在工具内部可以使用ToolTime进行短期记忆的获取和写入，前面在工具用法中已经详细介绍过，这里不再赘述，用法如下：

@tool
def search_weather(runtime: ToolRuntime, city: str):
    """查询今日天气"""
    print(runtime.state["messages"])
    return f"今日{city}天气晴，东风4级，体感舒适"

4.2 从中间件中读取

在被@before_model和@after_model包裹的中间件中，都可以获取到短期记忆信息，在前面修剪和删除消息的示例中已经展示过，这里不再赘述，用法如下：

@before_model包裹的中间件：

@before_model
def trim_messages(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    messages = state["messages"]
    if len(messages) <= 5:
        return None

    return {
        "messages": [
            RemoveMessage(id=REMOVE_ALL_MESSAGES),
            *messages[-5:]
        ]
    }

@after_model包裹的中间件：

@after_model
def delete_messages(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    messages = state["messages"]
    if len(messages) <= 2:
        return None

    return {
        "messages": [
            RemoveMessage(id=m.id) for m in messages[:2]
        ]
    }