LangGraph--带记忆和工具的聊天机器人

具体情况上一篇文章，这里就不啰嗦了，直接代码了

'''
    给聊天机器人添加记忆功能
'''

from typing import Annotated

from langchain_deepseek import ChatDeepSeek
from langchain_tavily import TavilySearch
from langchain_core.messages import BaseMessage
from typing_extensions import TypedDict

from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition

from langgraph.checkpoint.memory import MemorySaver

import os 
from dotenv import load_dotenv

# 加载.env文件中的环境变量
load_dotenv()

# 定义状态类型，使用 add_messages 注解来自动合并消息列表
class State(TypedDict):
    messages:Annotated[list, add_messages]


graph_builder = StateGraph(State)

llm = ChatDeepSeek(
    model="deepseek-chat",
    api_key=os.getenv("DEEPSEEK_API_KEY"))
# mess = llm.invoke(input = "你是谁？")

# 初始化搜索工具
tool = TavilySearch(max_results=2)
tools =[tool]

# 将工具绑定到llm中
llm_with_tools = llm.bind_tools(tools)
def chatbot(state:State):
    return {"messages":[llm_with_tools.invoke(state["messages"])]}

# 添加聊天机器人节点
graph_builder.add_node("chatbot",chatbot)
# 添加工具节点
tool_node = ToolNode(tools=[tool])
graph_builder.add_node("tools",tool_node)
# 添加条件边
graph_builder.add_conditional_edges("chatbot",tools_condition)

# 工具调用完成后，返回聊天机器人节点
graph_builder.add_edge("tools","chatbot")
graph_builder.set_entry_point("chatbot")

# 创建内存保存器
memory_saver = MemorySaver()

# 将内存保存器添加到图结构中
graph =graph_builder.compile(checkpointer=memory_saver)

# 打印图结构
print(graph.get_graph().draw_mermaid())
graph_png = graph.get_graph().draw_mermaid_png()
with open("chatbot_workflow.png", "wb") as f:
    f.write(graph_png)

# 定义对话线程ID
print("\n设置对话线程 ID = '1'...")
config = {"configurable": {"thread_id": "1"}}  # 使用线程ID来标识和区分不同的对话

# 示例 1: 第一次对话
print("\n示例 1: 第一次对话 - 用户介绍自己")
user_input = "Hi there! My name is Will."
print(f"\n用户输入: '{user_input}'")

# 注意: config 是 stream() 函数的第二个参数!
print("使用线程 ID '1' 调用图...")
events = graph.stream(
    {"messages": [{"role": "user", "content": user_input}]},
    config,  # 传递包含 thread_id 的配置
    stream_mode="values",
)

print("\n助理回应:")
for event in events:
    event["messages"][-1].pretty_print()  # 打印助理的回应

# 示例 2: 测试记忆功能
print("\n\n示例 2: 第二次对话 - 测试记忆功能")
user_input = "Remember my name?"
print(f"\n用户输入: '{user_input}'")

# 使用相同的线程ID再次调用图
print("使用相同的线程 ID '1' 再次调用图...")
events = graph.stream(
    {"messages": [{"role": "user", "content": user_input}]},
    config,  # 使用相同的配置，图将加载之前保存的状态
    stream_mode="values",
)

print("\n助理回应 (应该记得用户名字):")
for event in events:
    event["messages"][-1].pretty_print()

# 示例 3: 新对话线程
print("\n\n示例 3: 新对话线程 - 测试线程隔离")
print("创建新的线程 ID = '2'...")

# 使用不同的线程ID
print("使用新的线程 ID '2' 调用图...")
events = graph.stream(
    {"messages": [{"role": "user", "content": user_input}]},
    {"configurable": {"thread_id": "2"}},  # 使用新的线程ID
    stream_mode="values",
)

print("\n助理回应 (不应该记得用户名字):")
for event in events:
    event["messages"][-1].pretty_print()

# 示例 4: 返回第一个线程
print("\n\n示例 4: 返回第一个线程 - 验证记忆持久性")
print(f"\n用户输入: '{user_input}'")

# 再次使用第一个线程ID
print("再次使用线程 ID '1' 调用图...")
events = graph.stream(
    {"messages": [{"role": "user", "content": user_input}]},
    config,  # 使用原始线程ID
    stream_mode="values",
)

print("\n助理回应 (应该仍然记得用户名字):")
for event in events:
    event["messages"][-1].pretty_print()


print("done")

设置对话线程 ID = '1'...

示例 1: 第一次对话 - 用户介绍自己

用户输入: 'Hi there! My name is Will.'
使用线程 ID '1' 调用图...

助理回应:
================================ Human Message =================================

Hi there! My name is Will.
================================== Ai Message ==================================

Hi Will! It's great to meet you. How can I assist you today?


示例 2: 第二次对话 - 测试记忆功能

用户输入: 'Remember my name?'
使用相同的线程 ID '1' 再次调用图...

助理回应 (应该记得用户名字):
================================ Human Message =================================

Remember my name?
================================== Ai Message ==================================

Of course, Will! I'll remember your name for the rest of our conversation. How can I help you today? 😊


示例 3: 新对话线程 - 测试线程隔离
创建新的线程 ID = '2'...
使用新的线程 ID '2' 调用图...

助理回应 (不应该记得用户名字):
================================ Human Message =================================

Remember my name?
================================== Ai Message ==================================

I don't have the ability to remember personal details like names between interactions. However, you can tell me your name, and I'll happily use it for the rest of our current conversation! What should I call you?


示例 4: 返回第一个线程 - 验证记忆持久性

用户输入: 'Remember my name?'
再次使用线程 ID '1' 调用图...

助理回应 (应该仍然记得用户名字):
================================ Human Message =================================

Remember my name?
================================== Ai Message ==================================

Got it, Will! I'll keep your name in mind while we chat. Let me know how I can assist you! 😊

在这一节中，我们成功地为聊天机器人添加了记忆功能，使其能够记住对话历史并提供更连贯的回答。这是通过LangGraph的检查点功能实现的，它允许我们在每次交互后保存图的状态，并在后续交互中恢复该状态。这种能力对于构建真正智能的对话系统至关重要，因为它使AI能够维持上下文并提供个性化的用户体验。

我们学习了LangGraph的几个重要概念：

1. 检查点（Checkpointing） ：使用MemorySaver在内存中保存图的状态，为应用提供持久性。在生产环境中，可以替换为持久化存储解决方案。
2. 线程ID（Thread ID）：用于标识和区分不同的对话，实现多用户场景下的会话隔离。
3. 配置对象（Config）：在调用图时传递额外的配置信息，灵活控制执行行为。

这些功能共同构成了LangGraph强大的状态管理能力，为构建复杂的AI应用奠定了基础。