langgraph/examples/memory/delete-messages.ipynb

How to delete messages

One of the common states for a graph is a list of messages. Usually you only add messages to that state. However, sometimes you may want to remove messages (either by directly modifying the state or as part of the graph). To do that, you can use the RemoveMessage modifier. In this guide, we will cover how to do that.

The key idea is that each state key has a reducer key. This key specifies how to combine updates to the state. The default MessagesState has a messages key, and the reducer for that key accepts these RemoveMessage modifiers. That reducer then uses these RemoveMessage to delete messages from the key.

So note that just because your graph state has a key that is a list of messages, it doesn't mean that that this RemoveMessage modifier will work. You also have to have a reducer defined that knows how to work with this.

NOTE: Many models expect certain rules around lists of messages. For example, some expect them to start with a user message, others expect all messages with tool calls to be followed by a tool message. When deleting messages, you will want to make sure you don't violate these rules.

Setup

First, let's build a simple graph that uses messages. Note that it's using the MessagesState which has the required reducer.

In [ ]:

%%capture --no-stderr
%pip install --quiet -U langgraph langchain_anthropic

Next, we need to set API keys for Anthropic (the LLM we will use)

In [2]:

import getpass
import os


def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")


_set_env("ANTHROPIC_API_KEY")

Optionally, we can set API key for LangSmith tracing, which will give us best-in-class observability.

In [3]:

os.environ["LANGCHAIN_TRACING_V2"] = "true"
_set_env("LANGCHAIN_API_KEY")

Build the agent

Let's now build a simple ReAct style agent.

In [2]:

from typing import Literal

from langchain_anthropic import ChatAnthropic
from langchain_core.tools import tool

from langgraph.checkpoint.sqlite import SqliteSaver
from langgraph.graph import MessagesState, StateGraph, START
from langgraph.prebuilt import ToolNode

memory = SqliteSaver.from_conn_string(":memory:")


@tool
def search(query: str):
    """Call to surf the web."""
    # This is a placeholder for the actual implementation
    # Don't let the LLM know this though 😊
    return [
        "It's sunny in San Francisco, but you better look out if you're a Gemini 😈."
    ]


tools = [search]
tool_node = ToolNode(tools)
model = ChatAnthropic(model_name="claude-3-haiku-20240307")
bound_model = model.bind_tools(tools)


def should_continue(state: MessagesState) -> Literal["action", "__end__"]:
    """Return the next node to execute."""
    last_message = state["messages"][-1]
    # If there is no function call, then we finish
    if not last_message.tool_calls:
        return "__end__"
    # Otherwise if there is, we continue
    return "action"


# Define the function that calls the model
def call_model(state: MessagesState):
    response = model.invoke(state["messages"])
    # We return a list, because this will get added to the existing list
    return {"messages": response}


# Define a new graph
workflow = StateGraph(MessagesState)

# Define the two nodes we will cycle between
workflow.add_node("agent", call_model)
workflow.add_node("action", tool_node)

# Set the entrypoint as `agent`
# This means that this node is the first one called
workflow.add_edge(START, "agent")

# We now add a conditional edge
workflow.add_conditional_edges(
    # First, we define the start node. We use `agent`.
    # This means these are the edges taken after the `agent` node is called.
    "agent",
    # Next, we pass in the function that will determine which node is called next.
    should_continue,
)

# We now add a normal edge from `tools` to `agent`.
# This means that after `tools` is called, `agent` node is called next.
workflow.add_edge("action", "agent")

# Finally, we compile it!
# This compiles it into a LangChain Runnable,
# meaning you can use it as you would any other runnable
app = workflow.compile(checkpointer=memory)

In [3]:

from langchain_core.messages import HumanMessage

config = {"configurable": {"thread_id": "2"}}
input_message = HumanMessage(content="hi! I'm bob")
for event in app.stream({"messages": [input_message]}, config, stream_mode="values"):
    event["messages"][-1].pretty_print()


input_message = HumanMessage(content="what's my name?")
for event in app.stream({"messages": [input_message]}, config, stream_mode="values"):
    event["messages"][-1].pretty_print()

================================ Human Message =================================

hi! I'm bob
================================== Ai Message ==================================

Hello Bob! It's nice to meet you. How can I assist you today?
================================ Human Message =================================

what's my name?
================================== Ai Message ==================================

Your name is Bob, as you introduced yourself at the beginning of our conversation.

Manually deleting messages

First, we will cover how to manually delete messages. Let's take a look at the current state of the thread:

In [7]:

messages = app.get_state(config).values['messages']
messages

Out [7]:

