短期记忆主要是指将对话消息进行存储,用作下一次对话的记忆来源

短期记忆开启

短期记忆在开发环境中可以使用ImMemorySaver来存储,其数据是存储在内存中,应用重启则清空.
在生产环境中则是需要使用数据库来实现.

ImMemorySaver方案

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


agent = create_agent(
    "gpt-5.4",
    tools=[get_user_info],
    checkpointer=InMemorySaver(), # 创建一个内存记忆体
)

agent.invoke(
    {"messages": [{"role": "user", "content": "Hi! My name is Bob."}]},
    {"configurable": {"thread_id": "1"}}, # 此步骤很重要,对于同一个对话上下文,thread_id要保持一致,不然将无法使用历史消息
)

持久存储的记忆

在使用前需要安装对应的数据库模块,如pip install langgraph-checkpoint-postgres

• 使用

from langchain.agents import create_agent

from langgraph.checkpoint.postgres import PostgresSaver  


DB_URI = "postgresql://postgres:postgres@localhost:5432/postgres?sslmode=disable"
with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
    checkpointer.setup() # auto create tables in PostgreSQL
    agent = create_agent(
        "gpt-5.4",
        tools=[get_user_info],
        checkpointer=checkpointer,
    )

消息处理

过长的对话内容会影响输出质量以及增加消耗,过段则会让模型忘记上下文,因此对于消息处理有如下几种方案:

消息裁剪

此方案主要是移除消息列表中的部分消息,例如移除前几条消息,或者后几条消息.
消息移除需要使用Command来更新runtime.state['messages']中的消息,其中使用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

# 在模型处理前对消息做裁剪
@before_model
def trim_messages(state: AgentState, runtime: Runtime) -> dict[str, Any] | None:
    """Keep only the last few messages to fit context window."""
    messages = state["messages"]

    if len(messages) <= 3:
        return None  # No changes needed

    first_msg = messages[0]
    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
        ] # 此步骤的含义是,先将所有的历史消息清空,然后用裁剪后的消息列表来代替之前的消息列表
    }

agent = create_agent(
    your_model_here,
    tools=your_tools_here,
    middleware=[trim_messages],
    checkpointer=InMemorySaver(),
)

config: RunnableConfig = {"configurable": {"thread_id": "1"}}

agent.invoke({"messages": "hi, my name is bob"}, config)
agent.invoke({"messages": "write a short poem about cats"}, config)
agent.invoke({"messages": "now do the same but for dogs"}, config)
final_response = agent.invoke({"messages": "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.
"""

消息过滤

消息过滤是指删除特定的消息,也是使用RemoveMessage函数实现

from langchain.messages import RemoveMessage  

def delete_messages(state):
    messages = state["messages"]
    if len(messages) > 2:
        # remove the earliest two messages
        return {"messages": [RemoveMessage(id=m.id) for m in messages if 'password' in m.context]}

消息总结

对消息使用模型进行压缩总结,可以使用lanchain自带的SummarizationMiddleware模块
示例:

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


checkpointer = InMemorySaver()

agent = create_agent(
    model="gpt-5.4",
    tools=[],
    middleware=[
        SummarizationMiddleware(
            model="gpt-5.4-mini",
            trigger=("tokens", 4000),
            keep=("messages", 20)
        )
    ],
    checkpointer=checkpointer,
)

config: RunnableConfig = {"configurable": {"thread_id": "1"}}
agent.invoke({"messages": "hi, my name is bob"}, config)
agent.invoke({"messages": "write a short poem about cats"}, config)
agent.invoke({"messages": "now do the same but for dogs"}, config)
final_response = agent.invoke({"messages": "what's my name?"}, config)

final_response["messages"][-1].pretty_print()
"""
================================== Ai Message ==================================

Your name is Bob!
"""

状态数据的读取

对于状态(runtime.state)中的数据,可以用作对话级别的数据缓存,即在一次对话里,其数据是互通的.

• 示例

from langchain_core.messages import ToolMessage
from langchain.agents import AgentState, create_agent
from langchain.tools import tool
from langgraph.prebuilt import  ToolRuntime
from langgraph.types import Command

class CustomAgentState(AgentState):
    user_name: str | None # 添加 user_name 字段，并允许为 None

@tool
def recharge(phone:str, amount:int, runtime: ToolRuntime[None, CustomAgentState])->Command:
    """
    账号充值
    :param phone:需要充值的手机号
    :param amount: 需要充值的金额(单位为分)
    :param runtime: 工具上下文
    :return: 充值后的结果文案
    """
    print('[recharge]')
    return Command(
        update={
            'user_name': '小张',
            'messages': [
                ToolMessage(
                    content=f"充值成功。",
                    tool_call_id=runtime.tool_call_id,
                )
            ],
        }
    )

@tool
def query_user_balance(phone:str, runtime: ToolRuntime[None, CustomAgentState])->int:
    """
    查询用户余额
    :param phone:手机号
    :return: 余额(分)
    """
    print('[query_user_balance]', runtime.state.get('user_name'))
    return 2000

agent = create_agent(model, tools=[recharge,query_user_balance], state_schema=CustomAgentState,system_prompt="你是一个助手。每次只执行一个工具调用。在执行完充值操作后，请等待用户的进一步指令或观察状态变化后再进行查询。")

agent.invoke({"messages":[{'role':'user', 'content':"请为手机号1234567890充值10元,然后查询充值完成后的余额"}]})

• 输出

[recharge]
[query_user_balance] 小张