背景：

langchain v1.0版本默认带的记忆存储没法自定义处理，比如：

使用官方提供 postgres数据库记忆关联包

pip install langgraph-checkpoint-postgres

存储表

表内容：二进制存储

官方包没法自定义设计表。

所以我们通过自定义代码 middleware中间件方式来在调用模型前后进行处理。

from dotenv import load_dotenv
load_dotenv()

from langchain_openai import ChatOpenAI
from langchain.agents import create_agent

# 测试
SYSTEM_PROMPT = """你是一个乐于助人的助手。"""

# 创建工具
from dataclasses import dataclass
from langchain.tools import tool, ToolRuntime

# @tool(return_direct=True) # 定义工具 return_direct=True 表示直接返回工具的结果
@tool
def get_weather_for_location(city: str) -> str:
    """获取指定城市的天气"""
    return f"城市：{city}，天气一直都是晴天!"

@dataclass
class Context:
    """自定义运行时上下文模式."""
    user_id: str # 用户ID
    # ... 可以加些自定义属性

@tool
def get_user_location(runtime: ToolRuntime[Context]) -> str:
    """根据用户ID检索用户信息."""
    user_id = runtime.context.user_id
    return "长沙" if user_id == "1" else "广州"

# 配置模型
model = ChatOpenAI(api_key="xxxxxxxx",
                    model="Qwen3-30B-A3B-Instruct-2507",
                    temperature=0.1,
                    timeout=10,
                    max_tokens=1000,
                    base_url="http://127.0.0.1:8000/v1")


from langchain.agents import AgentState
from langchain.agents.middleware import AgentMiddleware
from typing import Any, List, TypedDict, Required
from langchain.messages import SystemMessage, HumanMessage, AIMessage, RemoveMessage, ToolMessage
from langchain_core.messages.base import BaseMessage

from langgraph.graph.message import REMOVE_ALL_MESSAGES

# 自定义消息对象
class MyAiMessage(TypedDict):
    user_id: str
    data: List[BaseMessage]

# 自定义状态
class CustomState(AgentState):
    user_preferences: dict

# 自定义中间件
class CustomMiddleware(AgentMiddleware):
    state_schema = CustomState
    #tools = [get_weather_for_location, get_user_location]

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        print("CustomMiddleware init...")
        # 模拟数据库，这里模拟一个列表来保存用户消息
        self.my_messages: List[MyAiMessage] = []

    def get_or_create_user_messages(self, user_id: str):
        """
        根据user_id判断，从my_messages读取数据
        如果不存在则创建新的列表，存在则直接返回
        """
        # 检查是否存在该用户的记录
        for msg in self.my_messages:
            if msg['user_id'] == user_id:  # 使用字典访问方式
                return msg['data']

        # 如果不存在，创建新的记录
        new_user_msg = {'user_id': user_id, 'data': []}
        self.my_messages.append(new_user_msg)
        return new_user_msg['data']

    def before_model(self, state: CustomState, runtime) -> dict[str, Any] | None:
        # 模型调用处理
        print("before_model:",state)
        #print(self.state_schema)
        #print(runtime)

        user_id = runtime.context.user_id
        user_preferences = state.get("user_preferences", {})
        print("user_id:", user_id)
        print("user_preferences:", user_preferences)

        user_msg_datas = self.get_or_create_user_messages(user_id)

        user_msg_datas.extend(state["messages"])

        # 赋值历史聊天记录
        return {
            "messages": [
                RemoveMessage(id=REMOVE_ALL_MESSAGES),
                *user_msg_datas
            ]
        }

    def after_model(self, state: CustomState, runtime) -> dict[str, Any] | None:
        # 模型调用后处理
        print("after_model:", state)
        user_id = runtime.context.user_id

        user_msg_datas = self.get_or_create_user_messages(user_id)
        user_msg_datas.append(state["messages"][-1])

agent = create_agent(
    model=model,
    middleware=[CustomMiddleware()]
)

# 代理现在可以追踪消息以外的其他状态
result = agent.invoke(
    input = {
    "messages": [SystemMessage(SYSTEM_PROMPT),
                HumanMessage("我是张三")],
    "user_preferences": {"style": "技术", "verbosity": "详细"}
    },
    context=Context(user_id="1")
)

print("result:",result)
print()
print()

result = agent.invoke(
    {
    "messages": [{"role": "user", "content": "你知道我是谁吗？"}]
    },
    context=Context(user_id="1")
)

print("result:",result)

结果：

CustomMiddleware init...

before_model: {'messages': [SystemMessage(content='你是一个乐于助人的助手。', additional_kwargs={}, response_metadata={}, id='877feef7-87e0-4b96-be8a-a1dd1cd09540'), HumanMessage(content='我是张三', additional_kwargs={}, response_metadata={}, id='0445de49-1b08-4773-8198-b9332285b78f')], 'user_preferences': {'style': '技术', 'verbosity': '详细'}}

user_id: 1

user_preferences: {'style': '技术', 'verbosity': '详细'}

