diff --git a/libs/langgraph/langgraph/graph/message.py b/libs/langgraph/langgraph/graph/message.py index 18c96d436..1b48a843d 100644 --- a/libs/langgraph/langgraph/graph/message.py +++ b/libs/langgraph/langgraph/graph/message.py @@ -63,6 +63,7 @@ def add_messages( right: Messages, *, format: Literal["langchain-openai"] | None = None, + mode: Literal["allow_everything", "append_only"] = "allow_everything", ) -> Messages: """Merges two lists of messages, updating existing messages by ID. @@ -83,10 +84,19 @@ def add_messages( Must have `langchain-core>=0.3.11` installed to use this feature. + mode: Controls how message updates and removals are handled. Options are: + - `allow_everything` (default): Allows adding new messages, updating + existing messages by ID, and removing messages. This is the standard behavior. + - `append_only`: Only allows adding new messages. If a message in `right` + has the same ID as a message in `left` (update or removal), a `ValueError` + will be raised. This mode is useful when you want to prevent any modification + of message history. + Returns: A new list of messages with the messages from `right` merged into `left`. If a message in `right` has the same ID as a message in `left`, the - message from `right` will replace the message from `left`. + message from `right` will replace the message from `left` (in + `allow_everything` mode) or raise a `ValueError` (in `append_only` mode). Example: Basic usage ```python @@ -183,6 +193,30 @@ def add_messages( # } ``` + Example: Use append_only mode to prevent message updates + ```python + from typing import Annotated + from typing_extensions import TypedDict + from langgraph.graph import StateGraph, add_messages + + + class State(TypedDict): + messages: Annotated[list, add_messages(mode="append_only")] + + + # This will work - adding new messages + msgs1 = [HumanMessage(content="Hello", id="1")] + msgs2 = [AIMessage(content="Hi there!", id="2")] + add_messages(msgs1, msgs2, mode="append_only") + # [HumanMessage(content='Hello', id='1'), AIMessage(content='Hi there!', id='2')] + + # This will raise an error - trying to update an existing message + msgs1 = [HumanMessage(content="Hello", id="1")] + msgs2 = [HumanMessage(content="Hello again", id="1")] + add_messages(msgs1, msgs2, mode="append_only") + # ValueError: Cannot update existing message with ID '1' in append_only mode + ``` + """ remove_all_idx = None # coerce to list @@ -207,6 +241,11 @@ def add_messages( if m.id is None: m.id = str(uuid.uuid4()) if isinstance(m, RemoveMessage) and m.id == REMOVE_ALL_MESSAGES: + if mode == "append_only": + raise ValueError( + "Cannot remove all messages in append_only mode. " + "Use mode='allow_everything' to allow message removal." + ) remove_all_idx = idx if remove_all_idx is not None: @@ -218,6 +257,18 @@ def add_messages( ids_to_remove = set() for m in right: if (existing_idx := merged_by_id.get(m.id)) is not None: + # Check for append_only mode violation + if mode == "append_only": + if isinstance(m, RemoveMessage): + raise ValueError( + f"Cannot remove existing message with ID '{m.id}' in append_only mode. " + f"Use mode='allow_everything' to allow message removal." + ) + else: + raise ValueError( + f"Cannot update existing message with ID '{m.id}' in append_only mode. " + f"Use mode='allow_everything' to allow message updates." + ) if isinstance(m, RemoveMessage): ids_to_remove.add(m.id) else: diff --git a/libs/langgraph/tests/test_messages_state.py b/libs/langgraph/tests/test_messages_state.py index f41552b16..7a694213f 100644 --- a/libs/langgraph/tests/test_messages_state.py +++ b/libs/langgraph/tests/test_messages_state.py @@ -367,3 +367,139 @@ def test_push_messages_in_graph(): messages.append(message) assert values["messages"] == messages + + +def test_append_only_mode_allows_new_messages(): + """Test that append_only mode allows adding new messages.""" + left = [HumanMessage(content="Hello", id="1")] + right = [AIMessage(content="Hi there!", id="2")] + result = add_messages(left, right, mode="append_only") + expected = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + assert result == expected + + +def test_append_only_mode_prevents_updates(): + """Test that append_only mode raises an error when trying to update an existing message.""" + left = [HumanMessage(content="Hello", id="1")] + right = [HumanMessage(content="Hello again", id="1")] + with pytest.raises( + ValueError, + match="Cannot update existing message with ID '1' in append_only mode", + ): + add_messages(left, right, mode="append_only") + + +def test_append_only_mode_prevents_remove_messages(): + """Test that append_only mode prevents RemoveMessage operations.""" + left = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + right = [RemoveMessage(id="2")] + with pytest.raises( + ValueError, + match="Cannot remove existing message with ID '2' in append_only mode", + ): + add_messages(left, right, mode="append_only") + + +def test_append_only_mode_prevents_remove_all_messages(): + """Test that append_only mode prevents REMOVE_ALL_MESSAGES operations.""" + left = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + right = [RemoveMessage(id=REMOVE_ALL_MESSAGES)] + with pytest.raises( + ValueError, + match="Cannot remove all messages in append_only mode", + ): + add_messages(left, right, mode="append_only") + + +def test_append_only_mode_with_multiple_new_messages(): + """Test that append_only mode allows adding multiple new messages.""" + left = [HumanMessage(content="Hello", id="1")] + right = [ + AIMessage(content="Hi there!", id="2"), + SystemMessage(content="System message", id="3"), + HumanMessage(content="Another message", id="4"), + ] + result = add_messages(left, right, mode="append_only") + expected = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + SystemMessage(content="System message", id="3"), + HumanMessage(content="Another message", id="4"), + ] + assert result == expected + + +def test_append_only_mode_with_mixed_operations(): + """Test that append_only mode fails when mixing new messages with updates.""" + left = [HumanMessage(content="Hello", id="1")] + right = [ + AIMessage(content="Hi there!", id="2"), # new message + HumanMessage(content="Updated hello", id="1"), # update attempt + ] + with pytest.raises( + ValueError, + match="Cannot update existing message with ID '1' in append_only mode", + ): + add_messages(left, right, mode="append_only") + + +def test_allow_everything_mode_default_behavior(): + """Test that allow_everything mode is the default and allows updates.""" + left = [HumanMessage(content="Hello", id="1")] + right = [HumanMessage(content="Hello again", id="1")] + # Test without specifying mode (should default to allow_everything) + result = add_messages(left, right) + expected = [HumanMessage(content="Hello again", id="1")] + assert result == expected + + # Test with explicit mode + result = add_messages(left, right, mode="allow_everything") + assert result == expected + + +def test_append_only_mode_in_state_graph(): + """Test append_only mode works correctly in a StateGraph.""" + + class State(TypedDict): + messages: Annotated[list[AnyMessage], add_messages(mode="append_only")] + + def add_message(state: State): + return {"messages": [AIMessage(content="Response", id="2")]} + + def try_update_message(state: State): + # This should fail because message with id="1" already exists + return {"messages": [HumanMessage(content="Updated", id="1")]} + + # Test successful case + builder = StateGraph(State) + builder.add_node("add_message", add_message) + builder.add_edge(START, "add_message") + builder.add_edge("add_message", END) + graph = builder.compile() + + result = graph.invoke({"messages": [HumanMessage(content="Hello", id="1")]}) + assert len(result["messages"]) == 2 + assert result["messages"][0].content == "Hello" + assert result["messages"][1].content == "Response" + + # Test failure case + builder2 = StateGraph(State) + builder2.add_node("try_update", try_update_message) + builder2.add_edge(START, "try_update") + builder2.add_edge("try_update", END) + graph2 = builder2.compile() + + with pytest.raises( + ValueError, + match="Cannot update existing message with ID '1' in append_only mode", + ): + graph2.invoke({"messages": [HumanMessage(content="Hello", id="1")]})