WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-11 12:17:53 +02:00
Code Issues Packages Projects Releases Wiki Activity

add more tests

Browse Source
This commit is contained in:
vbarda
2024-07-22 16:23:54 -04:00
parent 6e57fa5f9a
commit fbced023f3
4 changed files with 1822 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/checkpoint/memory.py
+2 -1
View File
@@ -72,7 +72,8 @@ class MemorySaver(BaseCheckpointSaver):
thread_id = config["configurable"]["thread_id"]
if thread_ts := config["configurable"].get("thread_ts"):
if checkpoints := self.storage[thread_id]:
ts = max(key for key in checkpoints.keys() if key <= thread_ts)
matching_keys = [key for key in checkpoints.keys() if key <= thread_ts]
ts = max(matching_keys) if matching_keys else None
if saved := self.storage[thread_id].get(ts):
checkpoint, metadata = saved
writes = self.writes[(thread_id, ts)]
libs/langgraph/tests/test_pregel.py
+893 -1
View File
@@ -67,6 +67,7 @@ from tests.memory_assert import (
MemorySaverAssertImmutable,
NoopSerializer,
)
from tests.utils import assert_state_history_equal
def test_graph_validation() -> None:
@@ -7586,10 +7587,161 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None:
assert app.invoke({"my_key": "my value"}, config, debug=True) == {
"my_key": "hi my value",
}
assert_state_history_equal(
list(app.get_state_history(config)),
[
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
assert app.invoke(None, config, debug=True) == {
"my_key": "hi my value here and there and back again",
}
assert_state_history_equal(
list(app.get_state_history(config)),
[
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"outer_2": {
"my_key": "hi my value here and there and back again"
}
},
"step": 4,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"inner": {"my_key": "hi my value here and there"}},
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# test stream updates w/ nested interrupt
config = {"configurable": {"thread_id": "2"}}
@@ -7631,8 +7783,102 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None:
"my_key": "hi my value",
},
]
assert_state_history_equal(
list(app.get_state_history(config)),
[
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# while we're waiting for the node w/ interrupt inside to finish
assert [*app.stream(None, config, stream_mode="values")] == []
assert_state_history_equal(
list(app.get_state_history(config)),
[
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
assert [*app.stream(None, config, stream_mode="values")] == [
{
"my_key": "hi my value here and there",
@@ -7641,6 +7887,106 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None:
"my_key": "hi my value here and there and back again",
},
]
assert_state_history_equal(
list(app.get_state_history(config)),
[
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "4",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"outer_2": {
"my_key": "hi my value here and there and back again"
}
},
"step": 4,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "4",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "4",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"inner": {"my_key": "hi my value here and there"}},
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "4",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# test interrupts AFTER the node w/ interrupts
app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"])
@@ -7653,16 +7999,562 @@ def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> None:
"my_key": "hi my value",
},
]
assert_state_history_equal(
list(app.get_state_history(config)),
[
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
assert [*app.stream(None, config, stream_mode="values")] == [
{
"my_key": "hi my value here and there",
},
]
assert_state_history_equal(
list(app.get_state_history(config)),
[
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"inner": {"my_key": "hi my value here and there"}},
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
assert [*app.stream(None, config, stream_mode="values")] == [
{
"my_key": "hi my value here and there and back again",
},
]
assert_state_history_equal(
list(app.get_state_history(config)),
[
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"outer_2": {
"my_key": "hi my value here and there and back again"
}
},
"step": 4,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"inner": {"my_key": "hi my value here and there"}},
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# test restarting from thread_ts
config = {"configurable": {"thread_id": "6"}}
app = graph.compile(checkpointer=checkpointer)
app.invoke({"my_key": "my value"}, config, debug=True)
state_history = [c for c in app.get_state_history(config)]
assert_state_history_equal(
state_history,
[
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
child_state_history = [
c for c in app.get_state_history({"configurable": {"thread_id": "6-inner"}})
]
assert_state_history_equal(
child_state_history,
[
StateSnapshot(
values={"my_key": "hi my value here"},
next=(),
config={
"configurable": {
"thread_id": "6-inner",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"inner_1": {
"my_key": "hi my value here",
"my_other_key": "hi my value",
}
},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "6-inner",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=(),
config={
"configurable": {
"thread_id": "6-inner",
"thread_ts": AnyStr(),
}
},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "6-inner",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={
"configurable": {
"thread_id": "6-inner",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "input",
"writes": {"my_key": "hi my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# check that parent snapshot is always older than child
child_snapshot = child_state_history[0]
assert (
child_snapshot.config["configurable"]["thread_ts"]
< state_history[0].config["configurable"]["thread_ts"]
)
# check resuming from interrupt w/ thread_ts
interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2]
before_interrupt_config = before_interrupt_state_snapshot.config
# going to get to interrupt again here, so the output is None
assert app.invoke(None, before_interrupt_config, debug=True) is None
assert_state_history_equal(
list(app.get_state_history(config)),
[
# NOTE: there is an identical snapshot here since we replayed from before interrupt
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# going to restart from interrupt
interrupt_config = interrupt_state_snapshot.config
assert app.invoke(None, interrupt_config, debug=True) == {
"my_key": "hi my value here and there and back again",
}
assert_state_history_equal(
list(app.get_state_history(config)),
[
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "6",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"outer_2": {
"my_key": "hi my value here and there and back again"
}
},
"step": 4,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "6",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "6",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"inner": {"my_key": "hi my value here and there"}},
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "6",
"thread_ts": AnyStr(),
}
},
),
# NOTE: there is an identical snapshot here since we replayed from before interrupt
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
finally:
if hasattr(checkpointer, "__exit__"):
checkpointer.__exit__(None, None, None)
libs/langgraph/tests/test_pregel_async.py
+910 -2
View File
@@ -61,6 +61,7 @@ from tests.memory_assert import (
MemorySaverAssertCheckpointMetadata,
MemorySaverAssertImmutable,
)
from tests.utils import assert_state_history_equal
async def test_checkpoint_errors() -> None:
@@ -6078,11 +6079,163 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non
assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == {
"my_key": "hi my value",
}
await asyncio.sleep(0.05)
assert_state_history_equal(
[s async for s in app.aget_state_history(config)],
[
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
assert await app.ainvoke(None, config, debug=True) == {
"my_key": "hi my value here and there and back again",
}
await asyncio.sleep(0.05)
assert_state_history_equal(
[s async for s in app.aget_state_history(config)],
[
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"outer_2": {
"my_key": "hi my value here and there and back again"
}
},
"step": 4,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"inner": {"my_key": "hi my value here and there"}},
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "1",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "1", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# test stream updates w/ nested interrupt
config = {"configurable": {"thread_id": "2"}}
assert [c async for c in app.astream({"my_key": "my value"}, config)] == [
@@ -6133,8 +6286,104 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non
"my_key": "hi my value",
},
]
await asyncio.sleep(0.05)
assert_state_history_equal(
[s async for s in app.aget_state_history(config)],
[
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# while we're waiting for the node w/ interrupt inside to finish
assert [c async for c in app.astream(None, config, stream_mode="values")] == []
await asyncio.sleep(0.05)
assert_state_history_equal(
[s async for s in app.aget_state_history(config)],
[
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
{
"my_key": "hi my value here and there",
@@ -6143,6 +6392,107 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non
"my_key": "hi my value here and there and back again",
},
]
await asyncio.sleep(0.05)
assert_state_history_equal(
[s async for s in app.aget_state_history(config)],
[
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "4",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"outer_2": {
"my_key": "hi my value here and there and back again"
}
},
"step": 4,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "4",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "4",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"inner": {"my_key": "hi my value here and there"}},
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "4",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "4", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "4", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# test interrupts AFTER the node w/ interrupts
app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"])
@@ -6160,16 +6510,574 @@ async def test_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> Non
"my_key": "hi my value",
},
]
await asyncio.sleep(0.05)
assert_state_history_equal(
[s async for s in app.aget_state_history(config)],
[
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
{
"my_key": "hi my value here and there",
},
]
await asyncio.sleep(0.05)
assert_state_history_equal(
[s async for s in app.aget_state_history(config)],
[
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"inner": {"my_key": "hi my value here and there"}},
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
assert [c async for c in app.astream(None, config, stream_mode="values")] == [
{
"my_key": "hi my value here and there and back again",
},
]
await asyncio.sleep(0.05)
assert_state_history_equal(
[s async for s in app.aget_state_history(config)],
[
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"outer_2": {
"my_key": "hi my value here and there and back again"
}
},
"step": 4,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"inner": {"my_key": "hi my value here and there"}},
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "5",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "5", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "5", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# test restarting from thread_ts
config = {"configurable": {"thread_id": "6"}}
app = graph.compile(checkpointer=checkpointer)
await app.ainvoke({"my_key": "my value"}, config, debug=True)
await asyncio.sleep(0.05)
state_history = [c async for c in app.aget_state_history(config)]
(
state_history,
[
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
# ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
await asyncio.sleep(0.05)
child_state_history = [
c
async for c in app.aget_state_history(
{"configurable": {"thread_id": "6-inner"}}
)
]
assert_state_history_equal(
child_state_history,
[
StateSnapshot(
values={"my_key": "hi my value here"},
next=(),
config={
"configurable": {
"thread_id": "6-inner",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"inner_1": {
"my_key": "hi my value here",
"my_other_key": "hi my value",
}
},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "6-inner",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=(),
config={
"configurable": {
"thread_id": "6-inner",
"thread_ts": AnyStr(),
}
},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "6-inner",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={
"configurable": {
"thread_id": "6-inner",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "input",
"writes": {"my_key": "hi my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# check that parent snapshot is always older than child
child_snapshot = child_state_history[0]
assert (
child_snapshot.config["configurable"]["thread_ts"]
< state_history[0].config["configurable"]["thread_ts"]
)
# check resuming from interrupt w/ thread_ts
interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2]
before_interrupt_config = before_interrupt_state_snapshot.config
# going to get to interrupt again here, so the output is None
assert (await app.ainvoke(None, before_interrupt_config, debug=True)) is None
await asyncio.sleep(0.05)
assert_state_history_equal(
[s async for s in app.aget_state_history(config)],
[
# NOTE: there is an identical snapshot here since we replayed from before interrupt
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
# going to restart from interrupt
interrupt_config = interrupt_state_snapshot.config
assert (await app.ainvoke(None, interrupt_config, debug=True)) == {
"my_key": "hi my value here and there and back again",
}
await asyncio.sleep(0.05)
assert_state_history_equal(
[s async for s in app.aget_state_history(config)],
[
StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
next=(),
config={
"configurable": {
"thread_id": "6",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {
"outer_2": {
"my_key": "hi my value here and there and back again"
}
},
"step": 4,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "6",
"thread_ts": AnyStr(),
}
},
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
next=("outer_2",),
config={
"configurable": {
"thread_id": "6",
"thread_ts": AnyStr(),
}
},
metadata={
"source": "loop",
"writes": {"inner": {"my_key": "hi my value here and there"}},
"step": 3,
},
created_at=AnyStr(),
parent_config={
"configurable": {
"thread_id": "6",
"thread_ts": AnyStr(),
}
},
),
# NOTE: there is an identical snapshot here since we replayed from before interrupt
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "interrupt", "step": 2},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "hi my value"},
next=("inner",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "loop",
"writes": {"outer_1": {"my_key": "hi my value"}},
"step": 1,
},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={"my_key": "my value"},
next=("outer_1",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={"source": "loop", "writes": None, "step": 0},
created_at=AnyStr(),
parent_config={
"configurable": {"thread_id": "6", "thread_ts": AnyStr()}
},
),
StateSnapshot(
values={},
next=("__start__",),
config={"configurable": {"thread_id": "6", "thread_ts": AnyStr()}},
metadata={
"source": "input",
"writes": {"my_key": "my value"},
"step": -1,
},
created_at=AnyStr(),
parent_config=None,
),
],
ignore_parent_config=isinstance(checkpointer, MemorySaver),
)
finally:
if hasattr(checkpointer, "__aexit__"):
await checkpointer.__aexit__(None, None, None)
libs/langgraph/tests/utils.py
+17
View File
@@ -0,0 +1,17 @@
from langgraph.pregel import StateSnapshot
def assert_state_history_equal(
actual_state_history: list[StateSnapshot],
expected_state_history: list[StateSnapshot],
ignore_parent_config: bool = False,
) -> None:
assert (
len(actual_state_history) == len(expected_state_history)
), f"Got different lengths for state history: {len(actual_state_history)} for actual, {len(expected_state_history)} for expected"
for actual, expected in zip(actual_state_history, expected_state_history):
if ignore_parent_config:
actual = actual._replace(parent_config=None)
expected = expected._replace(parent_config=None)
assert actual == expected