How to give your LLM agents long-term memory
Curious about how to replicate ChatGPT’s new functionality of remembering things in your own LangGraph agents?
Throughout this journey, we’ll walk you through this seamless integration of long-term memory for preventing repetition and enhancing future conversations.
From setting up nodes to fine-tuning memory tools, you’ll gain insight into optimizing agent communication and memory handling.
What is LangGraph?
We do have some other blog posts that explain it more in depth. You can check them out here, but let’s quickly go over it.
LangGraph is an extension of LangChain, whose main objective is creating agent and multi-agent projects. It adds the ability to create cyclical flows while providing developers a high degree of control, which is, obviously, very important for creating custom agents and flows.
That’s good enough for now. Let’s start coding!
Building the agent
In this section we’ll skip a few steps and focus more on the logic behind the agent, but if you want to check out the whole Jupyter notebook, it will also be available at the end of the blog post.
Step 1: create agents
To get started, let’s create helper functions to build our agents — in this case, our main agent and the long-term memory (LTM) agent. These will serve as nodes in the graph.
from langchain_core.messages import (
BaseMessage,
HumanMessage,
)
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langgraph.graph import END, StateGraph
def create_agent(llm, tools, system_prompt: str):
"""Create an agent."""
prompt = ChatPromptTemplate.from_messages(
[
("system", system_prompt),
MessagesPlaceholder(variable_name="messages"),
]
)
if len(tools) == 0:
return prompt | llm
return prompt | llm.bind_tools(tools)The create_agent function takes an LLM (language model), tools, and a system prompt to create an agent. It binds the tools to the agent if any are provided, allowing the agent to perform specific actions.
Step 2: define tools
Here we define the tools that our agents will use to interact with long-term memory (LTM). In this case, it’s a single tool with the ability to add/remove/update memories — ideally you’d modify the code to use a database like Redis.
import enum
from typing import Optional, Type
from langchain.pydantic_v1 import BaseModel, Field
from langchain.tools import BaseTool
from langchain.callbacks.manager import CallbackManagerForToolRun
LTM = []
class MemoryAction(enum.Enum):
ADD = "add"
REMOVE = "remove"
UPDATE = "update"
class ModifyMemoryInput(BaseModel):
memory: str = Field(..., description="the memory to add to the agent")
memory_old: str = Field(..., description="the memory to remove from the agent")
action: MemoryAction = Field(..., description="the action to perform on the memory")
class ModifyMemory(BaseTool):
name = "modify_memory"
description = "Modify the memory of the agent."
args_schema: Type[BaseModel] = ModifyMemoryInput
def _run(
self,
memory: str,
memory_old: str,
action: MemoryAction,
run_manager: Optional[CallbackManagerForToolRun] = None
):
if action == MemoryAction.ADD:
LTM.append(memory)
elif action == MemoryAction.REMOVE:
LTM.remove(memory_old)
elif action == MemoryAction.UPDATE:
LTM.remove(memory_old)
LTM.append(memory)
return LTMStep 3: define state
Here we define the state of the graph, which includes a list of messages and a key to track the most recent sender.
import operator
from typing import Annotated, Sequence, TypedDict
from langchain_openai import ChatOpenAI
from typing_extensions import TypedDict
class AgentState(TypedDict):
messages: Annotated[Sequence[BaseMessage], operator.add]
sender: strStep 4: define agent prompts
Now we need to define the agent prompts, one of the most important steps.
For this example, we’ll be creating a chef agent to wrap the LTM around, so we need to define its prompt here. For the LTM agent, the prompt is more straightforward and can be modified for different use cases if needed.
Note that we inject the memory in both prompts, so the chef agent can remember stuff, and the LTM agent can take action on deleting and updating older memories.
from datetime import datetime
chef_prompt = """{current_time}
You are an agent designed to assist in the kitchen, like a Chef.
You can help with recipes and cooking tips.
You can also help with meal planning and grocery shopping.
Current memory:
{current_memory}"""
chef_prompt = chef_prompt.format(
current_time=datetime.now().strftime("%d/%m/%Y %H:%M:%S"),
current_memory="\n".join([f"{i}: {memory}" for i, memory in enumerate(LTM)])
)
ltm_prompt = """{current_time}
You are an agent designed to assist with the storage of midterm and long-term memory for a user.
Evaluate incoming information to determine if it is relevant, such as a fact, event or preference.
If the information is relevant, store it in the memory. You may also update or delete information in the memory.
You have access to tools for interacting with the memory. (Insert/update/delete)
YOU MUST use the name "User" when referring to the user.
YOU MUST summarize the information before storing it in the memory.
Current memory:
{current_memory}"""
ltm_prompt = ltm_prompt.format(
current_time=datetime.now().strftime("%d/%m/%Y %H:%M:%S"),
current_memory="\n".join([f"{i}: {memory}" for i, memory in enumerate(LTM)])
)Step 5: define agent nodes
Now, we create the nodes for the agents.
So, using our previously created `create_agent` function, we define our agents by passing their respective models, tools, and prompts.
import functools
def agent_node(state, agent, name):
result = agent.invoke(state)
result.name = name
return {
"messages": [result],
"sender": name,
}
chef_agent = create_agent(
ChatOpenAI(model="gpt-4o", temperature=0),
[],
chef_prompt,
)
chef_node = functools.partial(agent_node, agent=chef_agent, name="chef")
ltm_agent = create_agent(
ChatOpenAI(model="gpt-4o", temperature=0),
[ModifyMemory()],
ltm_prompt,
)
ltm_node = functools.partial(agent_node, agent=ltm_agent, name="ltm")Step 6: define tool node
We also need to create a tool node that will take action when the agents decide. So, ideally it would have all the tools from all the agents, but in our case we only have the memory one.
from langgraph.prebuilt import ToolNode
tools = [ModifyMemory()]
tool_node = ToolNode(tools)Step 7: define graph
Now we just need to wrap everything up and connect the nodes with the correct “flow”.
For our use case, if needed, the chef agent can use call_tool, which will return the answer to the chef and so on until it’s done, and then it goes to the LTM node that follows that same pattern until it’s done.
from typing import Literal
def router(state) -> Literal["call_tool", "__end__", "continue"]:
messages = state["messages"]
last_message = messages[-1]
if last_message.tool_calls:
return "call_tool"
return "continue"
workflow = StateGraph(AgentState)
workflow.add_node("chef", chef_node)
workflow.add_node("ltm", ltm_node)
workflow.add_node("call_tool", tool_node)
workflow.add_conditional_edges(
"chef",
router,
{"continue": "ltm", "call_tool": "call_tool"},
)
workflow.add_conditional_edges(
"ltm",
router,
{"continue": END, "call_tool": "call_tool"},
)
workflow.add_conditional_edges(
"call_tool",
lambda x: x["sender"],
{
"chef": "chef",
"ltm": "ltm",
},
)
workflow.set_entry_point("chef")
graph = workflow.compile()Step 8: visualize graph
This step is optional, but it definitely helps understating the whole graph and its logic.
from IPython.display import Image, display
try:
display(Image(graph.get_graph(xray=True).draw_mermaid_png()))
except:
passFor our case, this code returns this graph, just as expected:
Step 9: invoke the graph
Finally, for our last step, we can invoke the graph with `graph.invoke()`.
There’s also some code to separate the messages from each agent, since you’re probably not going to want to store the LTM agent responses to your message history.
responses = graph.invoke(
{
"messages": [
HumanMessage(
content="Hey, can you help me with a seafood recipe? I'm allergic to shellfish.",
)
],
}
)
ltm_start = next((i for i, message in enumerate(responses['messages']) if message.name == 'ltm'))
agent_responses = responses['messages'][:ltm_start]
ltm_responses = responses['messages'][ltm_start:]
print("Agent responses:")
for response in agent_responses:
print(response)
print("\nLTM responses:")
for response in ltm_responses:
print(response)Agent responses:
content="Hey, can you help me with a seafood recipe? I'm allergic to shellfish."
content="Of course! How about a delicious Lemon Herb Grilled Salmon? It's..."
LTM responses:
content='' additional_kwargs={'tool_calls': [{'id': 'call_QhirSh7J46wKHZbLIrMQFhUf', 'function': {'arguments': '{"memory":"User is allergic to shellfish.","memory_old":"","action":"add"}', 'name': 'modify_memory'}, 'type': 'function'}]}
...So, just as expected, our chef answers with a recipe and our LTM agent calls the tool to add the memory — in our case, that the “User” is allergic to shellfish. I also cleaned a little bit of the output, so don’t be scared if your code outputs a lot of stuff.
Conclusion
In conclusion, integrating long-term memory into your LangGraph agents can significantly enhance their utility and user experience, similar to the new ChatGPT functionality. This capability allows your agents to retain critical information and adapt to user preferences without repeated inputs, leading to more seamless and personalized interactions.
Don’t forget to explore the complete Jupyter notebook linked at the end for a deeper dive into the implementation details. Happy coding!
