WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-30 11:49:38 +02:00
Code Issues Packages Projects Releases Wiki Activity

switch to validating on input only with a custom helper

Browse Source
This commit is contained in:
Josh Rogers
2026-02-09 15:07:06 -05:00
parent b60d68d69c
commit dd17e7fd6b
8 changed files with 833 additions and 181 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/__init__.py
+7 -1
View File
@@ -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",
)
libs/langgraph/langgraph/graph/message.py
+87 -42
View File
@@ -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,
libs/langgraph/langgraph/graph/state.py
+30
View File
@@ -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)
libs/langgraph/langgraph/pregel/_loop.py
+21
View File
@@ -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:
libs/langgraph/langgraph/pregel/main.py
+4
View File
@@ -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(
libs/langgraph/tests/test_messages_state.py
+6 -138
View File
@@ -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")]})
libs/langgraph/tests/test_messages_state.py.bak
+505
View File
@@ -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")]})
libs/langgraph/tests/test_validate_input_callable.py
+173
View File
@@ -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"