🎯 课程主题

将长期记忆（Store）集成到 LangGraph Agent 中，通过自定义工具在 Agent 运行时访问长期记忆，实现用户信息的存储（put）和读取（get），基于 InMemoryStore 进行演示。

📝 核心知识点

1. 在 Agent 中访问 Store 的方式

有两种主要方式在 Agent 中访问长期记忆：

方式说明
在工具中访问通过 ToolRuntime 获取 Store 实例，在工具函数内调用 put()、get()、search()
在中间件中访问通过钩子函数（如 before_model）在模型调用前后自动读写长期记忆

本节课聚焦第一种方式：在工具中访问。

2. 关键角色说明

角色说明
AgentStateLangGraph 内置的用户状态类，默认包含 messages 字段
CustomAgentState自定义状态类，继承 AgentState 并扩展字段（如 user_id），用于在工具中传递额外信息
ToolRuntime工具运行时环境对象，提供对 store 和 state 的访问
store (Store)长期记忆存储实例（InMemoryStore 或 PostgresStore）

3. 核心流程

invoke({"messages": [...], "user_id": "user-1"})
        │
        ▼
   CustomAgentState (user_id + messages)
        │
        ▼
   Agent 调用工具 (如 save_user_info)
        │
        ▼
   工具通过 ToolRuntime.state 获取 user_id
   工具通过 ToolRuntime.store 调用 put/get
        │
        ▼
   Store 读写长期记忆

🏗️ 架构与工作流

┌──────────────────────────────────────────────┐
│                  create_agent()               │
│  ┌─────────┐  ┌──────────┐  ┌─────────────┐ │
│  │  model   │  │  tools   │  │ store       │ │
│  │          │  │  ┌─────┐ │  │(InMemory)   │ │
│  │          │  │  │save_│ │  │             │ │
│  │          │  │  │user │─┼──▶ put()       │ │
│  │          │  │  │info │ │  │             │ │
│  │          │  │  ├─────┤ │  │             │ │
│  │          │  │  │get_ │ │  │             │ │
│  │          │  │  │user │─┼──▶ get()       │ │
│  │          │  │  │info │ │  │             │ │
│  │          │  │  └─────┘ │  │             │ │
│  └─────────┘  └──────────┘  └─────────────┘ │
│        state_schema: CustomAgentState         │
└──────────────────────────────────────────────┘

💻 代码实战

完整代码（InMemoryStore 版）

from langgraph.agent import create_agent
from langgraph.store.memory import InMemoryStore
from langgraph.prebuilt.tool_runtime import ToolRuntime
from typing import NotRequired
from typing_extensions import TypedDict
from langgraph.graph.state import AgentState

# 1. 初始化模型
from langchain.chat_models import init_chat_model
model = init_chat_model(model="gpt-4o-mini")

# 2. 初始化 InMemoryStore
store = InMemoryStore()

# 3. 自定义 AgentState — 扩展 user_id 字段
class CustomAgentState(AgentState):
    user_id: NotRequired[str]  # NotRequired 表示可选字段

# 4. 定义工具：保存用户信息到长期记忆
def save_user_info(
    name: str,
    runtime: ToolRuntime
) -> str:
    """将客户信息保存在长期记忆中。
    
    Args:
        name: 用户名
        runtime: 工具运行时环境
    Returns:
        保存状态字符串
    """
    user_id = runtime.state["user_id"]  # 从 state 获取 user_id
    runtime.store.put(
        namespace=("users",),
        key=user_id,
        value={"name": name}
    )
    return "saved"

# 5. 定义工具：从长期记忆读取用户信息
def get_user_info(runtime: ToolRuntime) -> str:
    """从长期记忆中读取客户信息。
    
    Args:
        runtime: 工具运行时环境
    Returns:
        用户的个人信息字符串
    """
    user_id = runtime.state["user_id"]
    item = runtime.store.get(namespace=("users",), key=user_id)
    if item:
        return str(item.value)
    return "INFO NOT FOUND"

# 6. 系统提示词
SYSTEM_PROMPT = """
用户提及个人信息时，使用工具保存用户信息。
如果用户询问个人信息时，尝试使用工具检索/读取用户信息。
"""

# 7. 创建 Agent
agent = create_agent(
    model=model,
    tools=[save_user_info, get_user_info],
    store=store,
    system_prompt=SYSTEM_PROMPT,
    state_schema=CustomAgentState  # 使用自定义状态
)

# 8. 调用 Agent — 第一次会话：存储信息
response = agent.invoke({
    "messages": [{"role": "user", "content": "你好，我是小花"}],
    "user_id": "user-1"  # 自定义字段，将存入 CustomAgentState
})
print(response["messages"][-1].content)
# 输出: "很高兴认识你，小花！"

# 9. 第二次会话：读取信息
response = agent.invoke({
    "messages": [{"role": "user", "content": "我是谁？"}],
    "user_id": "user-1"
})
print(response["messages"][-1].content)
# 输出: "你是小花。"

关键代码解析

自定义 AgentState：

class CustomAgentState(AgentState):
    user_id: NotRequired[str]
  • 继承 AgentState（默认含 messages 字段）。
  • 新增 user_id 字段，用 NotRequired 标记为可选。
  • 在 create_agent() 中通过 state_schema=CustomAgentState 指定。

工具中访问 Store 和 State：

def save_user_info(name: str, runtime: ToolRuntime) -> str:
    user_id = runtime.state["user_id"]          # 访问 state
    runtime.store.put(namespace=..., key=user_id, value=...)  # 访问 store
  • runtime.state — 获取当前 CustomAgentState 中的所有字段（包括 user_id）。
  • runtime.store — 获取 create_agent() 时传入的 Store 实例，直接调用 put/get/search。

invoke 传参：

agent.invoke({
    "messages": [...],
    "user_id": "user-1"   # 自定义字段，会存入 CustomAgentState
})
  • 除默认的 messages 外，可传入自定义字段，这些字段会保存在 CustomAgentState 中供工具读取。

⚠️ 常见问题与避坑指南

  1. 自定义 AgentState 必须传入 create_agent：state_schema=CustomAgentState 必须显式指定，否则 Agent 只会使用默认的 AgentState，自定义字段无法传递。
  2. user_id 需在每次 invoke 时传入：因为 InMemoryStore 中的 state 不持久化，每次 invoke 均需传递 user_id。
  3. ToolRuntime 是工具函数的特殊参数：不需要在工具装饰器中声明，LangGraph 会自动注入。参数名固定为 runtime。
  4. 工具返回值必须是字符串：工具函数返回 str 类型，Agent 会将返回值作为工具调用结果传给模型。
  5. 工具必须添加 docstring：通过 pass_docstring=True 或良好的 docstring 帮助 Agent 理解工具用途，决定何时调用。
  6. get 时检查 item 是否为 None：namespace 或 key 不存在时 get() 返回 None，需要判空处理。

💡 个人总结与延伸

  • 通过 ToolRuntime 在工具中访问 Store，是 Agent 集成长期记忆的标准模式。
  • 自定义 AgentState 扩展字段为 Agent 提供了"上下文参数"的传递通道（如 user_id、session_id 等），避免了在消息中混杂元数据。
  • 除了 put 和 get，同样可以在工具中使用 search() 实现更复杂的长期记忆检索逻辑。
  • 下一节将演示同样的逻辑切换到 PostgresStore，实现数据的持久化存储。
最近更新:
Contributors: yangwenguang