after_model: {'messages': [SystemMessage(content='你是一个乐于助人的助手。', additional_kwargs={}, response_metadata={}, id='877feef7-87e0-4b96-be8a-a1dd1cd09540'), HumanMessage(content='我是张三', additional_kwargs={}, response_metadata={}, id='0445de49-1b08-4773-8198-b9332285b78f'), AIMessage(content='你好，张三！很高兴认识你。有什么我可以帮你的吗？😊', additional_kwargs={'refusal': None}, id='lc_run--019c041d-5c16-7721-933b-e07322a8e353-0', usage_metadata={'input_tokens': 24, 'output_tokens': 17, 'total_tokens': 41, 'input_token_details': {}, 'output_token_details': {}})], 'user_preferences': {'style': '技术', 'verbosity': '详细'}}

result: {'messages': [SystemMessage(content='你是一个乐于助人的助手。', additional_kwargs={}, response_metadata={}, id='877feef7-87e0-4b96-be8a-a1dd1cd09540'), HumanMessage(content='我是张三', additional_kwargs={}, response_metadata={}, id='0445de49-1b08-4773-8198-b9332285b78f'), AIMessage(content='你好，张三！很高兴认识你。有什么我可以帮你的吗？😊', additional_kwargs={'refusal': None}, id='lc_run--019c041d-5c16-7721-933b-e07322a8e353-0', usage_metadata={'input_tokens': 24, 'output_tokens': 17, 'total_tokens': 41, 'input_token_details': {}, 'output_token_details': {}})], 'user_preferences': {'style': '技术', 'verbosity': '详细'}}

before_model: {'messages': [HumanMessage(content='你知道我是谁吗？', additional_kwargs={}, response_metadata={}, id='b3dd4493-9502-4218-bdf6-c39aed9d4763')]}

user_id: 1

user_preferences: {}

after_model: {'messages': [SystemMessage(content='你是一个乐于助人的助手。', additional_kwargs={}, response_metadata={}, id='877feef7-87e0-4b96-be8a-a1dd1cd09540'), HumanMessage(content='我是张三', additional_kwargs={}, response_metadata={}, id='0445de49-1b08-4773-8198-b9332285b78f'), AIMessage(content='你好，张三！很高兴认识你。有什么我可以帮你的吗？😊', additional_kwargs={'refusal': None}, id='lc_run--019c041d-5c16-7721-933b-e07322a8e353-0', usage_metadata={'input_tokens': 24, 'output_tokens': 17, 'total_tokens': 41, 'input_token_details': {}, 'output_token_details': {}}), HumanMessage(content='你知道我是谁吗？', additional_kwargs={}, response_metadata={}, id='b3dd4493-9502-4218-bdf6-c39aed9d4763'), AIMessage(content='当然知道啦，你就是张三呀！虽然我们是第一次聊天，但你已经告诉我你的名字了，我可不会忘记哦～✨ 有什么想聊的，或者需要帮忙的吗？', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 44, 'prompt_tokens': 55, 'total_tokens': 99, 'completion_tokens_details': None, 'prompt_tokens_details': None}, 'model_provider': 'openai', 'model_name': 'Qwen3-30B', 'system_fingerprint': None, 'id': 'chatcmpl-0798817249834dcd96f008353f675649', 'finish_reason': 'stop', 'logprobs': None}, id='lc_run--019c041d-5d04-7801-87ce-46ad8bc47b26-0', usage_metadata={'input_tokens': 55, 'output_tokens': 44, 'total_tokens': 99, 'input_token_details': {}, 'output_token_details': {}})]}

result: {'messages': [SystemMessage(content='你是一个乐于助人的助手。', additional_kwargs={}, response_metadata={}, id='877feef7-87e0-4b96-be8a-a1dd1cd09540'), HumanMessage(content='我是张三', additional_kwargs={}, response_metadata={}, id='0445de49-1b08-4773-8198-b9332285b78f'), AIMessage(content='你好，张三！很高兴认识你。有什么我可以帮你的吗？😊', additional_kwargs={'refusal': None}, id='lc_run--019c041d-5c16-7721-933b-e07322a8e353-0', usage_metadata={'input_tokens': 24, 'output_tokens': 17, 'total_tokens': 41, 'input_token_details': {}, 'output_token_details': {}}), HumanMessage(content='你知道我是谁吗？', additional_kwargs={}, response_metadata={}, id='b3dd4493-9502-4218-bdf6-c39aed9d4763'), AIMessage(content='当然知道啦，你就是张三呀！虽然我们是第一次聊天，但你已经告诉我你的名字了，我可不会忘记哦～✨ 有什么想聊的，或者需要帮忙的吗？', additional_kwargs={'refusal': None}, response_metadata={'token_usage': {'completion_tokens': 44, 'prompt_tokens': 55, 'total_tokens': 99, 'completion_tokens_details': None, 'prompt_tokens_details': None}, 'model_provider': 'openai', 'model_name': 'Qwen3-30B', 'system_fingerprint': None, 'id': 'chatcmpl-0798817249834dcd96f008353f675649', 'finish_reason': 'stop', 'logprobs': None}, id='lc_run--019c041d-5d04-7801-87ce-46ad8bc47b26-0', usage_metadata={'input_tokens': 55, 'output_tokens': 44, 'total_tokens': 99, 'input_token_details': {}, 'output_token_details': {}})]}