LangChain V1.0 核心解析(四):Short-term memory组件——让AI记住对话:上下文管理与历史追溯技术

47 阅读16分钟

Agent 记忆管理

概述

短期记忆让你的应用程序能够记住单个线程或对话中的先前交互。

核心要素

  • State (状态): 定义用来存储消息的结构(通常是 MessagesState
  • Checkpointer (检查点保存器): 负责在每一步结束后把状态保存下来(短期记忆通常用 MemorySaver
  • Thread ID (线程ID): 在调用时通过 config 传入,用来隔离不同用户的对话上下文

错误认知

  • 内存 = 短期记忆
  • 数据库 = 长期记忆

正确标准

  • 短期记忆:数据与会话(thread进程)生命周绑定,随会话结束而被清理或遗忘
  • 长期记忆:数据与用户/业务实体生命周期绑定,跨会话持久保留并可主动检索

短期记忆

短期记忆通过LangGrap的AgentState(一个TypedDict)来管理。对话历史、中间步骤等信息被保存在状态中,并通过检查点(Checkpoints)机制在每次迭代后持久化。这使得长对话和失败恢复成为可能。

Checkpointer机制

这是LangGrap记忆的灵魂

  • 不加这一行:Agent是无状态的。每次invoke都是全新的开始
  • 加上一行:LangGraph会在每一步执行后,把state序列化并存入MemorySaver
  • 原理:当你再次invoke并传入thread_id时,LangGraph会先去内存里查这个 ID上次停在哪里?状态是什么?然后加载转改,把消息append进去,在继续运行

InMemorySaver()内存记忆管理

import os
from langchain.agents import create_agent
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver
from langchain_core.messages import HumanMessage, SystemMessage, trim_messages

# ============ 初始化 LLM ============
model = load_chat_model(model="gpt-4o-mini",provider="openai")

# ============ 定义工具函数 ============
@tool
def get_user_info(name: str) -> str:
    """查询用户信息,返回姓名、年龄、爱好"""
    user_db = {
        "陈明": {"age": 28, "hobby": "旅游、滑雪、喝茶"},
        "张三": {"age": 32, "hobby": "编程、阅读、电影"}
    }
    info = user_db.get(name, {"age": "未知", "hobby": "未知"})
    return f"姓名: {name}, 年龄: {info['age']}岁, 爱好: {info['hobby']}"

# ============ 1. 基础短期记忆:InMemorySaver ============
"""
开发环境使用内存存储,重启后记忆丢失。
关键参数:
- checkpointer: 记忆存储对象
- thread_id: 会话唯一标识(用户级隔离)
"""
def demo_inmemory_memory():
    print("=" * 60)
    print("场景 1: 内存记忆(开发环境)")
    print("=" * 60)

    # 创建内存检查点
    memory = InMemorySaver()

    # 创建 Agent(自动继承对话记忆能力)
    agent = create_agent(
        model=model,
        tools=[get_user_info],
        checkpointer=memory  # 启用短期记忆
    )

    # 配置:thread_id 作为会话 ID
    config = {"configurable": {"thread_id": "user_123"}}

    # 第一轮对话:自我介绍
    response1 = agent.invoke(
        {"messages": [{"role": "user", "content": "你好,我叫陈明,好久不见!"}]},
        config=config
    )
    print(f"用户:你好,我叫陈明,好久不见!")
    print(f"AI: {response1['messages'][-1].content}")
    print("-" * 40)

    # 第二轮对话:测试记忆
    response2 = agent.invoke(
        {"messages": [{"role": "user", "content": "请问你还记得我叫什么名字吗?"}]},
        config=config  # 使用相同 thread_id,自动携带上下文
    )
    print(f"用户:请问你还记得我叫什么名字吗?")
    print(f"AI: {response2['messages'][-1].content}")
    print("-" * 40)

    # 验证记忆状态
    state = agent.get_state(config)
    print(f"当前记忆轮次: {len(state.values['messages'])} 条消息")

    # 新开一个会话(不同 thread_id)
    config2 = {"configurable": {"thread_id": "user_456"}}
    response3 = agent.invoke(
        {"messages": [{"role": "user", "content": "我们之前聊过吗?"}]},
        config=config2
    )
    print(f"新会话 AI: {response3['messages'][-1].content}")  # 应无记忆

demo_inmemory_memory()

PostgresSaver()数据库持久化记忆

PostgresSaver 即使存储到数据库,仍然属于短期记忆。仍然属于短期记忆的原因

  • 作用域限制:它只检索和加载当前 thread_id 的数据
  • 生命周期管理:默认不会主动清理,但数据语义上属于"本次会话"
  • 无跨会话检索能力:无法在新会话中自动访问旧会话数据(除非手动指定旧 thread_id)

pip install langgraph-checkpoint-postgres

from langgraph.checkpoint.postgres import PostgresSaver

"""
生产环境使用数据库存储,支持:
- 持久化(重启不丢失)
- 多实例共享(分布式部署)
- 大规模并发
"""

print("\n" + "=" * 60)
print("场景 2: Postgres 持久化记忆(生产环境)")
print("=" * 60)

# 定义工具函数:查询用户信息
@tool
def get_user_info(name: str) -> str:
    """查询用户信息,返回姓名、年龄、爱好"""
    user_db = {
        "陈明": {"age": 28, "hobby": "旅游、滑雪、喝茶"},
        "张三": {"age": 32, "hobby": "编程、阅读、电影"}
    }
    info = user_db.get(name, {"age": "未知", "hobby": "未知"})
    return f"姓名: {name}, 年龄: {info['age']}岁, 爱好: {info['hobby']}"

# 创建模型
model = load_chat_model(model="gpt-4o-mini",provider="openai")

# 数据库连接字符串
DB_URI = "postgresql://myuser:123456@localhost:5432/mydatabase"

# 使用上下文管理器确保连接正确关闭
with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
    # 自动创建表结构(仅首次运行)
    checkpointer.setup()

    # 创建智能体
    agent = create_agent(
        model=model,
        tools=[get_user_info],
        checkpointer=checkpointer
    )

    # 配置线程 ID(用于区分不同用户)
    config = {"configurable": {"thread_id": "production_user_001"}}

    # 模拟用户注册流程
    agent.invoke(
        {"messages": [{"role": "user", "content": "我是新用户张三,请记录我的信息"}]},
        config=config
    )

    response = agent.invoke(
        {"messages": [{"role": "user", "content": "我是谁?"}]},
        config=config
    )
    print(f"AI: {response['messages'][-1].content}")

image.png

上下文裁剪

启用短期记忆后,长对话可能会超出LLM的上下文窗口

  • 需要配合 trim_messages 使用
  • start_on="human": 这是一个很细节的最佳实践。如果截断导致第一条消息是 AI 的回复(没有对应的 User 问题),某些模型会感到困惑。这个参数确保截断后的对话总是以 User 开始。
# ============ 核心模块导入 ============
import os
from langchain.chat_models import init_chat_model
from langchain.agents import create_agent
from langchain_core.tools import tool
from langgraph.checkpoint.memory import InMemorySaver  # 短期记忆存储
from langchain_core.messages import trim_messages  # 消息裁剪工具
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage

# ============ 初始化 LLM 与工具 ============
model = init_chat_model(
    model="gpt-4o-mini",
    model_provider="openai",
    temperature=0.7
)

# 1. 定义天气查询工具
@tool
def search_weather(city: str) -> str:
    """查询城市天气"""
    return f"{city}天气:晴,25°C"

# ============ 配置裁剪参数 ============
MAX_TOKENS = 100  # 根据模型上下文长度调整,上下文128K的话,一般设置为4000左右
TRIM_STRATEGY = "last"  # 保留最新消息
INCLUDE_SYSTEM = True  # 系统消息不参与裁剪

"""
tiktoken 是 OpenAI 官方 token 编码库,支持精确计数。
不同模型需使用不同编码:
- gpt-4o, gpt-4o-mini: o200k_base
- gpt-4-turbo: cl100k_base
- text-davinci-003: p50k_base
"""
import tiktoken

# 2. 定义 Token 编码器获取函数
def get_token_encoder(model_name: str = "gpt-4o-mini"):
    """获取 tiktoken 编码器实例"""
    try:
        # 自动映射模型到编码器
        encoding = tiktoken.encoding_for_model(model_name)
        print(f"✅ 已加载模型 '{model_name}' 的 tiktoken 编码器")
        return encoding
    except KeyError:
        # 如果模型不在官方列表,使用默认编码
        print(f"⚠️  模型 '{model_name}' 未在映射表中,使用默认编码 'o200k_base'")
        return tiktoken.get_encoding("o200k_base")

# 3. 初始化编码器(全局复用提升性能)
TOKEN_ENCODER = get_token_encoder("gpt-4o-mini")

# ============ 2. 精确 Token 计数函数 ============
def count_tokens_tiktoken(messages):
    """
    使用 tiktoken 精确计算消息列表的 token 总数

    参数:
        messages: List[BaseMessage] - 消息对象列表

    返回:
        int: 总 token 数

    说明:
        - 每个消息包含角色、内容和元数据,需完整编码
        - 不同消息格式(Human/AI/System)的 token 开销不同
        - 此函数模拟真实 API 调用的 token 计数逻辑
    """
    total_tokens = 0

    # 遍历每条消息,累加 token 数
    for message in messages:
        # 消息格式:role + content + 特殊标记
        # 典型格式:<|im_start|>role<|im_end|>content<|im_end|>

        # 角色 token(如 "user", "assistant", "system")
        role_tokens = len(TOKEN_ENCODER.encode(message.type))

        # 内容 token(消息正文)
        content_tokens = len(TOKEN_ENCODER.encode(message.content))

        # 格式开销(OpenAI 消息边界标记)
        # 每条约 4 个特殊 token(开始、角色、结束、内容结束)
        format_overhead = 4

        total_tokens += role_tokens + content_tokens + format_overhead

    return total_tokens

# ============ 创建Agent ============
def create_trimmed_agent():
    """创建 Agent 并配置 InMemorySaver"""
    memory = InMemorySaver()

    agent = create_agent(
        model=model,
        tools=[search_weather],
        system_prompt="你是一个简洁的助手,记住用户提到的城市名称",
        checkpointer=memory  # 启用短期记忆
    )

    return agent

# ============ 手动裁剪并调用 Agent ============
def invoke_with_trim(agent, user_input: str, config: dict):
    """
    在调用 Agent 前手动裁剪上下文

    流程:
    1. 获取当前状态(所有历史消息)
    2. 使用 trim_messages 裁剪
    3. 构建新输入(裁剪后的消息 + 新消息)
    4. 调用 Agent
    """
    # 1. 获取当前记忆状态
    state = agent.get_state(config)
    existing_messages = state.values.get("messages", []) if state else []

    # 精确计算当前 token 数
    current_tokens = count_tokens_tiktoken(existing_messages)

    # 2. 如果有历史消息,先裁剪
    if existing_messages:
        print(f"裁剪前消息数: {len(existing_messages)}")

        # 核心:调用 trim_messages 进行裁剪
        trimmed_messages = trim_messages(
            existing_messages,                    # 待裁剪的消息列表
            max_tokens=MAX_TOKENS,                # 允许的最大 token 数,超过则触发裁剪
            token_counter=count_tokens_tiktoken,  # token 计数函数
            strategy=TRIM_STRATEGY,               # 裁剪策略,"last" 保留最新消息,"first" 保留最早消息
            include_system=INCLUDE_SYSTEM,        # 是否保留系统消息(通常必须保留)
            allow_partial=False,                  # False不允许部分消息,会尝试保留消息的完整性。如果无法在保持消息完整性的前提下将总token数裁剪到参数 max_tokens 以内,就会返回 空列表
            start_on="human"                      # 从 human 消息开始裁剪
        )

        # 计算裁剪后的 token 数
        new_tokens = count_tokens_tiktoken(trimmed_messages)

        print(f"裁剪后 token: {new_tokens},裁剪后消息数: {len(trimmed_messages)}")
    else:
        trimmed_messages = []

    # 3. 添加新消息
    new_messages = trimmed_messages + [HumanMessage(content=user_input)]

    # 4. 调用 Agent(checkpointer 会自动保存新状态)
    response = agent.invoke(
        {"messages": new_messages},
        config=config
    )

    return response

# ============ 演示多轮对话与裁剪 ============
def demo_manual_trim():
    """演示手动裁剪上下文的多轮对话"""
    print("=" * 60)
    print("场景:手动 trim_messages + InMemorySaver")
    print("=" * 60)

    agent = create_trimmed_agent()
    config = {"configurable": {"thread_id": "trim_user_001"}}

    # 模拟多轮对话
    conversations = [
        "你好,我叫陈明",
        "查询北京天气",
        "上海呢?",
        "明天北京天气如何?",  # 此时会触发裁剪
        "我是谁?",  # 测试记忆是否保留
    ]

    for i, query in enumerate(conversations, 1):
        print(f"\n--- 第 {i} 轮 ---")
        print(f"用户: {query}")

        # 每次调用前自动裁剪
        response = invoke_with_trim(agent, query, config)

        print(f"AI: {response['messages'][-1].content}")
        print("-" * 40)

# ============ 运行演示 ============
if __name__ == "__main__":
    demo_manual_trim()

自定义State扩展

在 LangGraph 中,AgentState 是一个 TypedDict,定义了 Agent 执行过程中流转的数据结构。扩展 State = 在基础结构上增加自定义字段,用于携带更多上下文和业务数据。

扩展 State 的本质:在 LangGraph 中,State 是 Agent 的 "内存" 和 "消息总线" ,扩展它就像给程序增加新的全局变量,但这些变量随执行流自动流转、隔离、持久化,是实现复杂 Agent 逻辑的基础。

扩展State核心目的

  • 跨步骤持久化上下文:Agent 执行是多步骤的(LLM调用 → 工具调用 → 结果解析),扩展的 State 字段能在所有步骤间共享。
  • 实现条件分支与动态路由:根据 State 中的字段值,决定 Agent 的下一步走向。
  • 支持多模态与复杂输入:现代 Agent 需要处理图片、文件等非文本数据,扩展到 State 中。
  • 实现记忆与持久化:扩展字段用于存储长期记忆,跨会话保持。
  • 性能监控与调试:扩展字段用于记录性能指标,便于分析优化。
class ExtendedState(TypedDict):
    messages: list[BaseMessage]
    user_id: str           # 扩展:用户身份,用于权限控制和个性化
    session_id: str        # 扩展:会话标识,用于对话历史管理
    retry_count: int       # 扩展:重试次数,用于错误处理策略
    original_query: str    # 扩展:原始查询,用于日志和审计
# ============ 自定义 State 扩展 ============
"""
通过 TypedDict 扩展 AgentState,添加业务字段(用户ID、偏好等)。
LangChain 1.0 推荐使用 TypedDict 而非 Pydantic。
"""
from typing import TypedDict, Optional
from langchain.agents import AgentState, create_agent
from langgraph.checkpoint.memory import InMemorySaver

# 定义自定义 State 结构
class CustomAgentState(AgentState):
    """扩展的 Agent 状态,包含业务上下文"""
    user_id: str  # 用户唯一标识
    preferences: dict  # 用户偏好(主题、语言等)
    visit_count: int  # 访问次数

# ============ 定义带状态访问的工具 ============
from langchain.tools import ToolRuntime
from langgraph.types import Command
from langchain.messages import ToolMessage

# 定义工具函数:更新用户偏好
@tool
def update_user_preference(runtime: ToolRuntime, theme: str) -> Command:
    """
    更新用户主题偏好,写入短期记忆

    ToolRuntime 提供对 state 和 context 的访问能力:
    - runtime.state: 当前状态(含自定义字段)
    - runtime.context: 调用上下文
    - runtime.tool_call_id: 工具调用ID
    """
    # 从当前状态获取偏好(如果不存在则初始化)
    current_prefs = runtime.state.get("preferences", {})
    current_prefs["theme"] = theme

    # 返回 Command 对象,指示状态更新
    return Command(update={
        "preferences": current_prefs,
        "messages": [
            ToolMessage(
                content=f"成功更新主题为: {theme}",
                tool_call_id=runtime.tool_call_id
            )
        ]
    })

# 定义工具函数:根据用户偏好生成问候
@tool
def greet_user(runtime: ToolRuntime) -> str:
    """根据用户偏好生成个性化问候"""
    user_name = runtime.state.get("user_id", "访客")
    prefs = runtime.state.get("preferences", {})
    theme = prefs.get("theme", "默认")

    return f"欢迎回来,{user_name}!当前主题: {theme}"

# ============ 创建带自定义状态的 Agent ============
def demo_custom_state():
    print("\n" + "=" * 60)
    print("场景 3: 自定义 State 扩展记忆维度")
    print("=" * 60)

    # 使用内存存储
    checkpointer = InMemorySaver()

    # 创建 Agent,指定自定义 state_schema
    agent = create_agent(
        model=model,
        tools=[update_user_preference, greet_user],
        state_schema=CustomAgentState,  # 关键:传入自定义状态类型
        checkpointer=checkpointer
    )

    # 配置线程 ID(用于区分不同用户)
    config = {"configurable": {"thread_id": "custom_state_user"}}

    # 第一轮:初始化用户信息
    result1 = agent.invoke(
        {
            "messages": [{"role": "user", "content": "设置主题为暗黑模式"}],
            "user_id": "user_789",  # 自定义字段
            "preferences": {"language": "zh-CN"},  # 初始偏好
            "visit_count": 1
        },
        config=config
    )
    print(f"第一轮: {result1['messages'][-1].content}")
    print("-" * 40)

    # 第二轮:读取记忆
    result2 = agent.invoke(
        {"messages": [{"role": "user", "content": "打个招呼"}]},
        config=config
    )
    print(f"第二轮: {result2['messages'][-1].content}")
    print("-" * 40)

    # 查看完整状态
    state = agent.get_state(config)
    print("当前记忆状态:")
    print(f"  用户ID: {state.values.get('user_id')}")
    print(f"  偏好: {state.values.get('preferences')}")
    print(f"  消息数: {len(state.values['messages'])}")

# ============ 运行自定义状态示例 ============
if __name__ == "__main__":
    demo_custom_state()

核心拓展-访问记忆

在工具中读取短期记忆

使用 ToolRuntime 参数在工具中访问短期记忆(状态)。tool_runtime 参数在工具签名中是隐藏的(因此模型看不到),但工具可以通过它访问状态。

from langchain.agents import create_agent, AgentState
from langchain.tools import tool, ToolRuntime


class CustomState(AgentState):
    user_id: str

@tool
def get_user_info(
    runtime: ToolRuntime
) -> str:
    """Look up user info."""
    user_id = runtime.state["user_id"]
    return "User is John Smith" if user_id == "user_123" else "Unknown user"

agent = create_agent(
    model="gpt-5-nano",
    tools=[get_user_info],
    state_schema=CustomState,
)

result = agent.invoke({
    "messages": "look up user information",
    "user_id": "user_123"
})
print(result["messages"][-1].content)
# > User is John Smith.

通过工具写入短期记忆

from langchain.tools import tool, ToolRuntime
from langchain_core.runnables import RunnableConfig
from langchain.messages import ToolMessage
from langchain.agents import create_agent, AgentState
from langgraph.types import Command
from pydantic import BaseModel


class CustomState(AgentState):  
    user_name: str

class CustomContext(BaseModel):
    user_id: str

@tool
def update_user_info(
    runtime: ToolRuntime[CustomContext, CustomState],
) -> Command:
    """Look up and update user info."""
    user_id = runtime.context.user_id
    name = "John Smith" if user_id == "user_123" else "Unknown user"
    return Command(update={  
        "user_name": name,
        # update the message history
        "messages": [
            ToolMessage(
                "Successfully looked up user information",
                tool_call_id=runtime.tool_call_id
            )
        ]
    })

@tool
def greet(
    runtime: ToolRuntime[CustomContext, CustomState]
) -> str | Command:
    """Use this to greet the user once you found their info."""
    user_name = runtime.state.get("user_name", None)
    if user_name is None:
       return Command(update={
            "messages": [
                ToolMessage(
                    "Please call the 'update_user_info' tool it will get and update the user's name.",
                    tool_call_id=runtime.tool_call_id
                )
            ]
        })
    return f"Hello {user_name}!"

agent = create_agent(
    model="gpt-5-nano",
    tools=[update_user_info, greet],
    state_schema=CustomState, 
    context_schema=CustomContext,
)

agent.invoke(
    {"messages": [{"role": "user", "content": "greet the user"}]},
    context=CustomContext(user_id="user_123"),
)

长期记忆

长期记忆通过与外部向量数据库或键值存储集成来实现。可以在Agent执行的关键节点(如对话结束时)提取关键信息、用户偏好等,并存入长期记忆库,供未来的对话使用。

语义检索与向量数据库

  • 向量数据库是实现长期记忆的核心技术,通过语义相似度搜索实现知识的长期存储和检索。

技术实现:

  • 向量化存储:将对话内容、用户偏好等转换为向量表示
  • 语义检索:基于向量相似度实现智能搜索
  • 多模态支持:支持文本、图像、音频等多种数据类型
  • 高性能查询:支持大规模数据的快速检索

典型实现:

  • Milvus:开源向量数据库,支持大规模向量检索
  • Qdrant:高性能向量搜索引擎
  • Pinecone:云原生向量数据库服务
  • Chroma: 轻量级向量数据库,可本地持久化

pip install langchain-chroma

import os
import uuid
from typing import List

# --- 1. 导入组件 ---
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.tools import tool
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_chroma import Chroma # 向量数据库
from langchain_core.documents import Document
from langchain.agents import AgentState, create_agent
from langgraph.checkpoint.memory import MemorySaver

# 确保配置了 OPENAI_API_KEY
# os.environ["OPENAI_API_KEY"] = "sk-..."

# ==========================================
# 2. 初始化向量数据库 (长期记忆的物理载体)
# ==========================================
# 在生产环境中,这里应该是连接到 Pinecone, Milvus 或本地持久化的 Chroma
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vector_store = Chroma(
    collection_name="agent_long_term_memory",
    embedding_function=embeddings,
    #persist_directory="./chroma_db" # 如果想存到硬盘,取消注释这一行
)

# ==========================================
# 3. 定义记忆工具 (Agent 的手)
# ==========================================
# 3.1 定义记忆保存工具
# ==========================================
@tool
def save_memory(content: str):
    """
    将重要信息保存到长期记忆中。
    当你获知用户的喜好、职业、计划或其他长期有效的事实时,调用此工具。
    参数:
        content (str): 要保存的记忆内容。
    """
    print(f"\n[记忆操作] 正在保存记忆: '{content}'")
    # 将文本封装为 Document
    doc = Document(
        page_content=content,
        metadata={"source": "user_interaction", "timestamp": "simulated_time"}
    )
    # 写入向量库
    vector_store.add_documents([doc])
    return "记忆已成功保存。"

# 3.2 定义记忆搜索工具
# ==========================================
@tool
def search_memory(query: str):
    """
    从长期记忆中搜索相关信息。
    当你被问及关于用户过去的问题,或者你不确定答案时,使用此工具进行查找。
    参数:
        query (str): 要搜索的查询语句。
    """
    print(f"\n[记忆操作] 正在搜索记忆: '{query}'")

    # 执行语义搜索 (k=2 表示只取最相关的2条)
    results = vector_store.similarity_search(query, k=2)

    if not results:
        return "没有找到相关的记忆。"

    # 将搜索结果拼接成字符串返回给 Agent
    memory_content = "\n".join([f"- {doc.page_content}" for doc in results])
    return f"找到以下相关记忆:\n{memory_content}"

# 将工具放入列表
tools = [save_memory, search_memory]

# ==========================================
# 4. 创建 Agent
# ==========================================

# 定义系统提示词:教会 Agent 何时使用记忆工具
SYSTEM_PROMPT = """你是一个拥有长期记忆的私人助手。
你的目标是记住用户的喜好和重要信息,以便提供个性化服务。

1. 如果用户告诉你任何关于他们自己的事实(如名字、喜好、居住地),请务必调用 'save_memory' 工具保存。
2. 如果用户问你一个问题,而答案可能在你之前的记忆中,请先调用 'search_memory' 工具查找。
3. 如果只是闲聊,不需要调用工具。
"""

llm = ChatOpenAI(model="gpt-4o", temperature=0) # 建议使用 GPT-4 或更强的模型以保证工具调用准确率

# 使用 checkpointer 依然是必要的,用于维持当前这一轮对话的上下文
checkpointer = MemorySaver()

# 创建 Agent 应用
agent_app = create_agent(
    llm,
    tools,
    system_prompt=SYSTEM_PROMPT, # 注入系统提示词
    checkpointer=checkpointer
)

# ==========================================
# 5. 运行演示
# ==========================================

def run_demo():
    # === 场景 A:存入记忆 ===
    # 使用一个 thread_id,代表这是今天的对话
    config_a = {"configurable": {"thread_id": "session_today"}}

    print("--- 场景 A:用户告诉 Agent 喜好 ---")
    user_input_1 = "你好,记住我最喜欢的水果是草莓,而且我对花生过敏。"

    # 运行 Agent,stream_mode="values"参数,返回每个时间步的中间结果
    for chunk in agent_app.stream({"messages": [HumanMessage(content=user_input_1)]}, config=config_a, stream_mode="values"):
        # 只打印最后一条机器人的回复
        pass
    print(f"Agent: {chunk['messages'][-1].content}")

    # === 场景 B:模拟遗忘 (开启新线程) ===
    # 我们换一个 thread_id,这意味着 Agent 失去了“短期记忆” (MemorySaver 里的东西访问不到了)
    # 但是,长期记忆在 VectorStore 里,是可以跨 thread 访问的!
    config_b = {"configurable": {"thread_id": "session_tomorrow"}}

    print("\n--- 场景 B:第二天 (新的 Session,短期记忆已清空) ---")
    user_input_2 = "我想吃点零食,但我忘了我有什么忌口,你能帮我查查吗?"

    print(f"User: {user_input_2}")

    # 观察控制台输出,你会看到 Agent 自动调用 search_memory
    final_response = None
    for chunk in agent_app.stream({"messages": [HumanMessage(content=user_input_2)]}, config=config_b, stream_mode="values"):
        final_response = chunk['messages'][-1]

    print(f"Agent: {final_response.content}")

if __name__ == "__main__":
    run_demo()