【python 用LangGraph+Mem0创建个性化体验的AI Agent】

推荐原文链接：当LangGraph遇上Mem0：如何让你的AI Agent具有更智能的记忆与个性化的体验？

智能体：一个带有网络搜索功能的简单对话机器人。但这次我们增加了长期记忆选项，用来实现跨多次会话、多用户甚至多Agent的个性化交互能力。

【创建记忆】

创建一个Mem0的Memory对象，向量数据库使用嵌入式的Chroma，LLM使用OpenAI的gpt-4o-mini模型（OpenAI的Key在环境变量配置）：

from mem0 import Memory
from typing import Annotated, TypedDict, List
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI
from langchain_core.messages import SystemMessage, HumanMessage, AIMessage
from langchain_community.tools.tavily_search import TavilySearchResults
from langgraph.prebuilt import ToolNode

#个性化长期记忆，采用chrona向量库存储
mem0 = Memory.from_config({
    "version":"v1.1",
    "llm": {
        "provider": "openai",
        "config": {
            "model": "gpt-4o-mini"
        }    },
    "vector_store": {
        "provider": "chroma",
        "config": {
            "collection_name": "chat_memory",
            "path": "./db",
        }
    }
})

【创建LangGraph工作流与Agent】

定义LangGraph中的节点行为，最后创建工作流，并添加节点与边。这里针对Mem0的主要修改集中在chatbot这个节点方法中。简单解释如下：

• 为了区分不同用户的记忆，会在state中保留一个mem0_user_id。在添加记忆或者检索记忆时都需要携带这个user_id。（在实际应用中，这个id很可能是你的某个客户ID）

• 在chatbot回复之前，先根据输入消息内容检索关联的个性化记忆(search)，并把检索到的记忆组装成System Message。（这是一个常见的优化点，即如何检索出更相关的记忆，可以参考RAG优化中的一些方法）

• 将System Message与用户消息一起作为上下文输入给LLM，从而在生成响应时，大模型能够根据个性化的记忆作出响应。

• 在本轮交互结束后，调用add接口将该用户的本次对话信息添加到记忆中。Mem0会自动识别和合并，以用于下次检索。

#定义LangGraph的State
class State(TypedDict):
    messages: Annotated[List[HumanMessage | AIMessage], add_messages]
    mem0_user_id: str

# 调用搜索引擎的工具节点，利用ToolNode构建
tools = [TavilySearchResults(max_results=1)]
tool_node = ToolNode(tools)

#定义chatbot节点
def chatbot(state: State):
    messages = state["messages"]
    user_id = state["mem0_user_id"]

    # 取出关联的个性化记忆，并组装成context，放在system message中
    memories = mem0.search(messages[-1].content, user_id = user_id)["results"]
    context = "历史对话中的相关信息有：\n"
    for memory in memories:
        context += f"- {memory['memory']}\n"
    system_message = SystemMessage(content=f"你是一个乐于助人的客户支持助手。利用所提供的上下文来个性化你的回复，并会记住用户的偏好和过去的交互。\n{context}")

    # 组装消息，并调用LLM（注意绑定tools）
    full_messages = [system_message] + messages
    llm = ChatOpenAI(model="gpt-4o-mini") .bind_tools(tools)
    response = llm.invoke(full_messages)

    # 记住本地对话的信息
    mem0.add(f"User: {messages[-1].content}\nAssistant: {response.content}", user_id=user_id)
    return {"messages": [response]}

# 一个辅助方法：判断是否需要调用工具
def should_continue(state):
    messages = state["messages"]
    last_message = messages[-1]

    #根据大模型的反馈来决定是结束，还是调用工具
    if not last_message.tool_calls:
        return "end"
    else:
        return "continue"

# 定义一个graph
workflow = StateGraph(State)
workflow.add_node("llm", chatbot)
workflow.add_node("search", tool_node)
workflow.set_entry_point("llm")

# 一个条件边，即从agent出来的两个分支及条件
workflow.add_conditional_edges(
    "llm",
    should_continue,
    {
        "continue": "search",
        "end": END,
    },
)

# action调用后返回agent
workflow.add_edge("search", "llm")
graph = workflow.compile()

【测试Agent的记忆能力】

if __name__ == "__main__":
    print("AI: 你好！有什么可以帮助你？")
    mem0_user_id = "testuser"  # You can generate or retrieve this based on your user management system
    while True:
        user_input = input("输入: ")
        if user_input.lower() in ['quit', 'exit', 'bye']:
            break

        config = {"configurable": {"thread_id": mem0_user_id}}
        state = {"messages": [HumanMessage(content=user_input)], "mem0_user_id": mem0_user_id}

        response = graph.invoke(state,config)
        print("AI: ",response["messages"][-1].content)