langgraph: support removing all messages with RemoveMessage (#4117)

libs/langgraph/langgraph/graph/message.py

@@ -27,6 +27,8 @@ from langgraph.graph.state import StateGraph
 
 Messages = Union[list[MessageLikeRepresentation], MessageLikeRepresentation]
 
+REMOVE_ALL_MESSAGES = "__remove_all__"
+
 
 def _add_messages_wrapper(func: Callable) -> Callable[[Messages, Messages], Messages]:
     def _add_messages(
@@ -158,6 +160,7 @@ def add_messages(
 
         Support for 'format="langchain-openai"' flag added.
     """
+    remove_all_idx = None
     # coerce to list
     if not isinstance(left, list):
         left = [left]  # type: ignore[assignment]
@@ -176,9 +179,15 @@ def add_messages(
     for m in left:
         if m.id is None:
             m.id = str(uuid.uuid4())
-    for m in right:
+    for idx, m in enumerate(right):
         if m.id is None:
             m.id = str(uuid.uuid4())
+        if isinstance(m, RemoveMessage) and m.id == REMOVE_ALL_MESSAGES:
+            remove_all_idx = idx
+
+    if remove_all_idx is not None:
+        return right[remove_all_idx + 1 :]
+
     # merge
     merged = left.copy()
     merged_by_id = {m.id: i for i, m in enumerate(merged)}

libs/langgraph/tests/test_messages_state.py

@@ -16,7 +16,7 @@ from pydantic.v1 import BaseModel as BaseModelV1
 from typing_extensions import TypedDict
 
 from langgraph.graph import add_messages
-from langgraph.graph.message import MessagesState
+from langgraph.graph.message import REMOVE_ALL_MESSAGES, MessagesState
 from langgraph.graph.state import END, START, StateGraph
 from tests.conftest import IS_LANGCHAIN_CORE_030_OR_GREATER
 from tests.messages import _AnyIdHumanMessage
@@ -313,3 +313,32 @@ def test_messages_state_format_openai():
     for m in result["messages"]:
         m.id = None
     assert result == {"messages": expected}
+
+
+def test_remove_all_messages():
+    # simple removal
+    left = [HumanMessage(content="Hello"), AIMessage(content="Hi there!")]
+    right = [RemoveMessage(id=REMOVE_ALL_MESSAGES)]
+    result = add_messages(left, right)
+    assert result == []
+
+    # removal and update (i.e., overwriting)
+    left = [HumanMessage(content="Hello"), AIMessage(content="Hi there!")]
+    right = [
+        RemoveMessage(id=REMOVE_ALL_MESSAGES),
+        HumanMessage(content="Updated hello"),
+    ]
+    result = add_messages(left, right)
+    assert result == [_AnyIdHumanMessage(content="Updated hello")]
+
+    # test removing preceding messages in the right list
+    left = [HumanMessage(content="Hello"), AIMessage(content="Hi there!")]
+    right = [
+        HumanMessage(content="Updated hello"),
+        RemoveMessage(id=REMOVE_ALL_MESSAGES),
+        HumanMessage(content="Updated hi there"),
+    ]
+    result = add_messages(left, right)
+    assert result == [
+        _AnyIdHumanMessage(content="Updated hi there"),
+    ]