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

feat: Implement durability mode argument (#5432)

Browse Source 
- Replaces checkpoint_during: bool
- checkpoint_during is deprecated but still respected
- We implement three durability modes (from least to most durable):
- "exit" - save checkpoint only when the graph exits (equivalent to
checkpoint_during=False)
- "async" - save checkpoint asynchronously while the next step executes
(the default, equivalent to old checkpoint_during=True)
- "sync" - save checkpoint synchronously before the next step starts
(new mode, slower but most durable)

Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
This commit is contained in:
Nuno Campos
2025-07-20 15:42:18 +01:00
committed by GitHub
co-authored by Sydney Runkle
parent a5fe3316b6
commit cb7b924006
13 changed files with 369 additions and 384 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/checkpoint-postgres/langgraph/checkpoint/postgres/shallow.py
+2 -2
View File
@@ -191,7 +191,7 @@ class ShallowPostgresSaver(BasePostgresSaver):
) -> None:
warnings.warn(
"ShallowPostgresSaver is deprecated as of version 2.0.20 and will be removed in 3.0.0. "
"Use PostgresSaver instead, and invoke the graph with `graph.invoke(..., checkpoint_during=False)`.",
"Use PostgresSaver instead, and invoke the graph with `graph.invoke(..., durability='exit')`.",
DeprecationWarning,
stacklevel=2,
)
@@ -547,7 +547,7 @@ class AsyncShallowPostgresSaver(BasePostgresSaver):
) -> None:
warnings.warn(
"AsyncShallowPostgresSaver is deprecated as of version 2.0.20 and will be removed in 3.0.0. "
"Use AsyncPostgresSaver instead, and invoke the graph with `await graph.ainvoke(..., checkpoint_during=False)`.",
"Use AsyncPostgresSaver instead, and invoke the graph with `await graph.ainvoke(..., durability='exit')`.",
DeprecationWarning,
stacklevel=2,
)
libs/langgraph/bench/__main__.py
+4 -4
View File
@@ -26,7 +26,7 @@ async def arun(graph: Pregel, input: dict):
"configurable": {"thread_id": str(uuid4())},
"recursion_limit": 1000000000,
},
checkpoint_during=False,
durability="exit",
)
]
)
@@ -43,7 +43,7 @@ async def arun_first_event_latency(graph: Pregel, input: dict) -> None:
"configurable": {"thread_id": str(uuid4())},
"recursion_limit": 1000000000,
},
checkpoint_during=False,
durability="exit",
)
try:
@@ -63,7 +63,7 @@ def run(graph: Pregel, input: dict):
"configurable": {"thread_id": str(uuid4())},
"recursion_limit": 1000000000,
},
checkpoint_during=False,
durability="exit",
)
]
)
@@ -80,7 +80,7 @@ def run_first_event_latency(graph: Pregel, input: dict) -> None:
"configurable": {"thread_id": str(uuid4())},
"recursion_limit": 1000000000,
},
checkpoint_during=False,
durability="exit",
)
try:
libs/langgraph/langgraph/_internal/_constants.py
+2 -2
View File
@@ -57,8 +57,8 @@ CONFIG_KEY_SCRATCHPAD = sys.intern("__pregel_scratchpad")
# holds a mutable dict for temporary storage scoped to the current task
CONFIG_KEY_RUNNER_SUBMIT = sys.intern("__pregel_runner_submit")
# holds a function that receives tasks from runner, executes them and returns results
CONFIG_KEY_CHECKPOINT_DURING = sys.intern("__pregel_checkpoint_during")
# holds a boolean indicating whether to checkpoint during the run (or only at the end)
CONFIG_KEY_DURABILITY = sys.intern("__pregel_durability")
# holds the durability mode, one of "sync", "async", or "exit"
CONFIG_KEY_RUNTIME = sys.intern("__pregel_runtime")
# holds a `Runtime` instance with context, store, stream writer, etc.
CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map")
libs/langgraph/langgraph/pregel/_loop.py
+11 -10
View File
@@ -113,6 +113,7 @@ from langgraph.types import (
All,
CachePolicy,
Command,
Durability,
PregelExecutableTask,
RetryPolicy,
StreamMode,
@@ -154,7 +155,7 @@ class PregelLoop:
manager: None | AsyncParentRunManager | ParentRunManager
interrupt_after: All | Sequence[str]
interrupt_before: All | Sequence[str]
checkpoint_during: bool
durability: Durability
retry_policy: Sequence[RetryPolicy]
cache_policy: CachePolicy | None
@@ -216,13 +217,13 @@ class PregelLoop:
output_keys: str | Sequence[str],
stream_keys: str | Sequence[str],
trigger_to_nodes: Mapping[str, Sequence[str]],
durability: Durability,
interrupt_after: All | Sequence[str] = EMPTY_SEQ,
interrupt_before: All | Sequence[str] = EMPTY_SEQ,
manager: None | AsyncParentRunManager | ParentRunManager = None,
migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
retry_policy: Sequence[RetryPolicy] = (),
cache_policy: CachePolicy | None = None,
checkpoint_during: bool = True,
) -> None:
self.stream = stream
self.config = config
@@ -246,7 +247,7 @@ class PregelLoop:
self.trigger_to_nodes = trigger_to_nodes
self.retry_policy = retry_policy
self.cache_policy = cache_policy
self.checkpoint_during = checkpoint_during
self.durability = durability
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)
@@ -323,7 +324,7 @@ class PregelLoop:
writes_to_save = writes
# save writes
self.checkpoint_pending_writes.extend((task_id, c, v) for c, v in writes)
if self.checkpoint_during and self.checkpointer_put_writes is not None:
if self.durability != "exit" and self.checkpointer_put_writes is not None:
config = patch_configurable(
self.checkpoint_config,
{
@@ -684,7 +685,7 @@ class PregelLoop:
self.checkpoint_metadata = metadata
# do checkpoint?
do_checkpoint = self._checkpointer_put_after_previous is not None and (
exiting or self.checkpoint_during
exiting or self.durability != "exit"
)
# create new checkpoint
self.checkpoint = create_checkpoint(
@@ -746,7 +747,7 @@ class PregelLoop:
traceback: TracebackType | None,
) -> bool | None:
# persist current checkpoint and writes
if not self.checkpoint_during and (
if self.durability == "exit" and (
# if it's a top graph
not self.is_nested
# or a nested graph with error or interrupt
@@ -891,6 +892,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
nodes: Mapping[str, PregelNode],
specs: Mapping[str, BaseChannel | ManagedValueSpec],
trigger_to_nodes: Mapping[str, Sequence[str]],
durability: Durability,
manager: None | AsyncParentRunManager | ParentRunManager = None,
interrupt_after: All | Sequence[str] = EMPTY_SEQ,
interrupt_before: All | Sequence[str] = EMPTY_SEQ,
@@ -900,7 +902,6 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
retry_policy: Sequence[RetryPolicy] = (),
cache_policy: CachePolicy | None = None,
checkpoint_during: bool = True,
) -> None:
super().__init__(
input,
@@ -921,7 +922,7 @@ class SyncPregelLoop(PregelLoop, AbstractContextManager):
trigger_to_nodes=trigger_to_nodes,
retry_policy=retry_policy,
cache_policy=cache_policy,
checkpoint_during=checkpoint_during,
durability=durability,
)
self.stack = ExitStack()
if checkpointer:
@@ -1062,6 +1063,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
nodes: Mapping[str, PregelNode],
specs: Mapping[str, BaseChannel | ManagedValueSpec],
trigger_to_nodes: Mapping[str, Sequence[str]],
durability: Durability,
interrupt_after: All | Sequence[str] = EMPTY_SEQ,
interrupt_before: All | Sequence[str] = EMPTY_SEQ,
manager: None | AsyncParentRunManager | ParentRunManager = None,
@@ -1071,7 +1073,6 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
migrate_checkpoint: Callable[[Checkpoint], None] | None = None,
retry_policy: Sequence[RetryPolicy] = (),
cache_policy: CachePolicy | None = None,
checkpoint_during: bool = True,
) -> None:
super().__init__(
input,
@@ -1092,7 +1093,7 @@ class AsyncPregelLoop(PregelLoop, AbstractAsyncContextManager):
trigger_to_nodes=trigger_to_nodes,
retry_policy=retry_policy,
cache_policy=cache_policy,
checkpoint_during=checkpoint_during,
durability=durability,
)
self.stack = AsyncExitStack()
if checkpointer:
libs/langgraph/langgraph/pregel/main.py
+80 -20
View File
@@ -11,7 +11,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
from dataclasses import is_dataclass
from functools import partial
from inspect import isclass
from typing import Any, Callable, Generic, Union, cast, get_type_hints
from typing import Any, Callable, Generic, Optional, Union, cast, get_type_hints
from uuid import UUID, uuid5
from langchain_core.globals import get_debug
@@ -40,10 +40,10 @@ from langgraph._internal._constants import (
CACHE_NS_WRITES,
CONF,
CONFIG_KEY_CACHE,
CONFIG_KEY_CHECKPOINT_DURING,
CONFIG_KEY_CHECKPOINT_ID,
CONFIG_KEY_CHECKPOINT_NS,
CONFIG_KEY_CHECKPOINTER,
CONFIG_KEY_DURABILITY,
CONFIG_KEY_NODE_FINISHED,
CONFIG_KEY_READ,
CONFIG_KEY_RUNNER_SUBMIT,
@@ -123,6 +123,7 @@ from langgraph.types import (
CachePolicy,
Checkpointer,
Command,
Durability,
Interrupt,
Send,
StateSnapshot,
@@ -2345,6 +2346,8 @@ class Pregel(
output_keys: str | Sequence[str] | None,
interrupt_before: All | Sequence[str] | None,
interrupt_after: All | Sequence[str] | None,
durability: Durability | None = None,
checkpoint_during: bool | None = None,
) -> tuple[
set[StreamMode],
str | Sequence[str],
@@ -2353,6 +2356,7 @@ class Pregel(
BaseCheckpointSaver | None,
BaseStore | None,
BaseCache | None,
Durability,
]:
if config["recursion_limit"] < 1:
raise ValueError("recursion_limit must be at least 1")
@@ -2391,6 +2395,17 @@ class Pregel(
cache: BaseCache | None = config[CONF][CONFIG_KEY_CACHE]
else:
cache = self.cache
if checkpoint_during is not None:
if durability is not None:
raise ValueError(
"Cannot use both `checkpoint_during` and `durability` parameters."
)
elif checkpoint_during:
durability = "async"
else:
durability = "exit"
if durability is None:
durability = config.get(CONF, {}).get(CONFIG_KEY_DURABILITY, "async")
return (
stream_modes,
output_keys,
@@ -2399,6 +2414,7 @@ class Pregel(
checkpointer,
store,
cache,
durability,
)
def stream(
@@ -2412,9 +2428,10 @@ class Pregel(
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,
durability: Durability | None = None,
subgraphs: bool = False,
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> Iterator[dict[str, Any] | Any]:
"""Stream graph steps for a single input.
@@ -2442,7 +2459,10 @@ class Pregel(
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 False. If False, only the final checkpoint is saved.
durability: The durability mode for the graph execution, defaults to "async". Options are:
- `"sync"`: Changes are persisted synchronously before the next step starts.
- `"async"`: Changes are persisted asynchronously while the next step executes.
- `"exit"`: Changes are persisted only when the graph exits.
subgraphs: Whether to stream events from inside subgraphs, defaults to False.
If True, the events will be emitted as tuples `(namespace, data)`,
or `(namespace, mode, data)` if `stream_mode` is a list,
@@ -2477,6 +2497,14 @@ class Pregel(
run_id=config.get("run_id"),
)
try:
deprecated_checkpoint_during = cast(
Optional[bool], kwargs.get("checkpoint_during")
)
if deprecated_checkpoint_during is not None:
warnings.warn(
"`checkpoint_during` is deprecated and will be removed. Please use `durability` instead.",
category=LangGraphDeprecatedSinceV10,
)
# assign defaults
(
stream_modes,
@@ -2486,6 +2514,7 @@ class Pregel(
checkpointer,
store,
cache,
durability_,
) = self._defaults(
config,
stream_mode=stream_mode,
@@ -2493,7 +2522,15 @@ class Pregel(
output_keys=output_keys,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
checkpoint_during=deprecated_checkpoint_during,
)
if checkpointer is None and (
durability is not None or deprecated_checkpoint_during is not None
):
warnings.warn(
"`durability` has no effect when no checkpointer is present.",
)
# set up subgraph checkpointing
if self.checkpointer is True:
ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS])
@@ -2526,9 +2563,9 @@ class Pregel(
def stream_writer(c: Any) -> None:
pass
# set checkpointing mode for subgraphs
if checkpoint_during is not None:
config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during
# set durability mode for subgraphs
if durability is not None or deprecated_checkpoint_during is not None:
config[CONF][CONFIG_KEY_DURABILITY] = durability_
config[CONF][CONFIG_KEY_RUNTIME] = Runtime(
context=context,
@@ -2551,9 +2588,7 @@ class Pregel(
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
checkpoint_during=checkpoint_during
if checkpoint_during is not None
else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True),
durability=durability_,
trigger_to_nodes=self.trigger_to_nodes,
migrate_checkpoint=self._migrate_checkpoint,
retry_policy=self.retry_policy,
@@ -2614,6 +2649,9 @@ class Pregel(
stream_mode, print_mode, subgraphs, stream.get, queue.Empty
)
loop.after_tick()
# wait for checkpoint
if durability_ == "sync":
loop._put_checkpoint_fut.result()
# emit output
yield from _output(
stream_mode, print_mode, subgraphs, stream.get, queue.Empty
@@ -2646,9 +2684,10 @@ class Pregel(
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,
durability: Durability | None = None,
subgraphs: bool = False,
debug: bool | None = None,
**kwargs: Unpack[DeprecatedKwargs],
) -> AsyncIterator[dict[str, Any] | Any]:
"""Asynchronously stream graph steps for a single input.
@@ -2675,7 +2714,10 @@ class Pregel(
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 False. If False, only the final checkpoint is saved.
durability: The durability mode for the graph execution, defaults to "async". Options are:
- `"sync"`: Changes are persisted synchronously before the next step starts.
- `"async"`: Changes are persisted asynchronously while the next step executes.
- `"exit"`: Changes are persisted only when the graph exits.
subgraphs: Whether to stream events from inside subgraphs, defaults to False.
If True, the events will be emitted as tuples `(namespace, data)`,
or `(namespace, mode, data)` if `stream_mode` is a list,
@@ -2729,6 +2771,14 @@ class Pregel(
else False
)
try:
deprecated_checkpoint_during = cast(
Optional[bool], kwargs.get("checkpoint_during")
)
if deprecated_checkpoint_during is not None:
warnings.warn(
"`checkpoint_during` is deprecated and will be removed. Please use `durability` instead.",
category=LangGraphDeprecatedSinceV10,
)
# assign defaults
(
stream_modes,
@@ -2738,6 +2788,7 @@ class Pregel(
checkpointer,
store,
cache,
durability_,
) = self._defaults(
config,
stream_mode=stream_mode,
@@ -2745,7 +2796,15 @@ class Pregel(
output_keys=output_keys,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
durability=durability,
checkpoint_during=deprecated_checkpoint_during,
)
if checkpointer is None and (
durability is not None or deprecated_checkpoint_during is not None
):
warnings.warn(
"`durability` has no effect when no checkpointer is present.",
)
# set up subgraph checkpointing
if self.checkpointer is True:
ns = cast(str, config[CONF][CONFIG_KEY_CHECKPOINT_NS])
@@ -2793,9 +2852,9 @@ class Pregel(
def stream_writer(c: Any) -> None:
pass
# set checkpointing mode for subgraphs
if checkpoint_during is not None:
config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during
# set durability mode for subgraphs
if durability is not None or deprecated_checkpoint_during is not None:
config[CONF][CONFIG_KEY_DURABILITY] = durability_
config[CONF][CONFIG_KEY_RUNTIME] = Runtime(
context=context,
@@ -2818,9 +2877,7 @@ class Pregel(
interrupt_before=interrupt_before_,
interrupt_after=interrupt_after_,
manager=run_manager,
checkpoint_during=checkpoint_during
if checkpoint_during is not None
else config[CONF].get(CONFIG_KEY_CHECKPOINT_DURING, True),
durability=durability_,
trigger_to_nodes=self.trigger_to_nodes,
migrate_checkpoint=self._migrate_checkpoint,
retry_policy=self.retry_policy,
@@ -2877,6 +2934,9 @@ class Pregel(
):
yield o
loop.after_tick()
# wait for checkpoint
if durability_ == "sync":
await cast(asyncio.Future, loop._put_checkpoint_fut)
# emit output
for o in _output(
stream_mode,
libs/langgraph/langgraph/types.py
+6
View File
@@ -55,9 +55,15 @@ __all__ = (
"StateSnapshot",
"Send",
"Command",
"Durability",
"interrupt",
)
Durability = Literal["sync", "async", "exit"]
"""Durability mode for the graph execution.
- `"sync"`: Changes are persisted synchronously before the next step starts.
- `"async"`: Changes are persisted asynchronously while the next step executes.
- `"exit"`: Changes are persisted only when the graph exits."""
All = Literal["*"]
"""Special value to indicate that graph should interrupt on all nodes."""
libs/langgraph/tests/conftest.py
+3 -2
View File
@@ -10,6 +10,7 @@ from langgraph.cache.memory import InMemoryCache
from langgraph.cache.sqlite import SqliteCache
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.store.base import BaseStore
from langgraph.types import Durability
from tests.conftest_checkpointer import (
_checkpointer_memory,
_checkpointer_memory_migrate_sends,
@@ -49,8 +50,8 @@ def deterministic_uuids(mocker: MockerFixture) -> MockerFixture:
return mocker.patch("uuid.uuid4", side_effect=side_effect)
@pytest.fixture(params=[True, False])
def checkpoint_during(request: pytest.FixtureRequest) -> bool:
@pytest.fixture(params=["sync", "async", "exit"])
def durability(request: pytest.FixtureRequest) -> Durability:
return request.param
libs/langgraph/tests/test_checkpoint_migration.py
+4 -4
View File
@@ -1513,7 +1513,7 @@ def test_latest_checkpoint_state_graph(
config = {"configurable": {"thread_id": "1"}}
assert [
*app.stream({"query": "what is weather in sf"}, config, checkpoint_during=True)
*app.stream({"query": "what is weather in sf"}, config, durability="async")
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
{"analyzer_one": {"query": "analyzed: query: what is weather in sf"}},
@@ -1529,7 +1529,7 @@ def test_latest_checkpoint_state_graph(
},
]
assert [*app.stream(Command(resume=""), config, checkpoint_during=True)] == [
assert [*app.stream(Command(resume=""), config, durability="async")] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
@@ -1556,7 +1556,7 @@ async def test_latest_checkpoint_state_graph_async(
assert [
c
async for c in app.astream(
{"query": "what is weather in sf"}, config, checkpoint_during=True
{"query": "what is weather in sf"}, config, durability="async"
)
] == [
{"rewrite_query": {"query": "query: what is weather in sf"}},
@@ -1574,7 +1574,7 @@ async def test_latest_checkpoint_state_graph_async(
]
assert [
c async for c in app.astream(Command(resume=""), config, checkpoint_during=True)
c async for c in app.astream(Command(resume=""), config, durability="async")
] == [
{"qa": {"answer": "doc1,doc2,doc3,doc4"}},
]
libs/langgraph/tests/test_interruption.py
+15 -14
View File
@@ -3,12 +3,13 @@ from typing_extensions import TypedDict
from langgraph.checkpoint.base import BaseCheckpointSaver
from langgraph.graph import END, START, StateGraph
from langgraph.types import Durability
pytestmark = pytest.mark.anyio
def test_interruption_without_state_updates(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
"""Test interruption without state updates. This test confirms that
interrupting doesn't require a state key having been updated in the prev step"""
@@ -33,24 +34,24 @@ def test_interruption_without_state_updates(
initial_input = {"input": "hello world"}
thread = {"configurable": {"thread_id": "1"}}
graph.invoke(initial_input, thread, checkpoint_during=checkpoint_during)
graph.invoke(initial_input, thread, durability=durability)
assert graph.get_state(thread).next == ("step_2",)
n_checkpoints = len([c for c in graph.get_state_history(thread)])
assert n_checkpoints == (3 if checkpoint_during else 1)
assert n_checkpoints == (3 if durability != "exit" else 1)
graph.invoke(None, thread, checkpoint_during=checkpoint_during)
graph.invoke(None, thread, durability=durability)
assert graph.get_state(thread).next == ("step_3",)
n_checkpoints = len([c for c in graph.get_state_history(thread)])
assert n_checkpoints == (4 if checkpoint_during else 2)
assert n_checkpoints == (4 if durability != "exit" else 2)
graph.invoke(None, thread, checkpoint_during=checkpoint_during)
graph.invoke(None, thread, durability=durability)
assert graph.get_state(thread).next == ()
n_checkpoints = len([c for c in graph.get_state_history(thread)])
assert n_checkpoints == (5 if checkpoint_during else 3)
assert n_checkpoints == (5 if durability != "exit" else 3)
async def test_interruption_without_state_updates_async(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
"""Test interruption without state updates. This test confirms that
interrupting doesn't require a state key having been updated in the prev step"""
@@ -75,17 +76,17 @@ async def test_interruption_without_state_updates_async(
initial_input = {"input": "hello world"}
thread = {"configurable": {"thread_id": "1"}}
await graph.ainvoke(initial_input, thread, checkpoint_during=checkpoint_during)
await graph.ainvoke(initial_input, thread, durability=durability)
assert (await graph.aget_state(thread)).next == ("step_2",)
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
assert n_checkpoints == (3 if checkpoint_during else 1)
assert n_checkpoints == (3 if durability != "exit" else 1)
await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during)
await graph.ainvoke(None, thread, durability=durability)
assert (await graph.aget_state(thread)).next == ("step_3",)
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
assert n_checkpoints == (4 if checkpoint_during else 2)
assert n_checkpoints == (4 if durability != "exit" else 2)
await graph.ainvoke(None, thread, checkpoint_during=checkpoint_during)
await graph.ainvoke(None, thread, durability=durability)
assert (await graph.aget_state(thread)).next == ()
n_checkpoints = len([c async for c in graph.aget_state_history(thread)])
assert n_checkpoints == (5 if checkpoint_during else 3)
assert n_checkpoints == (5 if durability != "exit" else 3)
libs/langgraph/tests/test_large_cases.py
+37 -40
View File
@@ -24,6 +24,7 @@ from langgraph.prebuilt.tool_node import ToolNode
from langgraph.pregel import NodeBuilder, Pregel
from langgraph.types import (
Command,
Durability,
Interrupt,
PregelTask,
RetryPolicy,
@@ -69,7 +70,7 @@ def test_invoke_two_processes_in_out_interrupt(
thread2 = {"configurable": {"thread_id": "2"}}
# start execution, stop at inbox
assert app.invoke(2, thread1, checkpoint_during=True) is None
assert app.invoke(2, thread1, durability="async") is None
# inbox == 3
checkpoint = sync_checkpointer.get(thread1)
@@ -77,10 +78,10 @@ def test_invoke_two_processes_in_out_interrupt(
assert checkpoint["channel_values"]["inbox"] == 3
# resume execution, finish
assert app.invoke(None, thread1, checkpoint_during=True) == 4
assert app.invoke(None, thread1, durability="async") == 4
# start execution again, stop at inbox
assert app.invoke(20, thread1, checkpoint_during=True) is None
assert app.invoke(20, thread1, durability="async") is None
# inbox == 21
checkpoint = sync_checkpointer.get(thread1)
@@ -88,11 +89,11 @@ def test_invoke_two_processes_in_out_interrupt(
assert checkpoint["channel_values"]["inbox"] == 21
# send a new value in, interrupting the previous execution
assert app.invoke(3, thread1, checkpoint_during=True) is None
assert app.invoke(None, thread1, checkpoint_during=True) == 5
assert app.invoke(3, thread1, durability="async") is None
assert app.invoke(None, thread1, durability="async") == 5
# start execution again, stopping at inbox
assert app.invoke(20, thread2, checkpoint_during=True) is None
assert app.invoke(20, thread2, durability="async") is None
# inbox == 21
snapshot = app.get_state(thread2)
@@ -299,9 +300,7 @@ def test_fork_always_re_runs_nodes(
# start execution, stop at inbox
assert [
*graph.stream(
1, thread1, stream_mode=["values", "updates"], checkpoint_during=True
)
*graph.stream(1, thread1, stream_mode=["values", "updates"], durability="async")
] == [
("values", 1),
("updates", {"add_one": 1}),
@@ -666,7 +665,7 @@ def test_conditional_state_graph(
assert [
c
for c in app_w_interrupt.stream(
{"input": "what is weather in sf"}, config, checkpoint_during=False
{"input": "what is weather in sf"}, config, durability="exit"
)
] == [
{
@@ -836,7 +835,7 @@ def test_conditional_state_graph(
assert [
c
for c in app_w_interrupt.stream(
{"input": "what is weather in sf"}, config, checkpoint_during=False
{"input": "what is weather in sf"}, config, durability="exit"
)
] == [
{
@@ -1003,7 +1002,7 @@ def test_conditional_state_graph(
assert [
c
for c in app_w_interrupt.stream(
{"input": "what is weather in sf"}, config, checkpoint_during=False
{"input": "what is weather in sf"}, config, durability="exit"
)
] == [
{"__interrupt__": ()},
@@ -1150,7 +1149,7 @@ def test_conditional_state_graph(
assert [
c
for c in app_w_interrupt.stream(
{"input": "what is weather in sf"}, config, checkpoint_during=False
{"input": "what is weather in sf"}, config, durability="exit"
)
] == [
{
@@ -1851,7 +1850,7 @@ def test_state_graph_packets(
for c in app_w_interrupt.stream(
{"messages": HumanMessage(content="what is weather in sf")},
config,
checkpoint_during=False,
durability="exit",
)
] == [
{
@@ -2116,7 +2115,7 @@ def test_state_graph_packets(
for c in app_w_interrupt.stream(
{"messages": HumanMessage(content="what is weather in sf")},
config,
checkpoint_during=False,
durability="exit",
)
] == [
{
@@ -2584,7 +2583,7 @@ def test_message_graph(
assert [
c
for c in app_w_interrupt.stream(
("human", "what is weather in sf"), config, checkpoint_during=False
("human", "what is weather in sf"), config, durability="exit"
)
] == [
{
@@ -2809,7 +2808,7 @@ def test_message_graph(
assert [
c
for c in app_w_interrupt.stream(
"what is weather in sf", config, checkpoint_during=False
"what is weather in sf", config, durability="exit"
)
] == [
{
@@ -3306,7 +3305,7 @@ def test_root_graph(
assert [
c
for c in app_w_interrupt.stream(
("human", "what is weather in sf"), config, checkpoint_during=False
("human", "what is weather in sf"), config, durability="exit"
)
] == [
{
@@ -3533,7 +3532,7 @@ def test_root_graph(
assert [
c
for c in app_w_interrupt.stream(
"what is weather in sf", config, checkpoint_during=False
"what is weather in sf", config, durability="exit"
)
] == [
{
@@ -4217,7 +4216,7 @@ def test_dynamic_interrupt(sync_checkpointer: BaseCheckpointSaver) -> None:
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert tool_two.invoke(
{"my_key": "value ⛰️", "market": "DE"}, thread1, checkpoint_during=False
{"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit"
) == {
"my_key": "value ⛰️",
"market": "DE",
@@ -4377,7 +4376,7 @@ def test_partial_pending_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> N
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert tool_two.invoke(
{"my_key": "value ⛰️", "market": "DE"}, thread1, checkpoint_during=False
{"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit"
) == {
"my_key": "value ⛰️ one",
"market": "DE",
@@ -4545,7 +4544,7 @@ def test_dynamic_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver) -> N
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert tool_two.invoke(
{"my_key": "value ⛰️", "market": "DE"}, thread1, checkpoint_during=False
{"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit"
) == {
"my_key": "value ⛰️",
"market": "DE",
@@ -4645,7 +4644,7 @@ def test_dynamic_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver) -> N
def test_send_dedupe_on_resume(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class InterruptOnce:
ticks: int = 0
@@ -4699,7 +4698,7 @@ def test_send_dedupe_on_resume(
graph = builder.compile(checkpointer=sync_checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
assert graph.invoke(["0"], thread1, checkpoint_during=checkpoint_during) == {
assert graph.invoke(["0"], thread1, durability=durability) == {
"__interrupt__": [
Interrupt(
value="Bahh",
@@ -4714,10 +4713,10 @@ def test_send_dedupe_on_resume(
assert state.next == ("flaky",)
# check history
history = [c for c in graph.get_state_history(thread1)]
assert len(history) == (4 if checkpoint_during else 1)
assert len(history) == (4 if durability != "exit" else 1)
# resume execution
assert graph.invoke(None, thread1, checkpoint_during=checkpoint_during) == [
assert graph.invoke(None, thread1, durability=durability) == [
"0",
"1",
"3.1",
@@ -4737,7 +4736,7 @@ def test_send_dedupe_on_resume(
assert state.next == ()
# check history
history = [c for c in graph.get_state_history(thread1)]
assert len(history) == (6 if checkpoint_during else 2)
assert len(history) == (6 if durability != "exit" else 2)
expected_history = [
StateSnapshot(
values=[
@@ -4866,7 +4865,7 @@ def test_send_dedupe_on_resume(
error=None,
interrupts=(Interrupt(value="Bahh", id=AnyStr()),),
state=None,
result=["flaky|4"] if checkpoint_during else None,
result=["flaky|4"] if durability != "exit" else None,
),
PregelTask(
id=AnyStr(),
@@ -5001,7 +5000,7 @@ def test_send_dedupe_on_resume(
),
),
]
if checkpoint_during:
if durability != "exit":
assert history == expected_history
else:
assert history[0] == expected_history[0]._replace(
@@ -5059,7 +5058,7 @@ def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None:
app = graph.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "1"}}
app.invoke({"my_key": "my value"}, config, checkpoint_during=False)
app.invoke({"my_key": "my value"}, config, durability="exit")
# test state w/ nested subgraph state (right after interrupt)
# first get_state without subgraph state
expected = StateSnapshot(
@@ -5183,7 +5182,7 @@ def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None:
assert child_history == expected_child_history
# resume
app.invoke(None, config, checkpoint_during=False)
app.invoke(None, config, durability="exit")
# test state w/ nested subgraph state (after resuming from interrupt)
assert app.get_state(config) == StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
@@ -5339,7 +5338,7 @@ def test_doubly_nested_graph_state(
assert [
c
for c in app.stream(
{"my_key": "my value"}, config, subgraphs=True, checkpoint_during=False
{"my_key": "my value"}, config, subgraphs=True, durability="exit"
)
] == [
((), {"parent_1": {"my_key": "hi my value"}}),
@@ -5558,9 +5557,7 @@ def test_doubly_nested_graph_state(
interrupts=(),
)
# # resume
assert [
c for c in app.stream(None, config, subgraphs=True, checkpoint_during=False)
] == [
assert [c for c in app.stream(None, config, subgraphs=True, durability="exit")] == [
(
(AnyStr("child:"), AnyStr("child_1:")),
{"grandchild_2": {"my_key": "hi my value here and there"}},
@@ -5918,7 +5915,7 @@ def test_send_react_interrupt(
graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"])
thread1 = {"configurable": {"thread_id": "2"}}
assert graph.invoke(
{"messages": [HumanMessage("hello")]}, thread1, checkpoint_during=False
{"messages": [HumanMessage("hello")]}, thread1, durability="exit"
) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
@@ -6042,7 +6039,7 @@ def test_send_react_interrupt(
graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"])
thread1 = {"configurable": {"thread_id": "3"}}
assert graph.invoke(
{"messages": [HumanMessage("hello")]}, thread1, checkpoint_during=False
{"messages": [HumanMessage("hello")]}, thread1, durability="exit"
) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
@@ -6308,7 +6305,7 @@ def test_send_react_interrupt_control(
graph = builder.compile(checkpointer=sync_checkpointer, interrupt_before=["foo"])
thread1 = {"configurable": {"thread_id": "2"}}
assert graph.invoke(
{"messages": [HumanMessage("hello")]}, thread1, checkpoint_during=False
{"messages": [HumanMessage("hello")]}, thread1, durability="exit"
) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
@@ -6567,7 +6564,7 @@ def test_weather_subgraph(
config=config,
stream_mode="updates",
subgraphs=True,
checkpoint_during=False,
durability="exit",
)
] == [
((), {"router_node": {"route": "weather"}}),
@@ -6654,7 +6651,7 @@ def test_weather_subgraph(
config=config,
stream_mode="updates",
subgraphs=True,
checkpoint_during=False,
durability="exit",
)
] == [
((), {"router_node": {"route": "weather"}}),
libs/langgraph/tests/test_large_cases_async.py
+18 -21
View File
@@ -63,7 +63,7 @@ async def test_invoke_two_processes_in_out_interrupt(
thread2 = {"configurable": {"thread_id": "2"}}
# start execution, stop at inbox
assert await app.ainvoke(2, thread1, checkpoint_during=True) is None
assert await app.ainvoke(2, thread1, durability="async") is None
# inbox == 3
checkpoint = await async_checkpointer.aget(thread1)
@@ -71,10 +71,10 @@ async def test_invoke_two_processes_in_out_interrupt(
assert checkpoint["channel_values"]["inbox"] == 3
# resume execution, finish
assert await app.ainvoke(None, thread1, checkpoint_during=True) == 4
assert await app.ainvoke(None, thread1, durability="async") == 4
# start execution again, stop at inbox
assert await app.ainvoke(20, thread1, checkpoint_during=True) is None
assert await app.ainvoke(20, thread1, durability="async") is None
# inbox == 21
checkpoint = await async_checkpointer.aget(thread1)
@@ -82,11 +82,11 @@ async def test_invoke_two_processes_in_out_interrupt(
assert checkpoint["channel_values"]["inbox"] == 21
# send a new value in, interrupting the previous execution
assert await app.ainvoke(3, thread1, checkpoint_during=True) is None
assert await app.ainvoke(None, thread1, checkpoint_during=True) == 5
assert await app.ainvoke(3, thread1, durability="async") is None
assert await app.ainvoke(None, thread1, durability="async") == 5
# start execution again, stopping at inbox
assert await app.ainvoke(20, thread2, checkpoint_during=True) is None
assert await app.ainvoke(20, thread2, durability="async") is None
# inbox == 21
snapshot = await app.aget_state(thread2)
@@ -301,7 +301,7 @@ async def test_fork_always_re_runs_nodes(
assert [
c
async for c in graph.astream(
1, thread1, stream_mode=["values", "updates"], checkpoint_during=True
1, thread1, stream_mode=["values", "updates"], durability="async"
)
] == [
("values", 1),
@@ -684,7 +684,7 @@ async def test_conditional_graph_state(async_checkpointer: BaseCheckpointSaver)
assert [
c
async for c in app_w_interrupt.astream(
{"input": "what is weather in sf"}, config, checkpoint_during=False
{"input": "what is weather in sf"}, config, durability="exit"
)
] == [
{
@@ -859,7 +859,7 @@ async def test_conditional_graph_state(async_checkpointer: BaseCheckpointSaver)
assert [
c
async for c in app_w_interrupt.astream(
{"input": "what is weather in sf"}, config, checkpoint_during=False
{"input": "what is weather in sf"}, config, durability="exit"
)
] == [
{
@@ -1577,7 +1577,7 @@ async def test_state_graph_packets(async_checkpointer: BaseCheckpointSaver) -> N
async for c in app_w_interrupt.astream(
{"messages": HumanMessage(content="what is weather in sf")},
config,
checkpoint_during=False,
durability="exit",
)
] == [
{
@@ -1828,7 +1828,7 @@ async def test_state_graph_packets(async_checkpointer: BaseCheckpointSaver) -> N
async for c in app_w_interrupt.astream(
{"messages": HumanMessage(content="what is weather in sf")},
config,
checkpoint_during=False,
durability="exit",
)
] == [
{
@@ -2257,7 +2257,7 @@ async def test_message_graph(async_checkpointer: BaseCheckpointSaver) -> None:
async for c in app_w_interrupt.astream(
HumanMessage(content="what is weather in sf"),
config,
checkpoint_during=False,
durability="exit",
)
] == [
{
@@ -2740,7 +2740,7 @@ async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> No
app = graph.compile(checkpointer=async_checkpointer)
config = {"configurable": {"thread_id": "1"}}
await app.ainvoke({"my_key": "my value"}, config, checkpoint_during=False)
await app.ainvoke({"my_key": "my value"}, config, durability="exit")
# test state w/ nested subgraph state (right after interrupt)
# first get_state without subgraph state
expected = StateSnapshot(
@@ -2871,7 +2871,7 @@ async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> No
assert child_history == expected_child_history
# resume
await app.ainvoke(None, config, checkpoint_during=False)
await app.ainvoke(None, config, durability="exit")
# test state w/ nested subgraph state (after resuming from interrupt)
assert await app.aget_state(config) == StateSnapshot(
values={"my_key": "hi my value here and there and back again"},
@@ -3030,7 +3030,7 @@ async def test_doubly_nested_graph_state(
assert [
c
async for c in app.astream(
{"my_key": "my value"}, config, subgraphs=True, checkpoint_during=False
{"my_key": "my value"}, config, subgraphs=True, durability="exit"
)
] == [
((), {"parent_1": {"my_key": "hi my value"}}),
@@ -3250,10 +3250,7 @@ async def test_doubly_nested_graph_state(
)
# resume
assert [
c
async for c in app.astream(
None, config, subgraphs=True, checkpoint_during=False
)
c async for c in app.astream(None, config, subgraphs=True, durability="exit")
] == [
(
(AnyStr("child:"), AnyStr("child_1:")),
@@ -3662,7 +3659,7 @@ async def test_weather_subgraph(
config=config,
stream_mode="updates",
subgraphs=True,
checkpoint_during=False,
durability="exit",
)
] == [
((), {"router_node": {"route": "weather"}}),
@@ -3751,7 +3748,7 @@ async def test_weather_subgraph(
config=config,
stream_mode="updates",
subgraphs=True,
checkpoint_during=False,
durability="exit",
)
] == [
((), {"router_node": {"route": "weather"}}),
libs/langgraph/tests/test_pregel.py
+84 -129
View File
@@ -57,6 +57,7 @@ from langgraph.store.base import BaseStore
from langgraph.types import (
CachePolicy,
Command,
Durability,
Interrupt,
PregelTask,
RetryPolicy,
@@ -185,7 +186,7 @@ def test_checkpoint_errors() -> None:
graph = builder.compile(checkpointer=FaultyPutWritesCheckpointer())
with pytest.raises(ValueError, match="Faulty put_writes"):
graph.invoke(
"", {"configurable": {"thread_id": "thread-1"}}, checkpoint_during=True
"", {"configurable": {"thread_id": "thread-1"}}, durability="async"
)
@@ -570,7 +571,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, checkpoint_during=True)
result = graph.invoke({"myval": 1}, thread1, durability="async")
assert result["myval"] == 4
history = [c for c in graph.get_state_history(thread1)]
@@ -827,7 +828,7 @@ def test_invoke_checkpoint_two(
def test_pending_writes_resume(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class State(TypedDict):
value: Annotated[int, operator.add]
@@ -864,7 +865,7 @@ def test_pending_writes_resume(
thread1: RunnableConfig = {"configurable": {"thread_id": "1"}}
with pytest.raises(ConnectionError, match="I'm not good"):
graph.invoke({"value": 1}, thread1, checkpoint_during=checkpoint_during)
graph.invoke({"value": 1}, thread1, durability=durability)
# both nodes should have been called once
assert one.calls == 1
@@ -908,7 +909,7 @@ def test_pending_writes_resume(
# resume execution
with pytest.raises(ConnectionError, match="I'm not good"):
graph.invoke(None, thread1, checkpoint_during=checkpoint_during)
graph.invoke(None, thread1, durability=durability)
# node "one" succeeded previously, so shouldn't be called again
assert one.calls == 1
@@ -922,14 +923,12 @@ def test_pending_writes_resume(
# resume execution, without exception
two.rtn = {"value": 3}
# both the pending write and the new write were applied, 1 + 2 + 3 = 6
assert graph.invoke(None, thread1, checkpoint_during=checkpoint_during) == {
"value": 6
}
assert graph.invoke(None, thread1, durability=durability) == {"value": 6}
# check all final checkpoints
checkpoints = [c for c in sync_checkpointer.list(thread1)]
# we should have 3
assert len(checkpoints) == (3 if checkpoint_during else 2)
assert len(checkpoints) == (3 if durability != "exit" else 2)
# the last one not too interesting for this test
assert checkpoints[0] == CheckpointTuple(
config={
@@ -1030,7 +1029,7 @@ def test_pending_writes_resume(
),
}
}
if checkpoint_during
if durability != "exit"
else None,
pending_writes=(
UnsortedSequence(
@@ -1038,7 +1037,7 @@ def test_pending_writes_resume(
(AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
(AnyStr(), "value", 3),
)
if checkpoint_during
if durability != "exit"
else UnsortedSequence(
(AnyStr(), "value", 2),
(AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
@@ -1047,7 +1046,7 @@ def test_pending_writes_resume(
)
),
)
if not checkpoint_during:
if durability == "exit":
return
assert checkpoints[2] == CheckpointTuple(
config={
@@ -1204,7 +1203,7 @@ def test_send_sequences() -> None:
def test_imp_task(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
mapper_calls = 0
@@ -1243,7 +1242,7 @@ def test_imp_task(
}
thread1 = {"configurable": {"thread_id": "1"}}
assert [*graph.stream([0, 1], thread1, checkpoint_during=checkpoint_during)] == [
assert [*graph.stream([0, 1], thread1, durability=durability)] == [
{"mapper": "00"},
{"mapper": "11"},
{
@@ -1257,9 +1256,7 @@ def test_imp_task(
]
assert mapper_calls == 2
assert graph.invoke(
Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
) == [
assert graph.invoke(Command(resume="answer"), thread1, durability=durability) == [
"00answer",
"11answer",
]
@@ -1267,7 +1264,7 @@ def test_imp_task(
def test_imp_nested(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
def mynode(input: list[str]) -> list[str]:
return [it + "a" for it in input]
@@ -1308,7 +1305,7 @@ def test_imp_nested(
}
thread1 = {"configurable": {"thread_id": "1"}}
assert [*graph.stream([0, 1], thread1, checkpoint_during=checkpoint_during)] == [
assert [*graph.stream([0, 1], thread1, durability=durability)] == [
{"submapper": "0"},
{"mapper": "00"},
{"submapper": "1"},
@@ -1323,16 +1320,14 @@ def test_imp_nested(
},
]
assert graph.invoke(
Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
) == [
assert graph.invoke(Command(resume="answer"), thread1, durability=durability) == [
"00answera",
"11answera",
]
def test_imp_stream_order(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
@task()
def foo(state: dict) -> tuple:
@@ -1354,10 +1349,7 @@ def test_imp_stream_order(
return fut_baz.result()
thread1 = {"configurable": {"thread_id": "1"}}
assert [
c
for c in graph.stream({"a": "0"}, thread1, checkpoint_during=checkpoint_during)
] == [
assert [c for c in graph.stream({"a": "0"}, thread1, durability=durability)] == [
{
"foo": (
"0foo",
@@ -1405,7 +1397,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, checkpoint_during=True) == 2
assert app.invoke(2, thread_1, durability="async") == 2
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 2
@@ -1415,7 +1407,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, checkpoint_during=True) == 5
assert app.invoke(3, thread_1, durability="async") == 5
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
@@ -1425,7 +1417,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, checkpoint_during=True)
app.invoke(4, thread_1, durability="async")
# checkpoint is updated with new input
state = app.get_state(thread_1)
assert state is not None
@@ -1433,7 +1425,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, checkpoint_during=True) == 9
assert app.invoke(2, thread_1, durability="async") == 9
state = app.get_state(thread_1)
assert state is not None
assert state.values.get("total") == 16, "total is now 7+9=16"
@@ -3176,7 +3168,7 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None:
def test_subgraph_checkpoint_true(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class InnerState(TypedDict):
my_key: Annotated[str, operator.add]
@@ -3210,7 +3202,7 @@ def test_subgraph_checkpoint_true(
assert [
c
for c in app.stream(
{"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during
{"my_key": ""}, config, subgraphs=True, durability=durability
)
] == [
(("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}),
@@ -3237,13 +3229,13 @@ def test_subgraph_checkpoint_true(
]
checkpoints = list(app.get_state_history(config))
if checkpoint_during:
if durability != "exit":
assert len(checkpoints) == 4
else:
assert len(checkpoints) == 1
def test_subgraph_checkpoint_during_false_inherited() -> None:
def test_subgraph_durability_inherited(durability: Durability) -> None:
sync_checkpointer = InMemorySaver()
class InnerState(TypedDict):
@@ -3274,22 +3266,19 @@ def test_subgraph_checkpoint_during_false_inherited() -> None:
"inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END
)
app = graph.compile(checkpointer=sync_checkpointer)
for checkpoint_during in [True, False]:
thread_id = str(uuid.uuid4())
config = {"configurable": {"thread_id": thread_id}}
app.invoke(
{"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during
)
if checkpoint_during:
checkpoints = list(sync_checkpointer.list(config))
assert len(checkpoints) == 12
else:
checkpoints = list(sync_checkpointer.list(config))
assert len(checkpoints) == 1
thread_id = str(uuid.uuid4())
config = {"configurable": {"thread_id": thread_id}}
app.invoke({"my_key": ""}, config, subgraphs=True, durability=durability)
if durability != "exit":
checkpoints = list(sync_checkpointer.list(config))
assert len(checkpoints) == 12
else:
checkpoints = list(sync_checkpointer.list(config))
assert len(checkpoints) == 1
def test_subgraph_checkpoint_true_interrupt(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
# Define subgraph
class SubgraphState(TypedDict):
@@ -3330,9 +3319,7 @@ def test_subgraph_checkpoint_true_interrupt(
graph = builder.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "1"}}
assert graph.invoke(
{"foo": "foo"}, config, checkpoint_during=checkpoint_during
) == {
assert graph.invoke({"foo": "foo"}, config, durability=durability) == {
"foo": "hi! foo",
"__interrupt__": [
Interrupt(
@@ -3344,9 +3331,9 @@ def test_subgraph_checkpoint_true_interrupt(
assert graph.get_state(config, subgraphs=True).tasks[0].state.values == {
"bar": "hi! foo"
}
assert graph.invoke(
Command(resume="baz"), config, checkpoint_during=checkpoint_during
) == {"foo": "hi! foobaz"}
assert graph.invoke(Command(resume="baz"), config, durability=durability) == {
"foo": "hi! foobaz"
}
def test_stream_subgraphs_during_execution(
@@ -3455,7 +3442,7 @@ def test_stream_buffering_single_node(sync_checkpointer: BaseCheckpointSaver) ->
def test_nested_graph_interrupts_parallel(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class InnerState(TypedDict):
my_key: Annotated[str, operator.add]
@@ -3501,11 +3488,11 @@ def test_nested_graph_interrupts_parallel(
# test invoke w/ nested interrupt
config = {"configurable": {"thread_id": "1"}}
assert app.invoke({"my_key": ""}, config, checkpoint_during=checkpoint_during) == {
assert app.invoke({"my_key": ""}, config, durability=durability) == {
"my_key": " and parallel",
}
assert app.invoke(None, config, checkpoint_during=checkpoint_during) == {
assert app.invoke(None, config, durability=durability) == {
"my_key": "got here and there and parallel and back again",
}
@@ -3515,16 +3502,14 @@ def test_nested_graph_interrupts_parallel(
# test stream updates w/ nested interrupt
config = {"configurable": {"thread_id": "2"}}
assert [
*app.stream(
{"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during
)
*app.stream({"my_key": ""}, config, subgraphs=True, durability=durability)
] == [
# we got to parallel node first
((), {"outer_1": {"my_key": " and parallel"}}),
((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}),
((), {"__interrupt__": ()}),
]
assert [*app.stream(None, config, checkpoint_during=checkpoint_during)] == [
assert [*app.stream(None, config, durability=durability)] == [
{"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}},
{"inner": {"my_key": "got here and there"}},
{"outer_2": {"my_key": " and back again"}},
@@ -3537,17 +3522,13 @@ def test_nested_graph_interrupts_parallel(
{"my_key": ""},
config,
stream_mode="values",
checkpoint_during=checkpoint_during,
durability=durability,
)
] == [
{"my_key": ""},
{"my_key": " and parallel"},
]
assert [
*app.stream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
)
] == [
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
{"my_key": ""},
{"my_key": "got here and there and parallel"},
{"my_key": "got here and there and parallel and back again"},
@@ -3561,23 +3542,15 @@ def test_nested_graph_interrupts_parallel(
{"my_key": ""},
config,
stream_mode="values",
checkpoint_during=checkpoint_during,
durability=durability,
)
] == [{"my_key": ""}]
# while we're waiting for the node w/ interrupt inside to finish
assert [
*app.stream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
)
] == [
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
{"my_key": ""},
{"my_key": " and parallel"},
]
assert [
*app.stream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
)
] == [
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
{"my_key": ""},
{"my_key": "got here and there and parallel"},
{"my_key": "got here and there and parallel and back again"},
@@ -3591,32 +3564,24 @@ def test_nested_graph_interrupts_parallel(
{"my_key": ""},
config,
stream_mode="values",
checkpoint_during=checkpoint_during,
durability=durability,
)
] == [
{"my_key": ""},
{"my_key": " and parallel"},
]
assert [
*app.stream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
)
] == [
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
{"my_key": ""},
{"my_key": "got here and there and parallel"},
]
assert [
*app.stream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
)
] == [
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
{"my_key": "got here and there and parallel"},
{"my_key": "got here and there and parallel and back again"},
]
def test_doubly_nested_graph_interrupts(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class State(TypedDict):
my_key: str
@@ -3669,13 +3634,11 @@ def test_doubly_nested_graph_interrupts(
# test invoke w/ nested interrupt
config = {"configurable": {"thread_id": "1"}}
assert app.invoke(
{"my_key": "my value"}, config, checkpoint_during=checkpoint_during
) == {
assert app.invoke({"my_key": "my value"}, config, durability=durability) == {
"my_key": "hi my value",
}
assert app.invoke(None, config, checkpoint_during=checkpoint_during) == {
assert app.invoke(None, config, durability=durability) == {
"my_key": "hi my value here and there and back again",
}
@@ -3684,14 +3647,12 @@ def test_doubly_nested_graph_interrupts(
config = {
"configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append}
}
assert [
*app.stream({"my_key": "my value"}, config, checkpoint_during=checkpoint_during)
] == [
assert [*app.stream({"my_key": "my value"}, config, durability=durability)] == [
{"parent_1": {"my_key": "hi my value"}},
{"__interrupt__": ()},
]
assert nodes == ["parent_1", "grandchild_1"]
assert [*app.stream(None, config, checkpoint_during=checkpoint_during)] == [
assert [*app.stream(None, config, durability=durability)] == [
{"child": {"my_key": "hi my value here and there"}},
{"parent_2": {"my_key": "hi my value here and there and back again"}},
]
@@ -3711,17 +3672,13 @@ def test_doubly_nested_graph_interrupts(
{"my_key": "my value"},
config,
stream_mode="values",
checkpoint_during=checkpoint_during,
durability=durability,
)
] == [
{"my_key": "my value"},
{"my_key": "hi my value"},
]
assert [
*app.stream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
)
] == [
assert [*app.stream(None, config, stream_mode="values", durability=durability)] == [
{"my_key": "hi my value"},
{"my_key": "hi my value here and there"},
{"my_key": "hi my value here and there and back again"},
@@ -4372,7 +4329,7 @@ def test_debug_retry(sync_checkpointer: BaseCheckpointSaver):
graph = builder.compile(checkpointer=sync_checkpointer)
config = {"configurable": {"thread_id": "1"}}
graph.invoke({"messages": []}, config=config, checkpoint_during=True)
graph.invoke({"messages": []}, config=config, durability="async")
# re-run step: 1
target_config = next(
@@ -4384,7 +4341,7 @@ def test_debug_retry(sync_checkpointer: BaseCheckpointSaver):
events = [
*graph.stream(
None, config=update_config, stream_mode="debug", checkpoint_during=True
None, config=update_config, stream_mode="debug", durability="async"
)
]
@@ -4417,7 +4374,7 @@ def test_debug_retry(sync_checkpointer: BaseCheckpointSaver):
def test_debug_subgraphs(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
):
class State(TypedDict):
messages: Annotated[list[str], operator.add]
@@ -4451,14 +4408,14 @@ def test_debug_subgraphs(
{"messages": []},
config=config,
stream_mode="debug",
checkpoint_during=checkpoint_during,
durability=durability,
)
]
checkpoint_events = list(
reversed([e["payload"] for e in events if e["type"] == "checkpoint"])
)
if not checkpoint_during:
if durability == "exit":
checkpoint_events = checkpoint_events[:1]
checkpoint_history = list(graph.get_state_history(config))
@@ -4489,7 +4446,7 @@ def test_debug_subgraphs(
def test_debug_nested_subgraphs(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
):
from collections import defaultdict
@@ -4533,7 +4490,7 @@ def test_debug_nested_subgraphs(
config=config,
stream_mode="debug",
subgraphs=True,
checkpoint_during=checkpoint_during,
durability=durability,
)
]
@@ -4573,9 +4530,9 @@ def test_debug_nested_subgraphs(
for checkpoint_events, checkpoint_history, ns in zip(
stream_ns.values(), history_ns.values(), stream_ns.keys()
):
if not checkpoint_during:
if durability == "exit":
checkpoint_events = checkpoint_events[-1:]
if ns: # Save no checkpoints for subgraphs when checkpoint_during=False
if ns: # Save no checkpoints for subgraphs when durability="exit"
assert not checkpoint_history
continue
assert len(checkpoint_events) == len(checkpoint_history)
@@ -4777,7 +4734,7 @@ def test_parent_command(
config = {"configurable": {"thread_id": "1"}}
assert graph.invoke(
{"messages": [("user", "get user name")]}, config, checkpoint_during=False
{"messages": [("user", "get user name")]}, config, durability="exit"
) == {
"messages": [
_AnyIdHumanMessage(
@@ -5363,7 +5320,7 @@ def test_concurrent_execution_thread_safety():
def test_checkpoint_recovery(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
):
"""Test recovery from checkpoints after failures."""
@@ -5394,7 +5351,7 @@ def test_checkpoint_recovery(
graph.invoke(
{"steps": ["start"], "attempt": 1},
config,
checkpoint_during=checkpoint_during,
durability=durability,
)
# Verify checkpoint state
@@ -5405,14 +5362,12 @@ def test_checkpoint_recovery(
assert "RuntimeError('Simulated failure')" in state.tasks[0].error
# Retry with updated attempt count
result = graph.invoke(
{"steps": [], "attempt": 2}, config, checkpoint_during=checkpoint_during
)
result = graph.invoke({"steps": [], "attempt": 2}, config, durability=durability)
assert result == {"steps": ["start", "node1", "node2"], "attempt": 2}
# Verify checkpoint history shows both attempts
history = list(graph.get_state_history(config))
if checkpoint_during:
if durability != "exit":
assert len(history) == 6 # Initial + failed attempt + successful attempt
else:
assert len(history) == 2 # error + success
@@ -5495,7 +5450,7 @@ def test_falsy_return_from_task(sync_checkpointer: BaseCheckpointSaver):
assert [
chunk
for chunk in graph.stream(
{"a": 5}, configurable, stream_mode="debug", checkpoint_during=False
{"a": 5}, configurable, stream_mode="debug", durability="exit"
)
] == [
{
@@ -5598,7 +5553,7 @@ def test_falsy_return_from_task(sync_checkpointer: BaseCheckpointSaver):
Command(resume="123"),
configurable,
stream_mode="debug",
checkpoint_during=False,
durability="exit",
)
] == [
{
@@ -7585,7 +7540,7 @@ def test_pregel_node_copy() -> None:
def test_update_as_input(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class State(TypedDict):
foo: str
@@ -7608,13 +7563,13 @@ def test_update_as_input(
assert graph.invoke(
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
durability=durability,
) == {"foo": "tool"}
assert graph.invoke(
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
durability=durability,
) == {"foo": "tool"}
def map_snapshot(i: StateSnapshot) -> dict:
@@ -7653,14 +7608,14 @@ def test_update_as_input(
for s in graph.get_state_history({"configurable": {"thread_id": "2"}})
]
if checkpoint_during:
if durability != "exit":
assert new_history == history
else:
assert [new_history[0], new_history[4]] == history
def test_batch_update_as_input(
sync_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
sync_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class State(TypedDict):
foo: str
@@ -7695,7 +7650,7 @@ def test_batch_update_as_input(
assert graph.invoke(
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
durability=durability,
) == {
"foo": "map",
"tasks": [0, 1, 2],
@@ -7749,7 +7704,7 @@ def test_batch_update_as_input(
for s in graph.get_state_history({"configurable": {"thread_id": "2"}})
]
if checkpoint_during:
if durability != "exit":
assert new_history == history
else:
assert new_history[:1] == history
libs/langgraph/tests/test_pregel_async.py
+103 -136
View File
@@ -58,6 +58,7 @@ from langgraph.store.base import BaseStore
from langgraph.types import (
CachePolicy,
Command,
Durability,
Interrupt,
PregelTask,
RetryPolicy,
@@ -177,11 +178,11 @@ async def test_checkpoint_errors() -> None:
graph = builder.compile(checkpointer=FaultyPutWritesCheckpointer())
with pytest.raises(ValueError, match="Faulty put_writes"):
await graph.ainvoke(
"", {"configurable": {"thread_id": "thread-1"}}, checkpoint_during=True
"", {"configurable": {"thread_id": "thread-1"}}, durability="async"
)
with pytest.raises(ValueError, match="Faulty put_writes"):
async for _ in graph.astream(
"", {"configurable": {"thread_id": "thread-2"}}, checkpoint_during=True
"", {"configurable": {"thread_id": "thread-2"}}, durability="async"
):
pass
with pytest.raises(ValueError, match="Faulty put_writes"):
@@ -189,7 +190,7 @@ async def test_checkpoint_errors() -> None:
"",
{"configurable": {"thread_id": "thread-3"}},
version="v2",
checkpoint_during=True,
durability="async",
):
pass
@@ -316,7 +317,7 @@ async def test_checkpoint_put_after_cancellation() -> None:
# start the task
t = asyncio.create_task(
graph.ainvoke({"hello": "world"}, thread1, checkpoint_during=False)
graph.ainvoke({"hello": "world"}, thread1, durability="exit")
)
# cancel after 0.2 seconds
await asyncio.sleep(0.2)
@@ -383,7 +384,7 @@ async def test_checkpoint_put_after_cancellation_stream_anext() -> None:
thread1 = {"configurable": {"thread_id": "1"}}
# start the task
s = graph.astream({"hello": "world"}, thread1, checkpoint_during=False)
s = graph.astream({"hello": "world"}, thread1, durability="exit")
t = asyncio.create_task(s.__anext__())
# cancel after 0.2 seconds
await asyncio.sleep(0.2)
@@ -455,7 +456,7 @@ async def test_checkpoint_put_after_cancellation_stream_events_anext() -> None:
thread1,
version="v2",
include_names=["LangGraph"],
checkpoint_during=False,
durability="exit",
)
# skip first event (happens right away)
await s.__anext__()
@@ -640,7 +641,7 @@ async def test_dynamic_interrupt(async_checkpointer: BaseCheckpointSaver) -> Non
assert [
c
async for c in tool_two.astream(
{"my_key": "value ⛰️", "market": "DE"}, thread1, checkpoint_during=False
{"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit"
)
] == [
{
@@ -807,7 +808,7 @@ async def test_dynamic_interrupt_subgraph(
assert [
c
async for c in tool_two.astream(
{"my_key": "value ⛰️", "market": "DE"}, thread1, checkpoint_during=False
{"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit"
)
] == [
{
@@ -980,7 +981,7 @@ async def test_partial_pending_checkpoint(
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert await tool_two.ainvoke(
{"my_key": "value ⛰️", "market": "DE"}, thread1, checkpoint_during=False
{"my_key": "value ⛰️", "market": "DE"}, thread1, durability="exit"
) == {
"my_key": "value ⛰️ one",
"market": "DE",
@@ -1763,7 +1764,7 @@ async def test_invoke_checkpoint(
async def test_pending_writes_resume(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class State(TypedDict):
value: Annotated[int, operator.add]
@@ -1800,7 +1801,7 @@ async def test_pending_writes_resume(
thread1: RunnableConfig = {"configurable": {"thread_id": "1"}}
with pytest.raises(ConnectionError, match="I'm not good"):
await graph.ainvoke({"value": 1}, thread1, checkpoint_during=checkpoint_during)
await graph.ainvoke({"value": 1}, thread1, durability=durability)
# both nodes should have been called once
assert one.calls == 1
@@ -1849,7 +1850,7 @@ async def test_pending_writes_resume(
# resume execution
with pytest.raises(ConnectionError, match="I'm not good"):
await graph.ainvoke(None, thread1, checkpoint_during=checkpoint_during)
await graph.ainvoke(None, thread1, durability=durability)
# node "one" succeeded previously, so shouldn't be called again
assert one.calls == 1
@@ -1863,14 +1864,12 @@ async def test_pending_writes_resume(
# resume execution, without exception
two.rtn = {"value": 3}
# both the pending write and the new write were applied, 1 + 2 + 3 = 6
assert await graph.ainvoke(None, thread1, checkpoint_during=checkpoint_during) == {
"value": 6
}
assert await graph.ainvoke(None, thread1, durability=durability) == {"value": 6}
# check all final checkpoints
checkpoints = [c async for c in async_checkpointer.alist(thread1)]
# we should have 3
assert len(checkpoints) == (3 if checkpoint_during else 2)
assert len(checkpoints) == (3 if durability != "exit" else 2)
# the last one not too interesting for this test
assert checkpoints[0] == CheckpointTuple(
config={
@@ -1969,14 +1968,14 @@ async def test_pending_writes_resume(
"checkpoint_id": checkpoints[2].config["configurable"]["checkpoint_id"],
}
}
if checkpoint_during
if durability != "exit"
else None,
pending_writes=UnsortedSequence(
(AnyStr(), "value", 2),
(AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
(AnyStr(), "value", 3),
)
if checkpoint_during
if durability != "exit"
else UnsortedSequence(
(AnyStr(), "value", 2),
(AnyStr(), "__error__", 'ConnectionError("I\'m not good")'),
@@ -1984,7 +1983,7 @@ async def test_pending_writes_resume(
# produced in a run where only the next checkpoint (the last) is saved
),
)
if not checkpoint_during:
if durability == "exit":
return
assert checkpoints[2] == CheckpointTuple(
config={
@@ -2058,7 +2057,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, checkpoint_during=True)
result = await graph.ainvoke({"myval": 1}, thread1, durability="async")
assert result["myval"] == 4
history = [c async for c in graph.aget_state_history(thread1)]
@@ -2240,7 +2239,7 @@ async def test_send_sequences(async_checkpointer: BaseCheckpointSaver) -> None:
@NEEDS_CONTEXTVARS
async def test_imp_task(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
mapper_calls = 0
@@ -2260,12 +2259,7 @@ async def test_imp_task(
tracer = FakeTracer()
thread1 = {"configurable": {"thread_id": "1"}, "callbacks": [tracer]}
assert [
c
async for c in graph.astream(
[0, 1], thread1, checkpoint_during=checkpoint_during
)
] == [
assert [c async for c in graph.astream([0, 1], thread1, durability=durability)] == [
{"mapper": "00"},
{"mapper": "11"},
{
@@ -2288,7 +2282,7 @@ async def test_imp_task(
assert any(r.inputs == {"input": 1} for r in mapper_runs)
assert await graph.ainvoke(
Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
Command(resume="answer"), thread1, durability=durability
) == [
"00answer",
"11answer",
@@ -2298,7 +2292,7 @@ async def test_imp_task(
@NEEDS_CONTEXTVARS
async def test_imp_nested(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
async def mynode(input: list[str]) -> list[str]:
return [it + "a" for it in input]
@@ -2337,12 +2331,7 @@ async def test_imp_nested(
}
thread1 = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in graph.astream(
[0, 1], thread1, checkpoint_during=checkpoint_during
)
] == [
assert [c async for c in graph.astream([0, 1], thread1, durability=durability)] == [
{"submapper": "0"},
{"mapper": "00"},
{"submapper": "1"},
@@ -2358,7 +2347,7 @@ async def test_imp_nested(
]
assert await graph.ainvoke(
Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
Command(resume="answer"), thread1, durability=durability
) == [
"00answera",
"11answera",
@@ -2367,7 +2356,7 @@ async def test_imp_nested(
@NEEDS_CONTEXTVARS
async def test_imp_task_cancel(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
mapper_calls = 0
mapper_cancels = 0
@@ -2393,12 +2382,7 @@ async def test_imp_task_cancel(
return [m + answer for m in mapped]
thread1 = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in graph.astream(
[0, 1], thread1, checkpoint_during=checkpoint_during
)
] == [
assert [c async for c in graph.astream([0, 1], thread1, durability=durability)] == [
{"mapper": "00"},
{
"__interrupt__": (
@@ -2413,7 +2397,7 @@ async def test_imp_task_cancel(
assert mapper_cancels == 1
assert await graph.ainvoke(
Command(resume="answer"), thread1, checkpoint_during=checkpoint_during
Command(resume="answer"), thread1, durability=durability
) == [
"00answer",
]
@@ -2423,7 +2407,7 @@ async def test_imp_task_cancel(
@NEEDS_CONTEXTVARS
async def test_imp_sync_from_async(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
@task()
def foo(state: dict) -> dict:
@@ -2446,10 +2430,7 @@ async def test_imp_sync_from_async(
thread1 = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in graph.astream(
{"a": "0"}, thread1, checkpoint_during=checkpoint_during
)
c async for c in graph.astream({"a": "0"}, thread1, durability=durability)
] == [
{"foo": {"a": "0foo", "b": "bar"}},
{"bar": {"a": "0foobar", "c": "bark"}},
@@ -2460,7 +2441,7 @@ async def test_imp_sync_from_async(
@NEEDS_CONTEXTVARS
async def test_imp_stream_order(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
@task()
async def foo(state: dict) -> dict:
@@ -2484,10 +2465,7 @@ async def test_imp_stream_order(
thread1 = {"configurable": {"thread_id": "1"}}
assert [
c
async for c in graph.astream(
{"a": "0"}, thread1, checkpoint_during=checkpoint_during
)
c async for c in graph.astream({"a": "0"}, thread1, durability=durability)
] == [
{"foo": {"a": "0foo", "b": "bar"}},
{"bar": {"a": "0foobar", "c": "bark"}},
@@ -2501,7 +2479,7 @@ async def test_imp_stream_order(
reason="Requires Python 3.11 or higher for context management",
)
async def test_send_dedupe_on_resume(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class InterruptOnce:
ticks: int = 0
@@ -2552,7 +2530,7 @@ async def test_send_dedupe_on_resume(
graph = builder.compile(checkpointer=async_checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
assert await graph.ainvoke(["0"], thread1, checkpoint_during=checkpoint_during) == {
assert await graph.ainvoke(["0"], thread1, durability=durability) == {
"__interrupt__": [
Interrupt(
value="Bahh",
@@ -2563,7 +2541,7 @@ async def test_send_dedupe_on_resume(
assert builder.nodes["2"].runnable.func.ticks == 3
assert builder.nodes["flaky"].runnable.func.ticks == 1
# resume execution
assert await graph.ainvoke(None, thread1, checkpoint_during=checkpoint_during) == [
assert await graph.ainvoke(None, thread1, durability=durability) == [
"0",
"1",
"3.1",
@@ -2580,7 +2558,7 @@ async def test_send_dedupe_on_resume(
assert builder.nodes["flaky"].runnable.func.ticks == 2
# check history
history = [c async for c in graph.aget_state_history(thread1)]
assert len(history) == (6 if checkpoint_during else 2)
assert len(history) == (6 if durability != "exit" else 2)
expected_history = [
StateSnapshot(
values=[
@@ -2709,7 +2687,7 @@ async def test_send_dedupe_on_resume(
error=None,
interrupts=(Interrupt(value="Bahh", id=AnyStr()),),
state=None,
result=["flaky|4"] if checkpoint_during else None,
result=["flaky|4"] if durability != "exit" else None,
),
PregelTask(
id=AnyStr(),
@@ -2844,7 +2822,7 @@ async def test_send_dedupe_on_resume(
interrupts=(),
),
]
if checkpoint_during:
if durability != "exit":
assert history == expected_history
else:
assert history[0] == expected_history[0]._replace(
@@ -2955,7 +2933,7 @@ async def test_send_react_interrupt(async_checkpointer: BaseCheckpointSaver) ->
graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"])
thread1 = {"configurable": {"thread_id": "2"}}
assert await graph.ainvoke(
{"messages": [HumanMessage("hello")]}, thread1, checkpoint_during=False
{"messages": [HumanMessage("hello")]}, thread1, durability="exit"
) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
@@ -3079,7 +3057,7 @@ async def test_send_react_interrupt(async_checkpointer: BaseCheckpointSaver) ->
graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"])
thread1 = {"configurable": {"thread_id": "3"}}
assert await graph.ainvoke(
{"messages": [HumanMessage("hello")]}, thread1, checkpoint_during=False
{"messages": [HumanMessage("hello")]}, thread1, durability="exit"
) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
@@ -3344,7 +3322,7 @@ async def test_send_react_interrupt_control(
graph = builder.compile(checkpointer=async_checkpointer, interrupt_before=["foo"])
thread1 = {"configurable": {"thread_id": "2"}}
assert await graph.ainvoke(
{"messages": [HumanMessage("hello")]}, thread1, checkpoint_during=False
{"messages": [HumanMessage("hello")]}, thread1, durability="exit"
) == {
"messages": [
_AnyIdHumanMessage(content="hello"),
@@ -3631,7 +3609,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, checkpoint_during=True) == 2
assert await app.ainvoke(2, thread_1, durability="async") == 2
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 2
@@ -3640,7 +3618,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, checkpoint_during=True) == 5
assert await app.ainvoke(3, thread_1, durability="async") == 5
state = await app.aget_state(thread_1)
assert state is not None
assert state.values.get("total") == 7
@@ -3650,7 +3628,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, checkpoint_during=True)
await app.ainvoke(4, thread_1, durability="async")
# checkpoint is not updated
state = await app.aget_state(thread_1)
assert state is not None
@@ -3658,7 +3636,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, checkpoint_during=True) == 9
assert await app.ainvoke(2, thread_1, durability="async") == 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"
@@ -4960,7 +4938,7 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None:
async def test_subgraph_checkpoint_true(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class InnerState(TypedDict):
my_key: Annotated[str, operator.add]
@@ -4998,7 +4976,7 @@ async def test_subgraph_checkpoint_true(
{"my_key": ""},
config,
subgraphs=True,
checkpoint_during=checkpoint_during,
durability=durability,
)
] == [
(("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}),
@@ -5025,7 +5003,9 @@ async def test_subgraph_checkpoint_true(
]
async def test_subgraph_checkpoint_during_false_inherited() -> None:
async def test_subgraph_durability_inherited(
durability: Durability,
) -> None:
async_checkpointer = InMemorySaver()
class InnerState(TypedDict):
@@ -5056,23 +5036,20 @@ async def test_subgraph_checkpoint_during_false_inherited() -> None:
"inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END
)
app = graph.compile(checkpointer=async_checkpointer)
for checkpoint_during in [True, False]:
thread_id = str(uuid.uuid4())
config = {"configurable": {"thread_id": thread_id}}
await app.ainvoke(
{"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during
)
if checkpoint_during:
checkpoints = list(async_checkpointer.list(config))
assert len(checkpoints) == 12
else:
checkpoints = list(async_checkpointer.list(config))
assert len(checkpoints) == 1
thread_id = str(uuid.uuid4())
config = {"configurable": {"thread_id": thread_id}}
await app.ainvoke({"my_key": ""}, config, subgraphs=True, durability=durability)
if durability != "exit":
checkpoints = list(async_checkpointer.list(config))
assert len(checkpoints) == 12
else:
checkpoints = list(async_checkpointer.list(config))
assert len(checkpoints) == 1
@NEEDS_CONTEXTVARS
async def test_subgraph_checkpoint_true_interrupt(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
# Define subgraph
class SubgraphState(TypedDict):
@@ -5113,9 +5090,7 @@ async def test_subgraph_checkpoint_true_interrupt(
graph = builder.compile(checkpointer=async_checkpointer)
config = {"configurable": {"thread_id": "1"}}
assert await graph.ainvoke(
{"foo": "foo"}, config, checkpoint_during=checkpoint_during
) == {
assert await graph.ainvoke({"foo": "foo"}, config, durability=durability) == {
"foo": "hi! foo",
"__interrupt__": [
Interrupt(
@@ -5128,7 +5103,7 @@ async def test_subgraph_checkpoint_true_interrupt(
"bar": "hi! foo"
}
assert await graph.ainvoke(
Command(resume="baz"), config, checkpoint_during=checkpoint_during
Command(resume="baz"), config, durability=durability
) == {"foo": "hi! foobaz"}
@@ -5242,7 +5217,7 @@ async def test_stream_buffering_single_node(
async def test_nested_graph_interrupts_parallel(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class InnerState(TypedDict):
my_key: Annotated[str, operator.add]
@@ -5291,13 +5266,11 @@ async def test_nested_graph_interrupts_parallel(
# test invoke w/ nested interrupt
config = {"configurable": {"thread_id": "1"}}
assert await app.ainvoke(
{"my_key": ""}, config, checkpoint_during=checkpoint_during
) == {
assert await app.ainvoke({"my_key": ""}, config, durability=durability) == {
"my_key": " and parallel",
}
assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == {
assert await app.ainvoke(None, config, durability=durability) == {
"my_key": "got here and there and parallel and back again",
}
@@ -5312,7 +5285,7 @@ async def test_nested_graph_interrupts_parallel(
{"my_key": ""},
config,
subgraphs=True,
checkpoint_during=checkpoint_during,
durability=durability,
)
] == [
# we got to parallel node first
@@ -5323,9 +5296,7 @@ async def test_nested_graph_interrupts_parallel(
),
((), {"__interrupt__": ()}),
]
assert [
c async for c in app.astream(None, config, checkpoint_during=checkpoint_during)
] == [
assert [c async for c in app.astream(None, config, durability=durability)] == [
{"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}},
{"inner": {"my_key": "got here and there"}},
{"outer_2": {"my_key": " and back again"}},
@@ -5339,7 +5310,7 @@ async def test_nested_graph_interrupts_parallel(
{"my_key": ""},
config,
stream_mode="values",
checkpoint_during=checkpoint_during,
durability=durability,
)
] == [
{"my_key": ""},
@@ -5348,7 +5319,7 @@ async def test_nested_graph_interrupts_parallel(
assert [
c
async for c in app.astream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
None, config, stream_mode="values", durability=durability
)
] == [
{"my_key": ""},
@@ -5365,7 +5336,7 @@ async def test_nested_graph_interrupts_parallel(
{"my_key": ""},
config,
stream_mode="values",
checkpoint_during=checkpoint_during,
durability=durability,
)
] == [
{"my_key": ""},
@@ -5374,7 +5345,7 @@ async def test_nested_graph_interrupts_parallel(
assert [
c
async for c in app.astream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
None, config, stream_mode="values", durability=durability
)
] == [
{"my_key": ""},
@@ -5383,7 +5354,7 @@ async def test_nested_graph_interrupts_parallel(
assert [
c
async for c in app.astream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
None, config, stream_mode="values", durability=durability
)
] == [
{"my_key": ""},
@@ -5400,7 +5371,7 @@ async def test_nested_graph_interrupts_parallel(
{"my_key": ""},
config,
stream_mode="values",
checkpoint_during=checkpoint_during,
durability=durability,
)
] == [
{"my_key": ""},
@@ -5409,7 +5380,7 @@ async def test_nested_graph_interrupts_parallel(
assert [
c
async for c in app.astream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
None, config, stream_mode="values", durability=durability
)
] == [
{"my_key": ""},
@@ -5418,7 +5389,7 @@ async def test_nested_graph_interrupts_parallel(
assert [
c
async for c in app.astream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
None, config, stream_mode="values", durability=durability
)
] == [
{"my_key": "got here and there and parallel"},
@@ -5427,7 +5398,7 @@ async def test_nested_graph_interrupts_parallel(
async def test_doubly_nested_graph_interrupts(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class State(TypedDict):
my_key: str
@@ -5480,13 +5451,11 @@ async def test_doubly_nested_graph_interrupts(
# test invoke w/ nested interrupt
config = {"configurable": {"thread_id": "1"}}
assert await app.ainvoke(
{"my_key": "my value"}, config, checkpoint_during=checkpoint_during
) == {
assert await app.ainvoke({"my_key": "my value"}, config, durability=durability) == {
"my_key": "hi my value",
}
assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == {
assert await app.ainvoke(None, config, durability=durability) == {
"my_key": "hi my value here and there and back again",
}
@@ -5498,16 +5467,14 @@ async def test_doubly_nested_graph_interrupts(
assert [
c
async for c in app.astream(
{"my_key": "my value"}, config, checkpoint_during=checkpoint_during
{"my_key": "my value"}, config, durability=durability
)
] == [
{"parent_1": {"my_key": "hi my value"}},
{"__interrupt__": ()},
]
assert nodes == ["parent_1", "grandchild_1"]
assert [
c async for c in app.astream(None, config, checkpoint_during=checkpoint_during)
] == [
assert [c async for c in app.astream(None, config, durability=durability)] == [
{"child": {"my_key": "hi my value here and there"}},
{"parent_2": {"my_key": "hi my value here and there and back again"}},
]
@@ -5528,7 +5495,7 @@ async def test_doubly_nested_graph_interrupts(
{"my_key": "my value"},
config,
stream_mode="values",
checkpoint_during=checkpoint_during,
durability=durability,
)
] == [
{"my_key": "my value"},
@@ -5537,7 +5504,7 @@ async def test_doubly_nested_graph_interrupts(
assert [
c
async for c in app.astream(
None, config, stream_mode="values", checkpoint_during=checkpoint_during
None, config, stream_mode="values", durability=durability
)
] == [
{"my_key": "hi my value"},
@@ -5837,7 +5804,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, checkpoint_during=True)
await graph.ainvoke({"messages": []}, config=config, durability="async")
# re-run step: 1
async for c in async_checkpointer.alist(config):
@@ -5851,7 +5818,7 @@ async def test_debug_retry(async_checkpointer: BaseCheckpointSaver):
events = [
c
async for c in graph.astream(
None, config=update_config, stream_mode="debug", checkpoint_during=True
None, config=update_config, stream_mode="debug", durability="async"
)
]
@@ -5884,7 +5851,7 @@ async def test_debug_retry(async_checkpointer: BaseCheckpointSaver):
async def test_debug_subgraphs(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
):
class State(TypedDict):
messages: Annotated[list[str], operator.add]
@@ -5919,14 +5886,14 @@ async def test_debug_subgraphs(
{"messages": []},
config=config,
stream_mode="debug",
checkpoint_during=checkpoint_during,
durability=durability,
)
]
checkpoint_events = list(
reversed([e["payload"] for e in events if e["type"] == "checkpoint"])
)
if not checkpoint_during:
if durability == "exit":
checkpoint_events = checkpoint_events[:1]
checkpoint_history = [c async for c in graph.aget_state_history(config)]
@@ -5955,7 +5922,7 @@ async def test_debug_subgraphs(
async def test_debug_nested_subgraphs(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
from collections import defaultdict
@@ -6000,7 +5967,7 @@ async def test_debug_nested_subgraphs(
config=config,
stream_mode="debug",
subgraphs=True,
checkpoint_during=checkpoint_during,
durability=durability,
)
]
@@ -6045,9 +6012,9 @@ async def test_debug_nested_subgraphs(
for checkpoint_events, checkpoint_history, ns in zip(
stream_ns.values(), history_ns.values(), stream_ns.keys()
):
if not checkpoint_during:
if durability == "exit":
checkpoint_events = checkpoint_events[-1:]
if ns: # Save no checkpoints for subgraphs when checkpoint_during=False
if ns: # Save no checkpoints for subgraphs when durability="exit"
assert not checkpoint_history
continue
assert len(checkpoint_events) == len(checkpoint_history)
@@ -6101,7 +6068,7 @@ async def test_parent_command(
config = {"configurable": {"thread_id": "1"}}
assert await graph.ainvoke(
{"messages": [("user", "get user name")]}, config, checkpoint_during=False
{"messages": [("user", "get user name")]}, config, durability="exit"
) == {
"messages": [
_AnyIdHumanMessage(
@@ -6606,7 +6573,7 @@ async def test_concurrent_execution():
async def test_checkpoint_recovery_async(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
"""Test recovery from checkpoints after failures with async nodes."""
@@ -6639,7 +6606,7 @@ async def test_checkpoint_recovery_async(
await graph.ainvoke(
{"steps": ["start"], "attempt": 1},
config,
checkpoint_during=checkpoint_during,
durability=durability,
)
# Verify checkpoint state
@@ -6650,13 +6617,13 @@ async def test_checkpoint_recovery_async(
# Retry with updated attempt count
result = await graph.ainvoke(
{"steps": [], "attempt": 2}, config, checkpoint_during=checkpoint_during
{"steps": [], "attempt": 2}, config, durability=durability
)
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)]
if checkpoint_during:
if durability != "exit":
assert len(history) == 6 # Initial + failed attempt + successful attempt
else:
assert len(history) == 2 # error + success
@@ -8112,7 +8079,7 @@ async def test_bulk_state_updates(async_checkpointer: BaseCheckpointSaver) -> No
async def test_update_as_input(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class State(TypedDict):
foo: str
@@ -8135,13 +8102,13 @@ async def test_update_as_input(
assert await graph.ainvoke(
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
durability=durability,
) == {"foo": "tool"}
assert await graph.ainvoke(
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
durability=durability,
) == {"foo": "tool"}
def map_snapshot(i: StateSnapshot) -> dict:
@@ -8180,14 +8147,14 @@ async def test_update_as_input(
async for s in graph.aget_state_history({"configurable": {"thread_id": "2"}})
]
if checkpoint_during:
if durability != "exit":
assert new_history == history
else:
assert [new_history[0], new_history[4]] == history
async def test_batch_update_as_input(
async_checkpointer: BaseCheckpointSaver, checkpoint_during: bool
async_checkpointer: BaseCheckpointSaver, durability: Durability
) -> None:
class State(TypedDict):
foo: str
@@ -8222,7 +8189,7 @@ async def test_batch_update_as_input(
assert await graph.ainvoke(
{"foo": "input"},
{"configurable": {"thread_id": "1"}},
checkpoint_during=checkpoint_during,
durability=durability,
) == {"foo": "map", "tasks": [0, 1, 2]}
def map_snapshot(i: StateSnapshot) -> dict:
@@ -8273,7 +8240,7 @@ async def test_batch_update_as_input(
async for s in graph.aget_state_history({"configurable": {"thread_id": "2"}})
]
if checkpoint_during:
if durability != "exit":
assert new_history == history
else:
assert new_history[:1] == history