短期記憶主要是指將對(duì)話消息進(jìn)行存儲(chǔ),用作下一次對(duì)話的記憶來源

短期記憶開啟

短期記憶在開發(fā)環(huán)境中可以使用ImMemorySaver來存儲(chǔ),其數(shù)據(jù)是存儲(chǔ)在內(nèi)存中,應(yīng)用重啟則清空.
在生產(chǎn)環(huán)境中則是需要使用數(shù)據(jù)庫來實(shí)現(xiàn).

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(), # 創(chuàng)建一個(gè)內(nèi)存記憶體
)

agent.invoke(
    {"messages": [{"role": "user", "content": "Hi! My name is Bob."}]},
    {"configurable": {"thread_id": "1"}}, # 此步驟很重要,對(duì)于同一個(gè)對(duì)話上下文,thread_id要保持一致,不然將無法使用歷史消息
)

持久存儲(chǔ)的記憶

在使用前需要安裝對(duì)應(yīng)的數(shù)據(jù)庫模塊,如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,
    )

消息處理

過長(zhǎng)的對(duì)話內(nèi)容會(huì)影響輸出質(zhì)量以及增加消耗,過段則會(huì)讓模型忘記上下文,因此對(duì)于消息處理有如下幾種方案:

消息裁剪

此方案主要是移除消息列表中的部分消息,例如移除前幾條消息,或者后幾條消息.
消息移除需要使用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

# 在模型處理前對(duì)消息做裁剪
@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函數(shù)實(shí)現(xiàn)

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]}

消息總結(jié)

對(duì)消息使用模型進(jìn)行壓縮總結(jié),可以使用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!
"""

狀態(tài)數(shù)據(jù)的讀取

對(duì)于狀態(tài)(runtime.state)中的數(shù)據(jù),可以用作對(duì)話級(jí)別的數(shù)據(jù)緩存,即在一次對(duì)話里,其數(shù)據(jù)是互通的.

• 示例

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:
    """
    賬號(hào)充值
    :param phone:需要充值的手機(jī)號(hào)
    :param amount: 需要充值的金額(單位為分)
    :param runtime: 工具上下文
    :return: 充值后的結(jié)果文案
    """
    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:手機(jī)號(hào)
    :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="你是一個(gè)助手。每次只執(zhí)行一個(gè)工具調(diào)用。在執(zhí)行完充值操作后，請(qǐng)等待用戶的進(jìn)一步指令或觀察狀態(tài)變化后再進(jìn)行查詢。")

agent.invoke({"messages":[{'role':'user', 'content':"請(qǐng)為手機(jī)號(hào)1234567890充值10元,然后查詢充值完成后的余額"}]})

• 輸出

[recharge]
[query_user_balance] 小張