diff --git a/libs/langgraph/langgraph/graph/message.py b/libs/langgraph/langgraph/graph/message.py index b6ed16131..2892051c5 100644 --- a/libs/langgraph/langgraph/graph/message.py +++ b/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)} diff --git a/libs/langgraph/tests/test_messages_state.py b/libs/langgraph/tests/test_messages_state.py index 97f38051c..d88e6226b 100644 --- a/libs/langgraph/tests/test_messages_state.py +++ b/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"), + ]