diff --git a/libs/langgraph/langgraph/graph/__init__.py b/libs/langgraph/langgraph/graph/__init__.py index 7bea3fc82..1104132d6 100644 --- a/libs/langgraph/langgraph/graph/__init__.py +++ b/libs/langgraph/langgraph/graph/__init__.py @@ -1,5 +1,10 @@ from langgraph.constants import END, START -from langgraph.graph.message import MessageGraph, MessagesState, add_messages +from langgraph.graph.message import ( + MessageGraph, + MessagesState, + add_messages, + validate_messages_append_only, +) from langgraph.graph.state import StateGraph __all__ = ( @@ -9,4 +14,5 @@ __all__ = ( "add_messages", "MessagesState", "MessageGraph", + "validate_messages_append_only", ) diff --git a/libs/langgraph/langgraph/graph/message.py b/libs/langgraph/langgraph/graph/message.py index 1b48a843d..14b7037a7 100644 --- a/libs/langgraph/langgraph/graph/message.py +++ b/libs/langgraph/langgraph/graph/message.py @@ -31,6 +31,7 @@ __all__ = ( "MessagesState", "MessageGraph", "REMOVE_ALL_MESSAGES", + "validate_messages_append_only", ) Messages = list[MessageLikeRepresentation] | MessageLikeRepresentation @@ -63,7 +64,6 @@ 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. @@ -193,30 +193,6 @@ 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 @@ -241,11 +217,6 @@ 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: @@ -257,18 +228,6 @@ 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: @@ -295,6 +254,92 @@ def add_messages( return merged +def validate_messages_append_only( + input: dict[str, Any], current_state: dict[str, Any] +) -> None: + """Validates that incoming messages are append-only (no updates or removals). + + This validator can be passed to `compile(validate_input=...)` to enforce that + external inputs only append new messages and never update or remove existing ones. + + Internal node updates bypass this validation, allowing nodes to perform operations + like message compaction or cleanup without restriction. + + Args: + input: The raw input dictionary being provided externally + current_state: The current state loaded from the checkpoint + + Raises: + ValueError: If the input attempts to update or remove existing messages + + Example: + ```python + from langgraph.graph import StateGraph, MessagesState + from langgraph.graph.message import validate_messages_append_only + + builder = StateGraph(MessagesState) + builder.add_node("chatbot", my_chatbot_node) + builder.set_entry_point("chatbot") + builder.set_finish_point("chatbot") + + # Compile with validator to enforce append-only at the boundary + graph = builder.compile(validate_input=validate_messages_append_only) + + # This works - adding new messages + graph.invoke({"messages": [("user", "Hello")]}) + + # This raises ValueError - trying to update existing message + graph.invoke({"messages": [HumanMessage(content="Modified", id="existing-id")]}) + ``` + """ + # Only validate if input contains messages + if "messages" not in input: + return + + # Get current messages from state + current_messages = current_state.get("messages", []) + + # Build a set of existing message IDs + existing_ids: set[str] = set() + for msg in current_messages: + if hasattr(msg, "id") and msg.id is not None: + existing_ids.add(msg.id) + + # Coerce input messages to list + input_messages = input["messages"] + if not isinstance(input_messages, list): + input_messages = [input_messages] + + # Convert to messages to get proper IDs + try: + converted_messages = convert_to_messages(input_messages) + except Exception: + # If conversion fails, let it through - the reducer will handle it + return + + # Check each input message + for msg in converted_messages: + # Check for REMOVE_ALL_MESSAGES + if isinstance(msg, RemoveMessage) and msg.id == REMOVE_ALL_MESSAGES: + raise ValueError( + "Cannot remove all messages in append_only mode. " + "External inputs must only append new messages." + ) + + # Check for updates or removals of existing messages + if msg.id and msg.id in existing_ids: + if isinstance(msg, RemoveMessage): + raise ValueError( + f"Cannot remove existing message with ID '{msg.id}' in append_only mode. " + "External inputs must only append new messages." + ) + else: + raise ValueError( + f"Cannot update existing message with ID '{msg.id}' in append_only mode. " + "External inputs must only append new messages." + ) + + @deprecated( "MessageGraph is deprecated in langgraph 1.0.0, to be removed in 2.0.0. Please use StateGraph with a `messages` key instead.", category=None, diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 4d0c90457..c7da83334 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -1042,6 +1042,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): interrupt_after: All | list[str] | None = None, debug: bool = False, name: str | None = None, + validate_input: Callable[[Any, dict[str, Any]], None] | None = None, ) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]: """Compiles the `StateGraph` into a `CompiledStateGraph` object. @@ -1074,6 +1075,34 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): interrupt_after: An optional list of node names to interrupt after. debug: A flag indicating whether to enable debug mode. name: The name to use for the compiled graph. + validate_input: Optional validation function for external inputs. + The function receives `(input, current_state)` and should raise an + exception if the input is invalid, or return normally if valid. + + This validation only applies to external updates (from `invoke()`, `stream()`, + `update_state()`, or `Command.update`) - internal node updates bypass this + validation, allowing trusted nodes to perform privileged operations. + + Example: + ```python + from langgraph.graph import StateGraph + from langgraph.graph.message import validate_messages_append_only + + builder = StateGraph(MessagesState) + # ... add nodes ... + + # Use built-in validator + graph = builder.compile( + validate_input=validate_messages_append_only + ) + + # Or custom validator + def my_validator(input: dict, current_state: dict) -> None: + if "forbidden_key" in input: + raise ValueError("Forbidden key detected") + + graph = builder.compile(validate_input=my_validator) + ``` Returns: CompiledStateGraph: The compiled `StateGraph`. @@ -1134,6 +1163,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): store=store, cache=cache, name=name or "LangGraph", + validate_input=validate_input, ) compiled.attach_node(START, None) diff --git a/libs/langgraph/langgraph/pregel/_loop.py b/libs/langgraph/langgraph/pregel/_loop.py index 79201d0d6..7bd69f1d0 100644 --- a/libs/langgraph/langgraph/pregel/_loop.py +++ b/libs/langgraph/langgraph/pregel/_loop.py @@ -226,6 +226,7 @@ class PregelLoop: migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + validate_input: Callable[[Any, dict[str, Any]], None] | None = None, ) -> None: self.stream = stream self.config = config @@ -250,6 +251,7 @@ class PregelLoop: self.retry_policy = retry_policy self.cache_policy = cache_policy self.durability = durability + self.validate_input = validate_input if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) scratchpad: PregelScratchpad | None = config[CONF].get(CONFIG_KEY_SCRATCHPAD) @@ -637,6 +639,14 @@ class PregelLoop: # map command to writes if isinstance(self.input, Command): + # Validate command update if validator is provided + if self.validate_input is not None and self.input.update: + current_state = read_channels(self.channels, self.output_keys) + try: + self.validate_input(self.input.update, current_state) + except Exception as e: + raise ValueError(f"Input validation failed: {str(e)}") from e + if (resume := self.input.resume) is not None: if not self.checkpointer: raise RuntimeError( @@ -691,6 +701,13 @@ class PregelLoop: ) # map inputs to channel updates elif input_writes := deque(map_input(input_keys, self.input)): + # Validate input if validator is provided + if self.validate_input is not None: + current_state = read_channels(self.channels, self.output_keys) + try: + self.validate_input(self.input, current_state) + except Exception as e: + raise ValueError(f"Input validation failed: {str(e)}") from e # discard any unfinished tasks from previous checkpoint discard_tasks = prepare_next_tasks( self.checkpoint, @@ -985,6 +1002,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + validate_input: Callable[[Any, dict[str, Any]], None] | None = None, ) -> None: super().__init__( input, @@ -1006,6 +1024,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager): retry_policy=retry_policy, cache_policy=cache_policy, durability=durability, + validate_input=validate_input, ) self.stack = ExitStack() if checkpointer: @@ -1161,6 +1180,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): migrate_checkpoint: Callable[[Checkpoint], None] | None = None, retry_policy: Sequence[RetryPolicy] = (), cache_policy: CachePolicy | None = None, + validate_input: Callable[[Any, dict[str, Any]], None] | None = None, ) -> None: super().__init__( input, @@ -1182,6 +1202,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager): retry_policy=retry_policy, cache_policy=cache_policy, durability=durability, + validate_input=validate_input, ) self.stack = AsyncExitStack() if checkpointer: diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 37e8125f9..f5b5fc833 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -652,6 +652,7 @@ class Pregel( config: RunnableConfig | None = None, trigger_to_nodes: Mapping[str, Sequence[str]] | None = None, name: str = "LangGraph", + validate_input: Callable[[Any, dict[str, Any]], None] | None = None, **deprecated_kwargs: Unpack[DeprecatedKwargs], ) -> None: if ( @@ -698,6 +699,7 @@ class Pregel( self.config = config self.trigger_to_nodes = trigger_to_nodes or {} self.name = name + self.validate_input = validate_input if auto_validate: self.validate() @@ -2599,6 +2601,7 @@ class Pregel( migrate_checkpoint=self._migrate_checkpoint, retry_policy=self.retry_policy, cache_policy=self.cache_policy, + validate_input=self.validate_input, ) as loop: # create runner runner = PregelRunner( @@ -2908,6 +2911,7 @@ class Pregel( migrate_checkpoint=self._migrate_checkpoint, retry_policy=self.retry_policy, cache_policy=self.cache_policy, + validate_input=self.validate_input, ) as loop: # create runner runner = PregelRunner( diff --git a/libs/langgraph/tests/test_messages_state.py b/libs/langgraph/tests/test_messages_state.py index 7a694213f..971cc0ebd 100644 --- a/libs/langgraph/tests/test_messages_state.py +++ b/libs/langgraph/tests/test_messages_state.py @@ -20,7 +20,7 @@ from langgraph.graph.message import REMOVE_ALL_MESSAGES, MessagesState, push_mes from langgraph.graph.state import StateGraph from tests.messages import _AnyIdHumanMessage -_, CORE_MINOR, CORE_PATCH = ( +CORE_MAJOR, CORE_MINOR, CORE_PATCH = ( int("".join(c for c in v if c.isdigit())) for v in langchain_core.__version__.split(".") ) @@ -205,7 +205,11 @@ def test_messages_state(state_schema): @pytest.mark.skipif( - condition=not ((CORE_MINOR == 3 and CORE_PATCH >= 11) or CORE_MINOR > 3), + condition=not ( + (CORE_MAJOR == 0 and CORE_MINOR == 3 and CORE_PATCH >= 11) + or (CORE_MAJOR == 0 and CORE_MINOR > 3) + or CORE_MAJOR > 0 + ), reason="Requires langchain_core>=0.3.11.", ) def test_messages_state_format_openai(): @@ -367,139 +371,3 @@ 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")]}) diff --git a/libs/langgraph/tests/test_messages_state.py.bak b/libs/langgraph/tests/test_messages_state.py.bak new file mode 100644 index 000000000..7a694213f --- /dev/null +++ b/libs/langgraph/tests/test_messages_state.py.bak @@ -0,0 +1,505 @@ +from typing import Annotated +from uuid import UUID + +import langchain_core +import pytest +from langchain_core.messages import ( + AIMessage, + AnyMessage, + HumanMessage, + RemoveMessage, + SystemMessage, + ToolMessage, +) +from pydantic import BaseModel +from typing_extensions import TypedDict + +from langgraph.constants import END, START +from langgraph.graph import add_messages +from langgraph.graph.message import REMOVE_ALL_MESSAGES, MessagesState, push_message +from langgraph.graph.state import StateGraph +from tests.messages import _AnyIdHumanMessage + +_, CORE_MINOR, CORE_PATCH = ( + int("".join(c for c in v if c.isdigit())) + for v in langchain_core.__version__.split(".") +) + + +def test_add_single_message(): + left = [HumanMessage(content="Hello", id="1")] + right = AIMessage(content="Hi there!", id="2") + result = add_messages(left, right) + expected_result = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + assert result == expected_result + + +def test_add_multiple_messages(): + left = [HumanMessage(content="Hello", id="1")] + right = [ + AIMessage(content="Hi there!", id="2"), + SystemMessage(content="System message", id="3"), + ] + result = add_messages(left, right) + expected_result = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + SystemMessage(content="System message", id="3"), + ] + assert result == expected_result + + +def test_update_existing_message(): + left = [HumanMessage(content="Hello", id="1")] + right = HumanMessage(content="Hello again", id="1") + result = add_messages(left, right) + expected_result = [HumanMessage(content="Hello again", id="1")] + assert result == expected_result + + +def test_missing_ids(): + left = [HumanMessage(content="Hello")] + right = [AIMessage(content="Hi there!")] + result = add_messages(left, right) + assert len(result) == 2 + assert all(isinstance(m.id, str) and UUID(m.id, version=4) for m in result) + + +def test_duplicates_in_input(): + left = [] + right = [ + AIMessage(id="1", content="Hi there!"), + AIMessage(id="1", content="Hi there again!"), + ] + result = add_messages(left, right) + assert len(result) == 1 + assert result[0].id == "1" + assert result[0].content == "Hi there again!" + + +def test_duplicates_in_input_with_remove(): + left = [AIMessage(id="1", content="Hello!")] + right = [ + RemoveMessage(id="1"), + AIMessage(id="1", content="Hi there!"), + AIMessage(id="1", content="Hi there again!"), + ] + result = add_messages(left, right) + assert len(result) == 1 + assert result[0].id == "1" + assert result[0].content == "Hi there again!" + + +def test_remove_message(): + left = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + right = RemoveMessage(id="2") + result = add_messages(left, right) + expected_result = [HumanMessage(content="Hello", id="1")] + assert result == expected_result + + +def test_duplicate_remove_message(): + left = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + right = [RemoveMessage(id="2"), RemoveMessage(id="2")] + result = add_messages(left, right) + expected_result = [HumanMessage(content="Hello", id="1")] + assert result == expected_result + + +def test_remove_nonexistent_message(): + left = [HumanMessage(content="Hello", id="1")] + right = RemoveMessage(id="2") + with pytest.raises( + ValueError, match="Attempting to delete a message with an ID that doesn't exist" + ): + add_messages(left, right) + + +def test_mixed_operations(): + left = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + right = [ + HumanMessage(content="Updated hello", id="1"), + RemoveMessage(id="2"), + SystemMessage(content="New message", id="3"), + ] + result = add_messages(left, right) + expected_result = [ + HumanMessage(content="Updated hello", id="1"), + SystemMessage(content="New message", id="3"), + ] + assert result == expected_result + + +def test_empty_inputs(): + assert add_messages([], []) == [] + assert add_messages([], [HumanMessage(content="Hello", id="1")]) == [ + HumanMessage(content="Hello", id="1") + ] + assert add_messages([HumanMessage(content="Hello", id="1")], []) == [ + HumanMessage(content="Hello", id="1") + ] + + +def test_non_list_inputs(): + left = HumanMessage(content="Hello", id="1") + right = AIMessage(content="Hi there!", id="2") + result = add_messages(left, right) + expected_result = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + assert result == expected_result + + +def test_delete_all(): + left = [ + HumanMessage(content="Hello", id="1"), + AIMessage(content="Hi there!", id="2"), + ] + right = [ + RemoveMessage(id="1"), + RemoveMessage(id="2"), + ] + result = add_messages(left, right) + expected_result = [] + assert result == expected_result + + +class MessagesStatePydantic(BaseModel): + messages: Annotated[list[AnyMessage], add_messages] + + +MESSAGES_STATE_SCHEMAS = [MessagesState, MessagesStatePydantic] + + +@pytest.mark.parametrize("state_schema", MESSAGES_STATE_SCHEMAS) +def test_messages_state(state_schema): + def foo(state): + return {"messages": [HumanMessage("foo")]} + + graph = StateGraph(state_schema) + graph.add_edge(START, "foo") + graph.add_edge("foo", END) + graph.add_node(foo) + + app = graph.compile() + + assert app.invoke({"messages": [("user", "meow")]}) == { + "messages": [ + _AnyIdHumanMessage(content="meow"), + _AnyIdHumanMessage(content="foo"), + ] + } + + +@pytest.mark.skipif( + condition=not ((CORE_MINOR == 3 and CORE_PATCH >= 11) or CORE_MINOR > 3), + reason="Requires langchain_core>=0.3.11.", +) +def test_messages_state_format_openai(): + class State(TypedDict): + messages: Annotated[list[AnyMessage], add_messages(format="langchain-openai")] + + def foo(state): + messages = [ + HumanMessage( + content=[ + { + "type": "text", + "text": "Here's an image:", + "cache_control": {"type": "ephemeral"}, + }, + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/jpeg", + "data": "1234", + }, + }, + ] + ), + AIMessage( + content=[ + { + "type": "tool_use", + "name": "foo", + "input": {"bar": "baz"}, + "id": "1", + } + ] + ), + HumanMessage( + content=[ + { + "type": "tool_result", + "tool_use_id": "1", + "is_error": False, + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/jpeg", + "data": "1234", + }, + }, + ], + } + ] + ), + ] + return {"messages": messages} + + expected = [ + HumanMessage(content="meow"), + HumanMessage( + content=[ + {"type": "text", "text": "Here's an image:"}, + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,1234"}, + }, + ], + ), + AIMessage( + content="", + tool_calls=[ + { + "name": "foo", + "type": "tool_calls", + "args": {"bar": "baz"}, + "id": "1", + } + ], + ), + ToolMessage( + content=[ + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64,1234"}, + } + ], + tool_call_id="1", + ), + ] + + graph = StateGraph(State) + graph.add_edge(START, "foo") + graph.add_edge("foo", END) + graph.add_node(foo) + + app = graph.compile() + + result = app.invoke({"messages": [("user", "meow")]}) + 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"), + ] + + +def test_push_messages_in_graph(): + class MessagesState(TypedDict): + messages: Annotated[list[AnyMessage], add_messages] + + def chat(_: MessagesState) -> MessagesState: + with pytest.raises(ValueError, match="Message ID is required"): + push_message(AIMessage(content="No ID")) + + push_message(AIMessage(content="First", id="1")) + push_message(HumanMessage(content="Second", id="2")) + push_message(AIMessage(content="Third", id="3")) + + builder = StateGraph(MessagesState) + builder.add_node(chat) + builder.add_edge(START, "chat") + + graph = builder.compile() + + messages, values = [], None + for event, chunk in graph.stream( + {"messages": []}, stream_mode=["messages", "values"] + ): + if event == "values": + values = chunk + elif event == "messages": + message, _ = chunk + 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")]}) diff --git a/libs/langgraph/tests/test_validate_input_callable.py b/libs/langgraph/tests/test_validate_input_callable.py new file mode 100644 index 000000000..8749f0da2 --- /dev/null +++ b/libs/langgraph/tests/test_validate_input_callable.py @@ -0,0 +1,173 @@ +"""Tests for validate_input with callable approach.""" + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, RemoveMessage +from langgraph.checkpoint.memory import MemorySaver + +from langgraph.graph import StateGraph +from langgraph.graph.message import MessagesState, validate_messages_append_only +from langgraph.types import Command + + +def test_validate_messages_append_only_allows_new(): + """Test that validate_messages_append_only allows adding new messages.""" + + def chatbot(state: MessagesState) -> MessagesState: + return {"messages": [AIMessage(content="Response")]} + + builder = StateGraph(MessagesState) + builder.add_node("chatbot", chatbot) + builder.set_entry_point("chatbot") + builder.set_finish_point("chatbot") + + graph = builder.compile(validate_input=validate_messages_append_only) + + result = graph.invoke({"messages": [HumanMessage(content="Hello", id="1")]}) + assert len(result["messages"]) == 2 + + +def test_validate_messages_append_only_blocks_updates(): + """Test that validate_messages_append_only blocks message updates.""" + + def chatbot(state: MessagesState) -> MessagesState: + return {"messages": [AIMessage(content="Response")]} + + builder = StateGraph(MessagesState) + builder.add_node("chatbot", chatbot) + builder.set_entry_point("chatbot") + builder.set_finish_point("chatbot") + + checkpointer = MemorySaver() + graph = builder.compile( + checkpointer=checkpointer, validate_input=validate_messages_append_only + ) + + config = {"configurable": {"thread_id": "test"}} + graph.invoke({"messages": [HumanMessage(content="Hello", id="1")]}, config) + + # Try to update - should fail + with pytest.raises(ValueError, match="Cannot update existing message"): + graph.invoke({"messages": [HumanMessage(content="Modified", id="1")]}, config) + + +def test_validate_messages_append_only_blocks_removals(): + """Test that validate_messages_append_only blocks removals.""" + + def chatbot(state: MessagesState) -> MessagesState: + return {"messages": []} + + builder = StateGraph(MessagesState) + builder.add_node("chatbot", chatbot) + builder.set_entry_point("chatbot") + builder.set_finish_point("chatbot") + + checkpointer = MemorySaver() + graph = builder.compile( + checkpointer=checkpointer, validate_input=validate_messages_append_only + ) + + config = {"configurable": {"thread_id": "test"}} + graph.invoke({"messages": [HumanMessage(content="Hello", id="1")]}, config) + + # Try to remove - should fail + with pytest.raises(ValueError, match="Cannot remove existing message"): + graph.invoke({"messages": [RemoveMessage(id="1")]}, config) + + +def test_validate_with_command_update(): + """Test that Command.update is also validated.""" + + def chatbot(state: MessagesState) -> MessagesState: + return {"messages": [AIMessage(content="Response")]} + + builder = StateGraph(MessagesState) + builder.add_node("chatbot", chatbot) + builder.set_entry_point("chatbot") + builder.set_finish_point("chatbot") + + checkpointer = MemorySaver() + graph = builder.compile( + checkpointer=checkpointer, validate_input=validate_messages_append_only + ) + + config = {"configurable": {"thread_id": "test"}} + graph.invoke({"messages": [HumanMessage(content="Hello", id="1")]}, config) + + # Try to update via Command - should fail + with pytest.raises(ValueError, match="Cannot update existing message"): + graph.invoke( + Command(update={"messages": [HumanMessage(content="Modified", id="1")]}), + config, + ) + + +def test_custom_validator(): + """Test with custom validation function.""" + + def my_validator(input: dict, current_state: dict) -> None: + if "forbidden" in str(input): + raise ValueError("Forbidden content") + + def node(state: MessagesState) -> MessagesState: + return {"messages": [AIMessage(content="OK")]} + + builder = StateGraph(MessagesState) + builder.add_node("node", node) + builder.set_entry_point("node") + builder.set_finish_point("node") + + graph = builder.compile(validate_input=my_validator) + + # Should work + graph.invoke({"messages": [("user", "Hello")]}) + + # Should fail + with pytest.raises(ValueError, match="Forbidden content"): + graph.invoke({"messages": [("user", "forbidden word")]}) + + +def test_node_bypasses_validation(): + """Test that internal node operations bypass validation.""" + + def node_that_modifies(state: MessagesState) -> MessagesState: + # Node modifies existing message - should work even with strict validator + messages = state["messages"] + if messages: + updated = messages[0].model_copy(update={"content": "Modified by node"}) + return {"messages": [updated]} + return {"messages": []} + + builder = StateGraph(MessagesState) + builder.add_node("modifier", node_that_modifies) + builder.set_entry_point("modifier") + builder.set_finish_point("modifier") + + graph = builder.compile(validate_input=validate_messages_append_only) + + # Node can modify - validation only applies to external input + result = graph.invoke({"messages": [HumanMessage(content="Original", id="1")]}) + assert any("Modified by node" in m.content for m in result["messages"]) + + +def test_without_validator_allows_everything(): + """Test backwards compatibility - without validator, everything is allowed.""" + + def chatbot(state: MessagesState) -> MessagesState: + return {"messages": []} + + builder = StateGraph(MessagesState) + builder.add_node("chatbot", chatbot) + builder.set_entry_point("chatbot") + builder.set_finish_point("chatbot") + + checkpointer = MemorySaver() + graph = builder.compile(checkpointer=checkpointer) # No validator + + config = {"configurable": {"thread_id": "test"}} + graph.invoke({"messages": [HumanMessage(content="Hello", id="1")]}, config) + + # Should allow update (old behavior) + result = graph.invoke( + {"messages": [HumanMessage(content="Modified", id="1")]}, config + ) + assert result["messages"][0].content == "Modified"