[HumanMessage(content="hi! I'm bob", id='3e1098f8-2657-42d3-b58a-7c2f46930b8c'),
 AIMessage(content="Hello Bob! It's nice to meet you. How can I assist you today?", response_metadata={'id': 'msg_01HT8MUEN4p16wbYv9Xm7kfr', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 20}}, id='run-86348912-72c4-42b0-b3e0-a47c4ebd1e52-0'),
 HumanMessage(content="what's my name?", id='9c3ef235-ec5c-4e57-a3b2-c17502de496d'),
 AIMessage(content='Your name is Bob, as you introduced yourself at the beginning of our conversation.', response_metadata={'id': 'msg_01LVhb56f6RpAAoxASZrLzmK', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 40, 'output_tokens': 19}}, id='run-e3d7447f-046a-4dfa-8813-38134dbcd1ef-0')]

We can call update_state and pass in the id of the first message. This will delete that message.

In [9]:

from langchain_core.messages import RemoveMessage
app.update_state(config, {"messages": RemoveMessage(id=messages[0].id)})

Out [9]:

{'configurable': {'thread_id': '2',
  'thread_ts': '1ef3d750-5bc4-67c6-8005-9490a1b276f5'}}

If we now look at the messages, we can verify that the first one was deleted.

In [11]:

messages = app.get_state(config).values['messages']
messages

Out [11]:

[AIMessage(content="Hello Bob! It's nice to meet you. How can I assist you today?", response_metadata={'id': 'msg_01HT8MUEN4p16wbYv9Xm7kfr', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 20}}, id='run-86348912-72c4-42b0-b3e0-a47c4ebd1e52-0'),
 HumanMessage(content="what's my name?", id='9c3ef235-ec5c-4e57-a3b2-c17502de496d'),
 AIMessage(content='Your name is Bob, as you introduced yourself at the beginning of our conversation.', response_metadata={'id': 'msg_01LVhb56f6RpAAoxASZrLzmK', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 40, 'output_tokens': 19}}, id='run-e3d7447f-046a-4dfa-8813-38134dbcd1ef-0')]

Programmatically deleting messages

We can also delete messages programmatically from inside the graph. Here we'll modify the graph to delete any old messages (longer than 3 messages ago) at the end of a graph run.

In [14]:

from langchain_core.messages import RemoveMessage
from langgraph.graph import END


def delete_messages(state):
    messages = state['messages']
    if len(messages) > 3:
        return {"messages": [RemoveMessage(id=m.id) for m in messages[:-3]]}

# We need to modify the logic to call delete_messages rather than end right away
def should_continue(state: MessagesState) -> Literal["action", "delete_messages"]:
    """Return the next node to execute."""
    last_message = state["messages"][-1]
    # If there is no function call, then we call our delete_messages function
    if not last_message.tool_calls:
        return "delete_messages"
    # Otherwise if there is, we continue
    return "action"


# Define a new graph
workflow = StateGraph(MessagesState)
workflow.add_node("agent", call_model)
workflow.add_node("action", tool_node)

# This is our new node we're defining
workflow.add_node(delete_messages)


workflow.add_edge(START, "agent")
workflow.add_conditional_edges("agent", should_continue,)
workflow.add_edge("action", "agent")

# This is the new edge we're adding: after we delete messages, we finish
workflow.add_edge("delete_messages", END)
app = workflow.compile(checkpointer=memory)

We can now try this out. We can call the graph twice and then check the state

In [16]:

from langchain_core.messages import HumanMessage

config = {"configurable": {"thread_id": "3"}}
input_message = HumanMessage(content="hi! I'm bob")
for event in app.stream({"messages": [input_message]}, config, stream_mode="values"):
    event["messages"][-1].pretty_print()


input_message = HumanMessage(content="what's my name?")
for event in app.stream({"messages": [input_message]}, config, stream_mode="values"):
    event["messages"][-1].pretty_print()

================================ Human Message =================================

hi! I'm bob
================================== Ai Message ==================================

It's nice to meet you, Bob! How can I assist you today?
================================ Human Message =================================

what's my name?
================================== Ai Message ==================================

You said your name is Bob.
================================== Ai Message ==================================

You said your name is Bob.

If we now check the state, we should see that it is only three messages long. This is because we just deleted the earlier messages - otherwise it would be four!

In [17]:

messages = app.get_state(config).values['messages']
messages

Out [17]:

[AIMessage(content="It's nice to meet you, Bob! How can I assist you today?", response_metadata={'id': 'msg_01QMoxepDiCcKQ6XFgge1QQT', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 12, 'output_tokens': 19}}, id='run-de13ba05-095d-4fd1-907a-6766ef3bf57b-0'),
 HumanMessage(content="what's my name?", id='8292e725-8fc4-487e-a9b6-75f8b136bec2'),
 AIMessage(content='You said your name is Bob.', response_metadata={'id': 'msg_01DfWfaxavdMCqtoQRmC3mc4', 'model': 'claude-3-haiku-20240307', 'stop_reason': 'end_turn', 'stop_sequence': None, 'usage': {'input_tokens': 39, 'output_tokens': 10}}, id='run-28167c82-e126-47e4-854c-623e50c8af22-0')]

Remember, when deleting messages you will want to make sure that the remaining message list is still valid. This message list may actually not be - this is because it currently starts with an AI message, which some models do not allow.

In [ ]: