[Langgraph] Remove a message from the graph state

[HGInfoNancy]

Checked other resources

• I added a very descriptive title to this question.
• I searched the LangChain documentation with the integrated search.
• I used the GitHub search to find a similar question and didn't find it.

Commit to Help

• I commit to help with one of those options 👆

Example Code

import operator
from typing import Annotated, Sequence, TypedDict

class GraphState(TypedDict):
    # The add_messages function defines how an update should be processed
    # Default is to replace. add_messages says "append"
    messages: Annotated[Sequence[BaseMessage], operator.add]

workflow = StateGraph(GraphState)

workflow.add_node("node1", lambda x:{"messages":"A first LLM message"})
workflow.add_node("node2", In this node I want to forget the message from node1)
workflow.add_edge("node1", "node2")

Description

Hi, I want to forget/delete some messages previously stored in my graph state.
In my case I have "messages" and I want in another node delete this message.
I tried something like :
def function_node2(state):
state = state["messages"][:-1]

The state is now updated within my function, but as soon as I went into another node (let's say node3) it seems load the state from another way and it rolls back to the previous state

Can someone help me? Thanks

System Info

[dosubot[bot]]

Hello @HGInfoNancy! 👋 I'm here to help you with any bugs, questions, or contributions you have. Let's tackle this issue together!

To ensure that the state update persists across nodes in a LangChain graph state when you want to delete a message from the state, you can use the clear method provided in the Neo4jChatMessageHistory class. This method deletes all messages associated with a specific session from the Neo4j database, ensuring that the state is updated and persisted across nodes.

Here is the relevant code for the clear method:

def clear(self) -> None:
    """Clear session memory from Neo4j"""
    query = (
        f"MATCH (s:`{self._node_label}`)-[:LAST_MESSAGE]->(last_message) "
        "WHERE s.id = $session_id MATCH p=(last_message)<-[:NEXT]-() "
        "WITH p, length(p) AS length ORDER BY length DESC LIMIT 1 "
        "UNWIND nodes(p) as node DETACH DELETE node;"
    )
    self._driver.execute_query(query, {"session_id": self._session_id}).summary

This method constructs a Cypher query to match the session node and its associated messages, then deletes them. This ensures that the state update (i.e., deletion of messages) is persisted across nodes ^[1].

If you are not using Neo4j, you might need to implement a similar mechanism in your storage backend to ensure that the state changes are persisted correctly across nodes.

To continue talking to Dosu, mention @dosu.

[denizkenan]

So nobody cares?

[nikotang]

According to this doc page, you can use RemoveMessage exactly like how you would add messages of whatever type when you have the add_message reducer to your messages key. Basically, change the message type of the message you want to remove to the RemoveMessage type.

Say your class is

class GraphState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_message]

(@HGInfoNancy idk if operator.add works exactly the same)

While you can add messages like this:

def my_node(state: State):
    new_message = HumanMessage(content="message")
    return {"messages": new_message}

you can remove messages the same way, identifying the message you want to change to the RemoveMessage type by its ID:

def my_node(state: State):
    messages_to_delete = [RemoveMessage(id=m.id) for m in state['messages'][:-1]] # here only leaving the most recent message in memory
    return {"messages": messages_to_delete}

[marianocodes]

That partially works. The message gets deleted but when the graph ends, the other nodes receive all the messages

[m-swamina]

Did you find a solution for this ? Such that updated message array persists?

[PragyanTiwari]

A minor change in your code can solve this problem.

While defining the function function_node2, you're not returning the updated state after modifying it. As a result, the state remains unchanged, and you receive the same output as before.

Here is the correct approach:

def function_node2(state):
updated_messages = state["messages"][:-1]
return {"messages":updated_messages}

This would remove the last message, but doesn't reset the state completely.

I hope this can help...