langgraph: add support for deleting messages (#651)

* (RFC) add support for deleting messages

* alternative design

* lint

* Revert "lint"

This reverts commit 88b59d2aef.

* Revert "alternative design"

This reverts commit 2ad6009160.

* switch to lc core removemessage

* lint

* add unittests

libs/langgraph/langgraph/graph/message.py

@@ -4,6 +4,7 @@ from typing import Annotated, TypedDict, Union
 from langchain_core.messages import (
     AnyMessage,
     MessageLikeRepresentation,
+    RemoveMessage,
     convert_to_messages,
     message_chunk_to_message,
 )
@@ -79,8 +80,16 @@ def add_messages(left: Messages, right: Messages) -> Messages:
     merged = left.copy()
     for m in right:
         if (existing_idx := left_idx_by_id.get(m.id)) is not None:
-            merged[existing_idx] = m
+            if isinstance(m, RemoveMessage):
+                del merged[existing_idx]
+            else:
+                merged[existing_idx] = m
         else:
+            if isinstance(m, RemoveMessage):
+                raise ValueError(
+                    f"Attempting to delete a message with an ID that doesn't exist ('{m.id}')"
+                )
+
             merged.append(m)
     return merged
 

libs/langgraph/tests/test_pregel.py

@@ -7494,3 +7494,57 @@ def test_checkpoint_metadata() -> None:
         assert chkpnt_tuple.metadata["thread_id"] == "2"
         assert chkpnt_tuple.metadata["test_config_3"] == "foo"
         assert chkpnt_tuple.metadata["test_config_4"] == "bar"
+
+
+def test_remove_message_via_state_update():
+    from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
+
+    workflow = MessageGraph()
+    workflow.add_node(
+        "chatbot",
+        lambda state: [
+            AIMessage(
+                content="Hello! How can I help you",
+            )
+        ],
+    )
+
+    workflow.set_entry_point("chatbot")
+    workflow.add_edge("chatbot", END)
+
+    checkpointer = MemorySaver()
+    app = workflow.compile(checkpointer=checkpointer)
+    config = {"configurable": {"thread_id": "1"}}
+    output = app.invoke([HumanMessage(content="Hi")], config=config)
+    app.update_state(config, values=[RemoveMessage(id=output[-1].id)])
+
+    updated_state = app.get_state(config)
+
+    assert len(updated_state.values) == 1
+    assert updated_state.values[-1].content == "Hi"
+
+
+def test_remove_message_from_node():
+    from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage
+
+    workflow = MessageGraph()
+    workflow.add_node(
+        "chatbot",
+        lambda state: [
+            AIMessage(
+                content="Hello!",
+            ),
+            AIMessage(
+                content="How can I help you?",
+            ),
+        ],
+    )
+    workflow.add_node("delete_messages", lambda state: [RemoveMessage(id=state[-2].id)])
+    workflow.set_entry_point("chatbot")
+    workflow.add_edge("chatbot", "delete_messages")
+    workflow.add_edge("delete_messages", END)
+
+    app = workflow.compile()
+    output = app.invoke([HumanMessage(content="Hi")])
+    assert len(output) == 2
+    assert output[-1].content == "How can I help you?"