a tool for langgraph

[succulen-plants]

🚀 The feature

I want to use the memory unit as a tool for langgraph, similar to a database tool. How can I build it?
I tried this, but the effect was completely different from what I expected.

Motivation, pitch

mem0 = Memory()

class State(TypedDict):
    messages: Annotated[list, add_messages]

graph_builder = StateGraph(State)



@tool
def get_mem0(state: State) -> str:
    """
    Please check whether the memory unit contains the user's information, age, gender, hobbies, characteristics, etc. 
    Use this feature before personalizing communications.
    example: m.search(query="What are Alice's hobbies?", user_id="alice")
    Must carry user input
    """
    print(f"get_mem0===State content: {state}")  
    messages = state["messages"]
    user_id = 1
    relevant_memories = mem0.search(messages[-1]['content'], user_id,limit=4)
    # relevant_memories = mem0.search('我今年几岁了', user_id,limit=4)
    return relevant_memories

@tool
def set_mem0(state: State) -> str:
    """
    Please call this method when you need to record user information, and ensure the integrity of the user information and do not miss any information.
    example：m.add("I am working on improving my tennis skills. Suggest some online courses.", user_id="alice", metadata={"category": "hobbies"})
    """
    print(f"set_mem0===State content: {state}")  
    messages = state["messages"]
    user_id = 1
    print("messages====", messages)
    print("user_id====", user_id)
    # relevant_memories = mem0.search(messages[-1].content, user_id,limit=4)
    mem0.add(messages[-1]['content'], user_id=user_id,metadata={"category": "hobbies"})
    return '添加成功'

tool = TavilySearchResults(max_results=2)
tools = [tool,get_mem0]

from langchain_openai import ChatOpenAI

llm = ChatOpenAI()
llm_with_tools = llm.bind_tools(tools)

def chatbot(state: State):
    print(f"chatbot===State content: {state}") 
    # state = {**state, "user_info": passenger_id}
    return {"messages": [llm_with_tools.invoke(state["messages"])]}


graph_builder.add_node("chatbot", chatbot)

tool_node = ToolNode(tools=[tool,get_mem0,set_mem0])
graph_builder.add_node("tools", tool_node)

graph_builder.add_conditional_edges(
    "chatbot",
    tools_condition,
)
# Any time a tool is called, we return to the chatbot to decide the next step
graph_builder.add_edge("tools", "chatbot")
graph_builder.set_entry_point("chatbot")
graph = graph_builder.compile()

def stream_graph_updates(user_input: str):
    state = {"messages": [HumanMessage(content=user_input)]}
    for event in graph.stream(state):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)

while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break

        stream_graph_updates(user_input)
    except:
        # fallback if input() is not available
        user_input = "What do you know about LangGraph?"
        print("User: " + user_input)
        stream_graph_updates(user_input)
        break