WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-23 16:12:25 +02:00
Code Issues Packages Projects Releases Wiki Activity

Flip default for checkpoint_during

Browse Source 
- Now defaulting to False, ie. saving only the final checkpoint
- All features other than time travel into an intermediate step are supported by checkpoint_during=False so this is a better default for almost all use cases
This commit is contained in:
Nuno Campos
2025-05-29 11:45:54 -07:00
parent 045d6dfb82
commit 0afc4ebda3
8 changed files with 380 additions and 1992 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+4 -8
View File
@@ -2387,7 +2387,7 @@ class Pregel(PregelProtocol):
output_keys: The keys to stream, defaults to all non-context channels.
interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph.
interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph.
checkpoint_during: Whether to checkpoint intermediate steps, defaults to True. If False, only the final checkpoint is saved.
checkpoint_during: Whether to checkpoint intermediate steps, defaults to False. If False, only the final checkpoint is saved.
debug: Whether to print debug information during execution, defaults to False.
subgraphs: Whether to stream events from inside subgraphs, defaults to False.
If True, the events will be emitted as tuples `(namespace, data)`,
@@ -2494,7 +2494,7 @@ class Pregel(PregelProtocol):
debug=debug,
checkpoint_during=checkpoint_during
if checkpoint_during is not None
else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True),
else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, False),
trigger_to_nodes=self.trigger_to_nodes,
migrate_checkpoint=self._migrate_checkpoint,
retry_policy=self.retry_policy,
@@ -2608,7 +2608,7 @@ class Pregel(PregelProtocol):
output_keys: The keys to stream, defaults to all non-context channels.
interrupt_before: Nodes to interrupt before, defaults to all nodes in the graph.
interrupt_after: Nodes to interrupt after, defaults to all nodes in the graph.
checkpoint_during: Whether to checkpoint intermediate steps, defaults to True. If False, only the final checkpoint is saved.
checkpoint_during: Whether to checkpoint intermediate steps, defaults to False. If False, only the final checkpoint is saved.
debug: Whether to print debug information during execution, defaults to False.
subgraphs: Whether to stream events from inside subgraphs, defaults to False.
If True, the events will be emitted as tuples `(namespace, data)`,
@@ -2737,7 +2737,7 @@ class Pregel(PregelProtocol):
debug=debug,
checkpoint_during=checkpoint_during
if checkpoint_during is not None
else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True),
else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, False),
trigger_to_nodes=self.trigger_to_nodes,
migrate_checkpoint=self._migrate_checkpoint,
retry_policy=self.retry_policy,
@@ -2816,7 +2816,6 @@ class Pregel(PregelProtocol):
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
checkpoint_during: bool | None = None,
debug: bool | None = None,
**kwargs: Any,
) -> dict[str, Any] | Any:
@@ -2849,7 +2848,6 @@ class Pregel(PregelProtocol):
output_keys=output_keys,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
checkpoint_during=checkpoint_during,
debug=debug,
**kwargs,
):
@@ -2884,7 +2882,6 @@ class Pregel(PregelProtocol):
output_keys: str | Sequence[str] | None = None,
interrupt_before: All | Sequence[str] | None = None,
interrupt_after: All | Sequence[str] | None = None,
checkpoint_during: bool | None = None,
debug: bool | None = None,
**kwargs: Any,
) -> dict[str, Any] | Any:
@@ -2918,7 +2915,6 @@ class Pregel(PregelProtocol):
output_keys=output_keys,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
checkpoint_during=checkpoint_during,
debug=debug,
**kwargs,
):
libs/langgraph/langgraph/pregel/debug.py
+17 -2
View File
@@ -144,6 +144,19 @@ def map_debug_task_results(
}
def rm_pregel_keys(config: Optional[RunnableConfig]) -> Optional[RunnableConfig]:
"""Remove pregel-specific keys from the config."""
if config is None:
return config
return {
"configurable": {
k: v
for k, v in config.get("configurable", {}).items()
if not k.startswith("__pregel_")
}
}
def map_debug_checkpoint(
step: int,
config: RunnableConfig,
@@ -183,8 +196,10 @@ def map_debug_checkpoint(
"timestamp": checkpoint["ts"],
"step": step,
"payload": {
"config": patch_checkpoint_map(config, metadata),
"parent_config": patch_checkpoint_map(parent_config, metadata),
"config": rm_pregel_keys(patch_checkpoint_map(config, metadata)),
"parent_config": rm_pregel_keys(
patch_checkpoint_map(parent_config, metadata)
),
"values": read_channels(channels, stream_channels),
"metadata": metadata,
"next": [t.name for t in tasks],
libs/langgraph/langgraph/pregel/loop.py
+7 -2
View File
@@ -564,7 +564,13 @@ class PregelLoop:
"debug",
map_debug_checkpoint,
self.step - 1, # printing checkpoint for previous step
self.checkpoint_config,
{
**self.checkpoint_config,
CONF: {
**self.checkpoint_config[CONF],
CONFIG_KEY_CHECKPOINT_ID: self.checkpoint["id"],
},
},
self.channels,
self.stream_keys,
self.checkpoint_metadata,
@@ -819,7 +825,6 @@ class PregelLoop:
**self.checkpoint_config,
CONF: {
**self.checkpoint_config[CONF],
# this is guaranteed to be set by code above
CONFIG_KEY_CHECKPOINT_NS: self.config[CONF].get(
CONFIG_KEY_CHECKPOINT_NS, ""
),
libs/langgraph/tests/test_checkpoint_migration.py
+11 -4
View File
@@ -1541,7 +1541,9 @@ def test_latest_checkpoint_state_graph(
app = builder.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "1"}}
assert [*app.stream({"query": "what is weather in sf"}, config)] == [
assert [
*app.stream({"query": "what is weather in sf"}, config, checkpoint_during=True)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
{"retriever_two": {"docs": ["doc3", "doc4"]}},
@@ -1557,7 +1559,7 @@ def test_latest_checkpoint_state_graph(
},
]
assert [*app.stream(Command(resume=""), config)] == [
assert [*app.stream(Command(resume=""), config, checkpoint_during=True)] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
@@ -1582,7 +1584,10 @@ async def test_latest_checkpoint_state_graph_async(
config = {"configurable": {"thread_id": "1"}}
assert [
c async for c in app.astream({"query": "what is weather in sf"}, config)
c
async for c in app.astream(
{"query": "what is weather in sf"}, config, checkpoint_during=True
)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
@@ -1599,7 +1604,9 @@ async def test_latest_checkpoint_state_graph_async(
},
]
assert [c async for c in app.astream(Command(resume=""), config)] == [
assert [
c async for c in app.astream(Command(resume=""), config, checkpoint_during=True)
] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
libs/langgraph/tests/test_large_cases.py
+69 -903
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_large_cases_async.py
+119 -914
View File
File diff suppressed because it is too large Load Diff
libs/langgraph/tests/test_pregel.py
+74 -63
View File
@@ -263,7 +263,9 @@ def test_checkpoint_errors() -> None:
builder.add_edge(START, "parallel")
graph = builder.compile(checkpointer=FaultyPutWritesCheckpointer())
with pytest.raises(ValueError, match="Faulty put_writes"):
graph.invoke("", {"configurable": {"thread_id": "thread-1"}})
graph.invoke(
"", {"configurable": {"thread_id": "thread-1"}}, checkpoint_during=True
)
@pytest.mark.parametrize("use_node_builder", [True, False])
@@ -749,7 +751,7 @@ def test_run_from_checkpoint_id_retains_previous_writes(
thread_id = uuid.uuid4()
thread1 = {"configurable": {"thread_id": str(thread_id)}}
result = graph.invoke({"myval": 1}, thread1)
result = graph.invoke({"myval": 1}, thread1, checkpoint_during=True)
assert result["myval"] == 4
history = [c for c in graph.get_state_history(thread1)]
@@ -1611,7 +1613,7 @@ def test_invoke_checkpoint_three(
thread_1 = {"configurable": {"thread_id": "1"}}
# total starts out as 0, so output is 0+2=2
assert app.invoke(2, thread_1, debug=1) == 2
assert app.invoke(2, thread_1, checkpoint_during=True) == 2
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 2
@@ -1621,7 +1623,7 @@ def test_invoke_checkpoint_three(
== sync_checkpointer.get(thread_1)["id"]
)
# total is now 2, so output is 2+3=5
assert app.invoke(3, thread_1) == 5
assert app.invoke(3, thread_1, checkpoint_during=True) == 5
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
@@ -1631,7 +1633,7 @@ def test_invoke_checkpoint_three(
)
# total is now 2+5=7, so output would be 7+4=11, but raises ValueError
with pytest.raises(ValueError):
app.invoke(4, thread_1)
app.invoke(4, thread_1, checkpoint_during=True)
# checkpoint is updated with new input
state = app.get_state(thread_1)
assert state is not None
@@ -1639,7 +1641,7 @@ def test_invoke_checkpoint_three(
assert state.next == ("one",)
"""we checkpoint inputs and it failed on "one", so the next node is one"""
# we can recover from error by sending new inputs
assert app.invoke(2, thread_1) == 9
assert app.invoke(2, thread_1, checkpoint_during=True) == 9
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 16, "total is now 7+9=16"
@@ -1647,8 +1649,8 @@ def test_invoke_checkpoint_three(
thread_2 = {"configurable": {"thread_id": "2"}}
# on a new thread, total starts out as 0, so output is 0+5=5
assert app.invoke(5, thread_2, debug=True) == 5
state = app.get_state({"configurable": {"thread_id": "1"}})
assert app.invoke(5, thread_2) == 5
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 16
assert state.next == (), "checkpoint of other thread not touched"
@@ -4848,7 +4850,7 @@ def test_debug_retry(sync_checkpointer: BaseCheckpointSaver):
graph = builder.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "1"}}
graph.invoke({"messages": []}, config=config)
graph.invoke({"messages": []}, config=config, checkpoint_during=True)
# re-run step: 1
target_config = next(
@@ -4858,7 +4860,11 @@ def test_debug_retry(sync_checkpointer: BaseCheckpointSaver):
)
update_config = graph.update_state(target_config, values=None)
events = [*graph.stream(None, config=update_config, stream_mode="debug")]
events = [
*graph.stream(
None, config=update_config, stream_mode="debug", checkpoint_during=True
)
]
checkpoint_events = list(
reversed([e["payload"] for e in events if e["type"] == "checkpoint"])
@@ -4888,7 +4894,9 @@ def test_debug_retry(sync_checkpointer: BaseCheckpointSaver):
assert stream_parent_conf == history_parent_conf
def test_debug_subgraphs(sync_checkpointer: BaseCheckpointSaver):
def test_debug_subgraphs(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
):
class State(TypedDict):
messages: Annotated[list[str], operator.add]
@@ -4921,12 +4929,15 @@ def test_debug_subgraphs(sync_checkpointer: BaseCheckpointSaver):
{"messages": []},
config=config,
stream_mode="debug",
checkpoint_during=checkpoint_during,
)
]
checkpoint_events = list(
reversed([e["payload"] for e in events if e["type"] == "checkpoint"])
)
if not checkpoint_during:
checkpoint_events = checkpoint_events[:1]
checkpoint_history = list(graph.get_state_history(config))
assert len(checkpoint_events) == len(checkpoint_history)
@@ -4955,7 +4966,9 @@ def test_debug_subgraphs(sync_checkpointer: BaseCheckpointSaver):
assert stream_task.get("state") == history_task.state
def test_debug_nested_subgraphs(sync_checkpointer: BaseCheckpointSaver):
def test_debug_nested_subgraphs(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
):
from collections import defaultdict
class State(TypedDict):
@@ -4998,6 +5011,7 @@ def test_debug_nested_subgraphs(sync_checkpointer: BaseCheckpointSaver):
config=config,
stream_mode="debug",
subgraphs=True,
checkpoint_during=checkpoint_during,
)
]
@@ -5037,6 +5051,9 @@ def test_debug_nested_subgraphs(sync_checkpointer: BaseCheckpointSaver):
for checkpoint_events, checkpoint_history in zip(
stream_ns.values(), history_ns.values()
):
if not checkpoint_during:
checkpoint_events = checkpoint_events[-1:]
assert len(checkpoint_events) == len(checkpoint_history)
for stream, history in zip(checkpoint_events, checkpoint_history):
assert stream["values"] == history.values
assert stream["next"] == list(history.next)
@@ -5263,15 +5280,7 @@ def test_parent_command(sync_checkpointer: BaseCheckpointSaver) -> None:
"parents": {},
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
parent_config=None,
tasks=(),
interrupts=(),
)
@@ -5831,7 +5840,9 @@ def test_concurrent_execution_thread_safety():
assert result["counter"] == 1
def test_checkpoint_recovery(sync_checkpointer: BaseCheckpointSaver):
def test_checkpoint_recovery(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
):
"""Test recovery from checkpoints after failures."""
class State(TypedDict):
@@ -5858,7 +5869,11 @@ def test_checkpoint_recovery(sync_checkpointer: BaseCheckpointSaver):
# First attempt should fail
with pytest.raises(RuntimeError):
graph.invoke({"steps": ["start"], "attempt": 1}, config)
graph.invoke(
{"steps": ["start"], "attempt": 1},
config,
checkpoint_during=checkpoint_during,
)
# Verify checkpoint state
state = graph.get_state(config)
@@ -5868,12 +5883,17 @@ def test_checkpoint_recovery(sync_checkpointer: BaseCheckpointSaver):
assert "RuntimeError('Simulated failure')" in state.tasks[0].error
# Retry with updated attempt count
result = graph.invoke({"steps": [], "attempt": 2}, config)
result = graph.invoke(
{"steps": [], "attempt": 2}, config, checkpoint_during=checkpoint_during
)
assert result == {"steps": ["start", "node1", "node2"], "attempt": 2}
# Verify checkpoint history shows both attempts
history = list(graph.get_state_history(config))
assert len(history) == 6 # Initial + failed attempt + successful attempt
if checkpoint_during:
assert len(history) == 6 # Initial + failed attempt + successful attempt
else:
assert len(history) == 2 # error + success
# Verify the error was recorded in checkpoint
failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error)
@@ -5936,9 +5956,7 @@ def test_multiple_updates() -> None:
]
def test_falsy_return_from_task(
sync_checkpointer: BaseCheckpointSaver, snapshot: SnapshotAssertion
):
def test_falsy_return_from_task(sync_checkpointer: BaseCheckpointSaver):
"""Test with a falsy return from a task."""
@task
@@ -5958,15 +5976,11 @@ def test_falsy_return_from_task(
{
"payload": {
"config": {
"callbacks": None,
"configurable": {
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
"thread_id": AnyStr(),
},
"metadata": {},
"recursion_limit": 25,
"tags": [],
},
"metadata": {
"parents": {},
@@ -6057,7 +6071,6 @@ def test_falsy_return_from_task(
"type": "task_result",
},
]
print(type(configurable["configurable"]["thread_id"]))
assert [
c
for c in graph.stream(Command(resume="123"), configurable, stream_mode="debug")
@@ -6065,15 +6078,11 @@ def test_falsy_return_from_task(
{
"payload": {
"config": {
"callbacks": None,
"configurable": {
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
"thread_id": AnyStr(),
},
"metadata": {},
"recursion_limit": 25,
"tags": [],
},
"metadata": {
"parents": {},
@@ -6155,15 +6164,11 @@ def test_falsy_return_from_task(
{
"payload": {
"config": {
"callbacks": None,
"configurable": {
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
"thread_id": AnyStr(),
},
"metadata": {},
"recursion_limit": 25,
"tags": [],
},
"metadata": {
"parents": {},
@@ -6171,17 +6176,7 @@ def test_falsy_return_from_task(
"step": 0,
},
"next": [],
"parent_config": {
"callbacks": None,
"configurable": {
"checkpoint_id": AnyStr(),
"checkpoint_ns": "",
"thread_id": AnyStr(),
},
"metadata": {},
"recursion_limit": 25,
"tags": [],
},
"parent_config": None,
"tasks": [],
"values": None,
},
@@ -8089,7 +8084,9 @@ def test_pregel_node_copy() -> None:
graph.nodes["agent"].copy({})
def test_update_as_input(sync_checkpointer: BaseCheckpointSaver) -> None:
def test_update_as_input(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
) -> None:
class State(TypedDict):
foo: str
@@ -8108,13 +8105,17 @@ def test_update_as_input(sync_checkpointer: BaseCheckpointSaver) -> None:
.compile(checkpointer=sync_checkpointer)
)
assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == {
"foo": "tool"
}
assert graph.invoke(
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
) == {"foo": "tool"}
assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == {
"foo": "tool"
}
assert graph.invoke(
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
) == {"foo": "tool"}
def map_snapshot(i: StateSnapshot) -> dict:
return {
@@ -8152,11 +8153,14 @@ def test_update_as_input(sync_checkpointer: BaseCheckpointSaver) -> None:
for s in graph.get_state_history({"configurable": {"thread_id": "2"}})
]
assert new_history == history
if checkpoint_during:
assert new_history == history
else:
assert [new_history[0], new_history[4]] == history
def test_batch_update_as_input(
sync_checkpointer: BaseCheckpointSaver,
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
) -> None:
class State(TypedDict):
foo: str
@@ -8188,7 +8192,11 @@ def test_batch_update_as_input(
.compile(checkpointer=sync_checkpointer)
)
assert graph.invoke({"foo": "input"}, {"configurable": {"thread_id": "1"}}) == {
assert graph.invoke(
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
) == {
"foo": "map",
"tasks": [0, 1, 2],
}
@@ -8241,7 +8249,10 @@ def test_batch_update_as_input(
for s in graph.get_state_history({"configurable": {"thread_id": "2"}})
]
assert new_history == history
if checkpoint_during:
assert new_history == history
else:
assert new_history[:1] == history
def test_migration_graph(snapshot: SnapshotAssertion) -> None:
libs/langgraph/tests/test_pregel_async.py
+79 -96
View File
@@ -154,13 +154,20 @@ async def test_checkpoint_errors() -> None:
builder.add_edge(START, "parallel")
graph = builder.compile(checkpointer=FaultyPutWritesCheckpointer())
with pytest.raises(ValueError, match="Faulty put_writes"):
await graph.ainvoke("", {"configurable": {"thread_id": "thread-1"}})
await graph.ainvoke(
"", {"configurable": {"thread_id": "thread-1"}}, checkpoint_during=True
)
with pytest.raises(ValueError, match="Faulty put_writes"):
async for _ in graph.astream("", {"configurable": {"thread_id": "thread-2"}}):
async for _ in graph.astream(
"", {"configurable": {"thread_id": "thread-2"}}, checkpoint_during=True
):
pass
with pytest.raises(ValueError, match="Faulty put_writes"):
async for _ in graph.astream_events(
"", {"configurable": {"thread_id": "thread-3"}}, version="v2"
"",
{"configurable": {"thread_id": "thread-3"}},
version="v2",
checkpoint_during=True,
):
pass
@@ -271,7 +278,6 @@ async def test_checkpoint_put_after_cancellation() -> None:
# check logs before cancellation is handled
assert sorted(logs) == [
"awhile.start",
"checkpoint.aput.start",
], "Cancelled before checkpoint put started"
# wait for task to finish
try:
@@ -336,7 +342,6 @@ async def test_checkpoint_put_after_cancellation_stream_anext() -> None:
# check logs before cancellation is handled
assert sorted(logs) == [
"awhile.start",
"checkpoint.aput.start",
], "Cancelled before checkpoint put started"
# wait for task to finish
try:
@@ -403,7 +408,6 @@ async def test_checkpoint_put_after_cancellation_stream_events_anext() -> None:
t.cancel()
# check logs before cancellation is handled
assert logs == [
"checkpoint.aput.start",
"awhile.start",
], "Cancelled before checkpoint put started"
# wait for task to finish
@@ -412,9 +416,9 @@ async def test_checkpoint_put_after_cancellation_stream_events_anext() -> None:
except asyncio.CancelledError:
# check logs after cancellation is handled
assert logs == [
"checkpoint.aput.start",
"awhile.start",
"awhile.end",
"checkpoint.aput.start",
"checkpoint.aput.end",
], "Checkpoint put is not cancelled"
else:
@@ -590,12 +594,6 @@ async def test_dynamic_interrupt(async_checkpointer: BaseCheckpointSaver) -> Non
"step": 0,
"thread_id": "1",
},
{
"parents": {},
"source": "input",
"step": -1,
"thread_id": "1",
},
]
tup = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_state(thread1) == StateSnapshot(
@@ -623,9 +621,7 @@ async def test_dynamic_interrupt(async_checkpointer: BaseCheckpointSaver) -> Non
"step": 0,
"thread_id": "1",
},
parent_config=(
[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config
),
parent_config=None,
interrupts=(
Interrupt(
value="Just because...",
@@ -771,12 +767,6 @@ async def test_dynamic_interrupt_subgraph(
"step": 0,
"thread_id": "1",
},
{
"parents": {},
"source": "input",
"step": -1,
"thread_id": "1",
},
]
tup = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_state(thread1) == StateSnapshot(
@@ -810,11 +800,7 @@ async def test_dynamic_interrupt_subgraph(
"step": 0,
"thread_id": "1",
},
parent_config=(
[c async for c in tool_two.checkpointer.alist(thread1root, limit=2)][
-1
].config
),
parent_config=None,
interrupts=(
Interrupt(
value="Just because...",
@@ -959,12 +945,6 @@ async def test_copy_checkpoint(async_checkpointer: BaseCheckpointSaver) -> None:
"step": 0,
"thread_id": "1",
},
{
"parents": {},
"source": "input",
"step": -1,
"thread_id": "1",
},
]
tup = await tool_two.checkpointer.aget_tuple(thread1)
@@ -1002,9 +982,7 @@ async def test_copy_checkpoint(async_checkpointer: BaseCheckpointSaver) -> None:
"step": 0,
"thread_id": "1",
},
parent_config=(
[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config
),
parent_config=None,
interrupts=(
Interrupt(
value="Just because...",
@@ -1044,9 +1022,7 @@ async def test_copy_checkpoint(async_checkpointer: BaseCheckpointSaver) -> None:
"thread_id": "1",
},
parent_config=(
[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][
-1
].parent_config
[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config
),
interrupts=(),
)
@@ -2104,7 +2080,7 @@ async def test_run_from_checkpoint_id_retains_previous_writes(
thread_id = uuid.uuid4()
thread1 = {"configurable": {"thread_id": str(thread_id)}}
result = await graph.ainvoke({"myval": 1}, thread1)
result = await graph.ainvoke({"myval": 1}, thread1, checkpoint_during=True)
assert result["myval"] == 4
history = [c async for c in graph.aget_state_history(thread1)]
@@ -3061,15 +3037,7 @@ async def test_send_react_interrupt(async_checkpointer: BaseCheckpointSaver) ->
"thread_id": "2",
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
parent_config=None,
tasks=(
PregelTask(
id=AnyStr(),
@@ -3193,15 +3161,7 @@ async def test_send_react_interrupt(async_checkpointer: BaseCheckpointSaver) ->
"thread_id": "3",
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "3",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
parent_config=None,
tasks=(
PregelTask(
id=AnyStr(),
@@ -3466,15 +3426,7 @@ async def test_send_react_interrupt_control(
"thread_id": "2",
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "2",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
parent_config=None,
tasks=(
PregelTask(
id=AnyStr(),
@@ -3710,7 +3662,7 @@ async def test_invoke_checkpoint_three(
thread_1 = {"configurable": {"thread_id": "1"}}
# total starts out as 0, so output is 0+2=2
assert await app.ainvoke(2, thread_1) == 2
assert await app.ainvoke(2, thread_1, checkpoint_during=True) == 2
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 2
@@ -3719,7 +3671,7 @@ async def test_invoke_checkpoint_three(
== (await async_checkpointer.aget(thread_1))["id"]
)
# total is now 2, so output is 2+3=5
assert await app.ainvoke(3, thread_1) == 5
assert await app.ainvoke(3, thread_1, checkpoint_during=True) == 5
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
@@ -3729,7 +3681,7 @@ async def test_invoke_checkpoint_three(
)
# total is now 2+5=7, so output would be 7+4=11, but raises ValueError
with pytest.raises(ValueError):
await app.ainvoke(4, thread_1)
await app.ainvoke(4, thread_1, checkpoint_during=True)
# checkpoint is not updated
state = await app.aget_state(thread_1)
assert state is not None
@@ -3737,7 +3689,7 @@ async def test_invoke_checkpoint_three(
assert state.next == ("one",)
"""we checkpoint inputs and it failed on "one", so the next node is one"""
# we can recover from error by sending new inputs
assert await app.ainvoke(2, thread_1) == 9
assert await app.ainvoke(2, thread_1, checkpoint_during=True) == 9
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 16, "total is now 7+9=16"
@@ -3746,7 +3698,7 @@ async def test_invoke_checkpoint_three(
thread_2 = {"configurable": {"thread_id": "2"}}
# on a new thread, total starts out as 0, so output is 0+5=5
assert await app.ainvoke(5, thread_2) == 5
state = await app.aget_state({"configurable": {"thread_id": "1"}})
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 16
assert state.next == ()
@@ -6005,7 +5957,7 @@ async def test_debug_retry(async_checkpointer: BaseCheckpointSaver):
graph = builder.compile(checkpointer=async_checkpointer)
config = {"configurable": {"thread_id": "1"}}
await graph.ainvoke({"messages": []}, config=config)
await graph.ainvoke({"messages": []}, config=config, checkpoint_during=True)
# re-run step: 1
async for c in async_checkpointer.alist(config):
@@ -6017,7 +5969,10 @@ async def test_debug_retry(async_checkpointer: BaseCheckpointSaver):
update_config = await graph.aupdate_state(target_config, values=None)
events = [
c async for c in graph.astream(None, config=update_config, stream_mode="debug")
c
async for c in graph.astream(
None, config=update_config, stream_mode="debug", checkpoint_during=True
)
]
checkpoint_events = list(
@@ -6048,7 +6003,9 @@ async def test_debug_retry(async_checkpointer: BaseCheckpointSaver):
assert stream_parent_conf == history_parent_conf
async def test_debug_subgraphs(async_checkpointer: BaseCheckpointSaver):
async def test_debug_subgraphs(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
):
class State(TypedDict):
messages: Annotated[list[str], operator.add]
@@ -6082,12 +6039,15 @@ async def test_debug_subgraphs(async_checkpointer: BaseCheckpointSaver):
{"messages": []},
config=config,
stream_mode="debug",
checkpoint_during=checkpoint_during,
)
]
checkpoint_events = list(
reversed([e["payload"] for e in events if e["type"] == "checkpoint"])
)
if not checkpoint_during:
checkpoint_events = checkpoint_events[:1]
checkpoint_history = [c async for c in graph.aget_state_history(config)]
assert len(checkpoint_events) == len(checkpoint_history)
@@ -6114,7 +6074,9 @@ async def test_debug_subgraphs(async_checkpointer: BaseCheckpointSaver):
assert stream_task.get("state") == history_task.state
async def test_debug_nested_subgraphs(async_checkpointer: BaseCheckpointSaver):
async def test_debug_nested_subgraphs(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
) -> None:
from collections import defaultdict
class State(TypedDict):
@@ -6158,6 +6120,7 @@ async def test_debug_nested_subgraphs(async_checkpointer: BaseCheckpointSaver):
config=config,
stream_mode="debug",
subgraphs=True,
checkpoint_during=checkpoint_during,
)
]
@@ -6202,6 +6165,9 @@ async def test_debug_nested_subgraphs(async_checkpointer: BaseCheckpointSaver):
for checkpoint_events, checkpoint_history in zip(
stream_ns.values(), history_ns.values()
):
if not checkpoint_during:
checkpoint_events = checkpoint_events[-1:]
assert len(checkpoint_events) == len(checkpoint_history)
for stream, history in zip(checkpoint_events, checkpoint_history):
assert stream["values"] == history.values
assert stream["next"] == list(history.next)
@@ -6282,15 +6248,7 @@ async def test_parent_command(async_checkpointer: BaseCheckpointSaver) -> None:
"parents": {},
},
created_at=AnyStr(),
parent_config=(
{
"configurable": {
"thread_id": "1",
"checkpoint_ns": "",
"checkpoint_id": AnyStr(),
}
}
),
parent_config=None,
tasks=(),
interrupts=(),
)
@@ -6772,7 +6730,7 @@ async def test_concurrent_execution():
async def test_checkpoint_recovery_async(
async_checkpointer: BaseCheckpointSaver,
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
) -> None:
"""Test recovery from checkpoints after failures with async nodes."""
@@ -6802,7 +6760,11 @@ async def test_checkpoint_recovery_async(
# First attempt should fail
with pytest.raises(RuntimeError):
await graph.ainvoke({"steps": ["start"], "attempt": 1}, config)
await graph.ainvoke(
{"steps": ["start"], "attempt": 1},
config,
checkpoint_during=checkpoint_during,
)
# Verify checkpoint state
state = await graph.aget_state(config)
@@ -6811,12 +6773,17 @@ async def test_checkpoint_recovery_async(
assert state.next == ("node1",) # Should retry failed node
# Retry with updated attempt count
result = await graph.ainvoke({"steps": [], "attempt": 2}, config)
result = await graph.ainvoke(
{"steps": [], "attempt": 2}, config, checkpoint_during=checkpoint_during
)
assert result == {"steps": ["start", "node1", "node2"], "attempt": 2}
# Verify checkpoint history shows both attempts
history = [c async for c in graph.aget_state_history(config)]
assert len(history) == 6 # Initial + failed attempt + successful attempt
if checkpoint_during:
assert len(history) == 6 # Initial + failed attempt + successful attempt
else:
assert len(history) == 2 # error + success
# Verify the error was recorded in checkpoint
failed_checkpoint = next(c for c in history if c.tasks and c.tasks[0].error)
@@ -8291,7 +8258,9 @@ async def test_bulk_state_updates(async_checkpointer: BaseCheckpointSaver) -> No
)
async def test_update_as_input(async_checkpointer: BaseCheckpointSaver) -> None:
async def test_update_as_input(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
) -> None:
class State(TypedDict):
foo: str
@@ -8311,11 +8280,15 @@ async def test_update_as_input(async_checkpointer: BaseCheckpointSaver) -> None:
)
assert await graph.ainvoke(
{"foo": "input"}, {"configurable": {"thread_id": "1"}}
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
) == {"foo": "tool"}
assert await graph.ainvoke(
{"foo": "input"}, {"configurable": {"thread_id": "1"}}
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
) == {"foo": "tool"}
def map_snapshot(i: StateSnapshot) -> dict:
@@ -8354,10 +8327,15 @@ async def test_update_as_input(async_checkpointer: BaseCheckpointSaver) -> None:
async for s in graph.aget_state_history({"configurable": {"thread_id": "2"}})
]
assert new_history == history
if checkpoint_during:
assert new_history == history
else:
assert [new_history[0], new_history[4]] == history
async def test_batch_update_as_input(async_checkpointer: BaseCheckpointSaver) -> None:
async def test_batch_update_as_input(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
) -> None:
class State(TypedDict):
foo: str
tasks: Annotated[list[int], operator.add]
@@ -8389,7 +8367,9 @@ async def test_batch_update_as_input(async_checkpointer: BaseCheckpointSaver) ->
)
assert await graph.ainvoke(
{"foo": "input"}, {"configurable": {"thread_id": "1"}}
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
) == {"foo": "map", "tasks": [0, 1, 2]}
def map_snapshot(i: StateSnapshot) -> dict:
@@ -8440,7 +8420,10 @@ async def test_batch_update_as_input(async_checkpointer: BaseCheckpointSaver) ->
async for s in graph.aget_state_history({"configurable": {"thread_id": "2"}})
]
assert new_history == history
if checkpoint_during:
assert new_history == history
else:
assert new_history[:1] == history
async def test_draw_invalid():