From 0afc4ebda3eb219957d378a4ecf34da9745f1a8e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 28 May 2025 14:13:24 -0700 Subject: [PATCH] Flip default for checkpoint_during - 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 --- libs/langgraph/langgraph/pregel/__init__.py | 12 +- libs/langgraph/langgraph/pregel/debug.py | 19 +- libs/langgraph/langgraph/pregel/loop.py | 9 +- .../tests/test_checkpoint_migration.py | 15 +- libs/langgraph/tests/test_large_cases.py | 972 ++-------------- .../langgraph/tests/test_large_cases_async.py | 1033 ++--------------- libs/langgraph/tests/test_pregel.py | 137 ++- libs/langgraph/tests/test_pregel_async.py | 175 ++- 8 files changed, 380 insertions(+), 1992 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 1cd6a225d..532212dc2 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -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, ): diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index b5a213ec4..6429cf0c4 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -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], diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index eb56d5885..80aaab9bc 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -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, "" ), diff --git a/libs/langgraph/tests/test_checkpoint_migration.py b/libs/langgraph/tests/test_checkpoint_migration.py index 307195924..e284af3d5 100644 --- a/libs/langgraph/tests/test_checkpoint_migration.py +++ b/libs/langgraph/tests/test_checkpoint_migration.py @@ -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"}}, ] diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 6cc21979e..6a6317f80 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -72,7 +72,7 @@ def test_invoke_two_processes_in_out_interrupt( thread2 = {"configurable": {"thread_id": "2"}} # start execution, stop at inbox - assert app.invoke(2, thread1) is None + assert app.invoke(2, thread1, checkpoint_during=True) is None # inbox == 3 checkpoint = sync_checkpointer.get(thread1) @@ -80,10 +80,10 @@ def test_invoke_two_processes_in_out_interrupt( assert checkpoint["channel_values"]["inbox"] == 3 # resume execution, finish - assert app.invoke(None, thread1) == 4 + assert app.invoke(None, thread1, checkpoint_during=True) == 4 # start execution again, stop at inbox - assert app.invoke(20, thread1) is None + assert app.invoke(20, thread1, checkpoint_during=True) is None # inbox == 21 checkpoint = sync_checkpointer.get(thread1) @@ -91,11 +91,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) is None - assert app.invoke(None, thread1) == 5 + assert app.invoke(3, thread1, checkpoint_during=True) is None + assert app.invoke(None, thread1, checkpoint_during=True) == 5 # start execution again, stopping at inbox - assert app.invoke(20, thread2) is None + assert app.invoke(20, thread2, checkpoint_during=True) is None # inbox == 21 snapshot = app.get_state(thread2) @@ -112,6 +112,7 @@ def test_invoke_two_processes_in_out_interrupt( # list history history = [c for c in app.get_state_history(thread1)] + assert len(history) == 8 assert history == [ StateSnapshot( values={"inbox": 4, "output": 5, "input": 3}, @@ -308,7 +309,11 @@ def test_fork_always_re_runs_nodes( thread1 = {"configurable": {"thread_id": "1"}} # start execution, stop at inbox - assert [*graph.stream(1, thread1, stream_mode=["values", "updates"])] == [ + assert [ + *graph.stream( + 1, thread1, stream_mode=["values", "updates"], checkpoint_during=True + ) + ] == [ ("values", 1), ("updates", {"add_one": 1}), ("values", 2), @@ -741,9 +746,7 @@ def test_conditional_graph( "step": 0, "thread_id": "1", }, - parent_config=( - list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=None, interrupts=(), ) assert ( @@ -956,9 +959,7 @@ def test_conditional_graph( "step": 0, "thread_id": "2", }, - parent_config=( - list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=None, interrupts=(), ) @@ -1165,9 +1166,7 @@ def test_conditional_graph( "step": 0, "thread_id": "3", }, - parent_config=( - list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=None, interrupts=(), ) @@ -1513,9 +1512,7 @@ def test_conditional_state_graph( "step": 1, "thread_id": "1", }, - parent_config=( - list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=None, interrupts=(), ) @@ -1683,9 +1680,7 @@ def test_conditional_state_graph( "step": 1, "thread_id": "2", }, - parent_config=( - list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=None, interrupts=(), ) @@ -1842,9 +1837,7 @@ def test_conditional_state_graph( "step": 0, "thread_id": "3", }, - parent_config=( - list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=None, interrupts=(), ) @@ -2001,9 +1994,7 @@ def test_conditional_state_graph( "step": 1, "thread_id": "4", }, - parent_config=( - list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=None, interrupts=(), ) @@ -2722,9 +2713,7 @@ def test_state_graph_packets( "step": 1, "thread_id": "1", }, - parent_config=( - [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config - ), + parent_config=None, interrupts=(), ) @@ -2991,9 +2980,7 @@ def test_state_graph_packets( "step": 1, "thread_id": "2", }, - parent_config=( - [*app_w_interrupt.checkpointer.list(config, limit=2)][-1].config - ), + parent_config=None, interrupts=(), ) @@ -3458,9 +3445,7 @@ def test_message_graph( "step": 1, "thread_id": "1", }, - parent_config=( - list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=None, interrupts=(), ) @@ -3684,9 +3669,7 @@ def test_message_graph( "step": 1, "thread_id": "2", }, - parent_config=( - list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=None, interrupts=(), ) @@ -4185,9 +4168,7 @@ def test_root_graph( "step": 1, "thread_id": "1", }, - parent_config=( - list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=None, interrupts=(), ) @@ -4413,9 +4394,7 @@ def test_root_graph( "step": 1, "thread_id": "2", }, - parent_config=( - list(app_w_interrupt.checkpointer.list(config, limit=2))[-1].config - ), + parent_config=None, interrupts=(), ) @@ -5072,12 +5051,6 @@ def test_dynamic_interrupt(sync_checkpointer: BaseCheckpointSaver) -> None: "step": 0, "thread_id": "1", }, - { - "parents": {}, - "source": "input", - "step": -1, - "thread_id": "1", - }, ] assert tool_two.get_state(thread1) == StateSnapshot( @@ -5111,7 +5084,7 @@ def test_dynamic_interrupt(sync_checkpointer: BaseCheckpointSaver) -> None: "step": 0, "thread_id": "1", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), + parent_config=None, interrupts=( Interrupt( value="Just because...", @@ -5245,12 +5218,6 @@ def test_copy_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: "step": 0, "thread_id": "1", }, - { - "parents": {}, - "source": "input", - "step": -1, - "thread_id": "1", - }, ] assert tool_two.get_state(thread1) == StateSnapshot( @@ -5290,7 +5257,7 @@ def test_copy_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: "step": 0, "thread_id": "1", }, - parent_config=([*tool_two.checkpointer.list(thread1, limit=2)][-1].config), + parent_config=None, interrupts=( Interrupt( value="Just because...", @@ -5336,9 +5303,7 @@ def test_copy_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None: "step": 1, "thread_id": "1", }, - parent_config=( - [*tool_two.checkpointer.list(thread1, limit=2)][-1].parent_config - ), + parent_config=([*tool_two.checkpointer.list(thread1, limit=2)][-1].config), interrupts=(), ) @@ -5452,12 +5417,6 @@ def test_dynamic_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver) -> N "step": 0, "thread_id": "1", }, - { - "parents": {}, - "source": "input", - "step": -1, - "thread_id": "1", - }, ] assert tool_two.get_state(thread1) == StateSnapshot( @@ -5497,13 +5456,7 @@ def test_dynamic_interrupt_subgraph(sync_checkpointer: BaseCheckpointSaver) -> N "step": 0, "thread_id": "1", }, - parent_config=( - list( - tool_two.checkpointer.list( - {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2 - ) - )[-1].config - ), + parent_config=None, interrupts=( Interrupt( value="Just because...", @@ -5607,13 +5560,6 @@ def test_start_branch_then( "assistant_id": "a", "thread_id": "1", }, - { - "parents": {}, - "source": "input", - "step": -1, - "assistant_id": "a", - "thread_id": "1", - }, ] assert tool_two.get_state(thread1) == StateSnapshot( @@ -5635,7 +5581,7 @@ def test_start_branch_then( "assistant_id": "a", "thread_id": "1", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), + parent_config=None, interrupts=(), ) # resume, for same result as above @@ -5691,7 +5637,7 @@ def test_start_branch_then( "assistant_id": "a", "thread_id": "2", }, - parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config), + parent_config=None, interrupts=(), ) # resume, for same result as above @@ -5747,7 +5693,7 @@ def test_start_branch_then( "assistant_id": "b", "thread_id": "3", }, - parent_config=(list(tool_two.checkpointer.list(thread3, limit=2))[-1].config), + parent_config=None, interrupts=(), ) # update state @@ -5844,7 +5790,10 @@ def test_branch_then( res = [ *tool_two.stream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + {"my_key": "value", "market": "DE"}, + thread10, + stream_mode="debug", + checkpoint_during=True, ) ] @@ -5855,10 +5804,6 @@ def test_branch_then( "step": -1, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -5890,10 +5835,6 @@ def test_branch_then( "step": 0, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -5911,10 +5852,6 @@ def test_branch_then( "thread_id": "10", }, "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -5956,10 +5893,6 @@ def test_branch_then( "step": 1, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -5977,10 +5910,6 @@ def test_branch_then( "thread_id": "10", }, "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -6027,10 +5956,6 @@ def test_branch_then( "step": 2, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -6048,10 +5973,6 @@ def test_branch_then( "thread_id": "10", }, "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -6096,10 +6017,6 @@ def test_branch_then( "step": 3, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -6117,10 +6034,6 @@ def test_branch_then( "thread_id": "10", }, "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -6166,7 +6079,7 @@ def test_branch_then( "step": 1, "thread_id": "1", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), + parent_config=None, interrupts=(), ) # resume, for same result as above @@ -6220,7 +6133,7 @@ def test_branch_then( "step": 1, "thread_id": "2", }, - parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config), + parent_config=None, interrupts=(), ) # resume, for same result as above @@ -6282,7 +6195,7 @@ def test_branch_then( "step": 2, "thread_id": "11", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), + parent_config=None, interrupts=(), ) @@ -6345,7 +6258,7 @@ def test_branch_then( "step": 1, "thread_id": "21", }, - parent_config=(list(tool_two.checkpointer.list(thread1, limit=2))[-1].config), + parent_config=None, interrupts=(), ) # resume, for same result as above @@ -6399,7 +6312,7 @@ def test_branch_then( "step": 1, "thread_id": "22", }, - parent_config=(list(tool_two.checkpointer.list(thread2, limit=2))[-1].config), + parent_config=None, interrupts=(), ) # resume, for same result as above @@ -6934,7 +6847,7 @@ def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None: app.invoke({"my_key": "my value"}, config, debug=True) # test state w/ nested subgraph state (right after interrupt) # first get_state without subgraph state - assert app.get_state(config) == StateSnapshot( + expected = StateSnapshot( values={"my_key": "hi my value"}, tasks=( PregelTask( @@ -6959,17 +6872,11 @@ def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, interrupts=(), ) + assert app.get_state(config) == expected + assert list(app.get_state_history(config)) == [expected] # now, get_state with subgraphs state assert app.get_state(config, subgraphs=True) == StateSnapshot( values={"my_key": "hi my value"}, @@ -7015,18 +6922,7 @@ def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None: "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - } - ), + parent_config=None, interrupts=(), ), ), @@ -7046,129 +6942,12 @@ def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, interrupts=(), ) - # get_state_history returns outer graph checkpoints - history = list(app.get_state_history(config)) - expected_history = [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - } - }, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - interrupts=(), - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), - ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - interrupts=(), - ), - ] - - assert history == expected_history # get_state_history for a subgraph returns its checkpoints - child_history = [*app.get_state_history(history[0].tasks[0].state)] + child_history = [*app.get_state_history(app.get_state(config).tasks[0].state)] expected_child_history = [ StateSnapshot( values={"my_key": "hi my value here", "my_other_key": "hi my value"}, @@ -7195,107 +6974,11 @@ def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None: "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - } - ), + parent_config=None, interrupts=(), tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "step": 0, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=( - PregelTask( - AnyStr(), - "inner_1", - (PULL, "inner_1"), - result={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - ), - ), - interrupts=(), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config=None, - interrupts=(), - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), ] - assert child_history == expected_child_history # resume @@ -7362,40 +7045,6 @@ def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None: ), interrupts=(), ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=( - PregelTask( - AnyStr(), - "outer_2", - (PULL, "outer_2"), - result={"my_key": "hi my value here and there and back again"}, - ), - ), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), StateSnapshot( values={"my_key": "hi my value"}, tasks=( @@ -7406,7 +7055,7 @@ def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None: state={ "configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()} }, - result={"my_key": "hi my value here and there"}, + result=None, ), ), next=("inner",), @@ -7424,74 +7073,6 @@ def test_nested_graph_state(sync_checkpointer: BaseCheckpointSaver) -> None: "thread_id": "1", }, created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), - ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), parent_config=None, interrupts=(), ), @@ -7597,15 +7178,7 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, interrupts=(), ) child_state = app.get_state(outer_state.tasks[0].state) @@ -7650,21 +7223,7 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - } - } - ), + parent_config=None, interrupts=(), ) grandchild_state = app.get_state(child_state.tasks[0].state) @@ -7709,22 +7268,7 @@ def test_doubly_nested_graph_state( "langgraph_triggers": ["branch:to:child_1"], }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - } - ), + parent_config=None, interrupts=(), ) # get state with subgraphs @@ -7790,24 +7334,7 @@ def test_doubly_nested_graph_state( ], }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr( - re.compile(r"child:.+|child1:") - ): AnyStr(), - } - ), - } - } - ), + parent_config=None, interrupts=(), ), ), @@ -7835,18 +7362,7 @@ def test_doubly_nested_graph_state( "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - } - ), + parent_config=None, interrupts=(), ), ), @@ -7866,15 +7382,7 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, interrupts=(), ) # # resume @@ -7952,40 +7460,6 @@ def test_doubly_nested_graph_state( }, interrupts=(), ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("parent_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "step": 2, - "parents": {}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="parent_2", - path=(PULL, "parent_2"), - result={"my_key": "hi my value here and there and back again"}, - ), - ), - interrupts=(), - ), StateSnapshot( values={"my_key": "hi my value"}, tasks=( @@ -7999,7 +7473,7 @@ def test_doubly_nested_graph_state( "checkpoint_ns": AnyStr("child"), } }, - result={"my_key": "hi my value here and there"}, + result=None, ), ), next=("child",), @@ -8017,80 +7491,12 @@ def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("parent_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "step": 0, - "parents": {}, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="parent_1", - path=(PULL, "parent_1"), - result={"my_key": "hi my value"}, - ), - ), - interrupts=(), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "input", - "step": -1, - "parents": {}, - "thread_id": "1", - }, - created_at=AnyStr(), parent_config=None, interrupts=(), - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), ), ] # get child graph history - child_history = list(app.get_state_history(outer_history[2].tasks[0].state)) + child_history = list(app.get_state_history(outer_history[1].tasks[0].state)) assert child_history == [ StateSnapshot( values={"my_key": "hi my value here and there"}, @@ -8155,16 +7561,7 @@ def test_doubly_nested_graph_state( "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, + parent_config=None, tasks=( PregelTask( id=AnyStr(), @@ -8176,47 +7573,11 @@ def test_doubly_nested_graph_state( "checkpoint_ns": AnyStr("child:"), } }, - result={"my_key": "hi my value here and there"}, + result=None, ), ), interrupts=(), ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child"], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config=None, - interrupts=(), - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), ] # get grandchild graph history grandchild_history = list(app.get_state_history(child_history[1].tasks[0].state)) @@ -8312,150 +7673,17 @@ def test_doubly_nested_graph_state( "langgraph_triggers": ["branch:to:child_1"], }, created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, + parent_config=None, tasks=( PregelTask( id=AnyStr(), name="grandchild_2", path=(PULL, "grandchild_2"), - result={"my_key": "hi my value here and there"}, + result=None, ), ), interrupts=(), ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("grandchild_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "loop", - "step": 0, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": ["branch:to:child_1"], - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="grandchild_1", - path=(PULL, "grandchild_1"), - result={"my_key": "hi my value here"}, - ), - ), - interrupts=(), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "input", - "step": -1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": ["branch:to:child_1"], - }, - created_at=AnyStr(), - parent_config=None, - interrupts=(), - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - ), - ] - - # replay grandchild checkpoint - assert [ - c for c in app.stream(None, grandchild_history[2].config, subgraphs=True) - ] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, - ), - ((), {"__interrupt__": ()}), ] @@ -8677,15 +7905,7 @@ def test_send_react_interrupt( "thread_id": "2", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "2", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, interrupts=(), tasks=( PregelTask( @@ -8809,15 +8029,7 @@ def test_send_react_interrupt( "thread_id": "3", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "3", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, interrupts=(), tasks=( PregelTask( @@ -9083,15 +8295,7 @@ 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(), @@ -9325,15 +8529,7 @@ def test_weather_subgraph( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, tasks=( PregelTask( id=AnyStr(), @@ -9415,15 +8611,7 @@ def test_weather_subgraph( "thread_id": "14", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, tasks=( PregelTask( id=AnyStr(), @@ -9462,21 +8650,7 @@ def test_weather_subgraph( "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("weather_graph:"): AnyStr(), - } - ), - } - } - ), + parent_config=None, tasks=( PregelTask( id=AnyStr(), @@ -9516,15 +8690,7 @@ def test_weather_subgraph( "thread_id": "14", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, tasks=( PregelTask( id=AnyStr(), diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index d4daa7367..777e5c7e4 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -64,7 +64,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) is None + assert await app.ainvoke(2, thread1, checkpoint_during=True) is None # inbox == 3 checkpoint = await async_checkpointer.aget(thread1) @@ -72,10 +72,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) == 4 + assert await app.ainvoke(None, thread1, checkpoint_during=True) == 4 # start execution again, stop at inbox - assert await app.ainvoke(20, thread1) is None + assert await app.ainvoke(20, thread1, checkpoint_during=True) is None # inbox == 21 checkpoint = await async_checkpointer.aget(thread1) @@ -83,11 +83,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) is None - assert await app.ainvoke(None, thread1) == 5 + assert await app.ainvoke(3, thread1, checkpoint_during=True) is None + assert await app.ainvoke(None, thread1, checkpoint_during=True) == 5 # start execution again, stopping at inbox - assert await app.ainvoke(20, thread2) is None + assert await app.ainvoke(20, thread2, checkpoint_during=True) is None # inbox == 21 snapshot = await app.aget_state(thread2) @@ -104,6 +104,7 @@ async def test_invoke_two_processes_in_out_interrupt( # list history history = [c async for c in app.aget_state_history(thread1)] + assert len(history) == 8 assert history == [ StateSnapshot( values={"inbox": 4, "output": 5, "input": 3}, @@ -307,7 +308,10 @@ async def test_fork_always_re_runs_nodes( # start execution, stop at inbox assert [ - c async for c in graph.astream(1, thread1, stream_mode=["values", "updates"]) + c + async for c in graph.astream( + 1, thread1, stream_mode=["values", "updates"], checkpoint_during=True + ) ] == [ ("values", 1), ("updates", {"add_one": 1}), @@ -803,9 +807,7 @@ async def test_conditional_graph(async_checkpointer: BaseCheckpointSaver) -> Non "step": 0, "thread_id": "1", }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, + parent_config=None, interrupts=(), ) @@ -1019,9 +1021,7 @@ async def test_conditional_graph(async_checkpointer: BaseCheckpointSaver) -> Non "step": 0, "thread_id": "2", }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, + parent_config=None, interrupts=(), ) @@ -1235,9 +1235,7 @@ async def test_conditional_graph(async_checkpointer: BaseCheckpointSaver) -> Non "step": 0, "thread_id": "3", }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, + parent_config=None, interrupts=(), ) @@ -1607,9 +1605,7 @@ async def test_conditional_graph_state(async_checkpointer: BaseCheckpointSaver) "step": 1, "thread_id": "1", }, - parent_config=[ - c async for c in app_w_interrupt.checkpointer.alist(config, limit=2) - ][-1].config, + parent_config=None, interrupts=(), ) @@ -1785,11 +1781,7 @@ async def test_conditional_graph_state(async_checkpointer: BaseCheckpointSaver) "step": 1, "thread_id": "2", }, - parent_config=( - [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ - -1 - ].config - ), + parent_config=None, interrupts=(), ) @@ -2524,11 +2516,7 @@ async def test_state_graph_packets(async_checkpointer: BaseCheckpointSaver) -> N "step": 1, "thread_id": "1", }, - parent_config=( - [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ - -1 - ].config - ), + parent_config=None, interrupts=(), ) @@ -2775,11 +2763,7 @@ async def test_state_graph_packets(async_checkpointer: BaseCheckpointSaver) -> N "step": 1, "thread_id": "2", }, - parent_config=( - [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ - -1 - ].config - ), + parent_config=None, interrupts=(), ) @@ -3207,11 +3191,7 @@ async def test_message_graph(async_checkpointer: BaseCheckpointSaver) -> None: "step": 1, "thread_id": "1", }, - parent_config=( - [c async for c in app_w_interrupt.checkpointer.alist(config, limit=2)][ - -1 - ].config - ), + parent_config=None, interrupts=(), ) @@ -3658,13 +3638,6 @@ async def test_start_branch_then( "assistant_id": "a", "thread_id": "1", }, - { - "parents": {}, - "source": "input", - "step": -1, - "assistant_id": "a", - "thread_id": "1", - }, ] assert await tool_two.aget_state(thread1) == StateSnapshot( @@ -3686,9 +3659,7 @@ async def test_start_branch_then( "assistant_id": "a", "thread_id": "1", }, - parent_config=[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][ - -1 - ].config, + parent_config=None, interrupts=(), ) # resume, for same result as above @@ -3746,9 +3717,7 @@ async def test_start_branch_then( "assistant_id": "a", "thread_id": "2", }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][-1].config - ), + parent_config=None, interrupts=(), ) # resume, for same result as above @@ -3806,9 +3775,7 @@ async def test_start_branch_then( "assistant_id": "b", "thread_id": "3", }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread3, limit=2)][-1].config - ), + parent_config=None, interrupts=(), ) # update state @@ -3902,7 +3869,10 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: assert [ c async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread10, stream_mode="debug" + {"my_key": "value", "market": "DE"}, + thread10, + stream_mode="debug", + checkpoint_during=True, ) ] == [ { @@ -3911,10 +3881,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "step": -1, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -3946,10 +3912,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "step": 0, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -3967,10 +3929,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "thread_id": "10", }, "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -4017,10 +3975,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "step": 1, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -4038,10 +3992,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "thread_id": "10", }, "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -4088,10 +4038,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "step": 2, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -4109,10 +4055,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "thread_id": "10", }, "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -4162,10 +4104,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "step": 3, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -4183,10 +4121,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "thread_id": "10", }, "parent_config": { - "tags": [], - "metadata": {"thread_id": "10"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "10", "checkpoint_ns": "", @@ -4213,7 +4147,10 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: assert [ c async for c in tool_two.astream( - {"my_key": "value", "market": "DE"}, thread1, stream_mode="debug" + {"my_key": "value", "market": "DE"}, + thread1, + stream_mode="debug", + checkpoint_during=True, ) ] == [ { @@ -4222,10 +4159,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "step": -1, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "11", "checkpoint_ns": "", @@ -4257,10 +4190,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "step": 0, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "11", "checkpoint_ns": "", @@ -4278,10 +4207,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "thread_id": "11", }, "parent_config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "11", "checkpoint_ns": "", @@ -4328,10 +4253,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "step": 1, "payload": { "config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "11", "checkpoint_ns": "", @@ -4349,10 +4270,6 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "thread_id": "11", }, "parent_config": { - "tags": [], - "metadata": {"thread_id": "11"}, - "callbacks": None, - "recursion_limit": 25, "configurable": { "thread_id": "11", "checkpoint_ns": "", @@ -4447,9 +4364,7 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "step": 1, "thread_id": "12", }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][-1].config - ), + parent_config=None, interrupts=(), ) # resume, for same result as above @@ -4513,9 +4428,7 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "step": 1, "thread_id": "21", }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config - ), + parent_config=None, interrupts=(), ) # resume, for same result as above @@ -4571,9 +4484,7 @@ async def test_branch_then(async_checkpointer: BaseCheckpointSaver) -> None: "step": 1, "thread_id": "22", }, - parent_config=( - [c async for c in tool_two.checkpointer.alist(thread2, limit=2)][-1].config - ), + parent_config=None, interrupts=(), ) # resume, for same result as above @@ -4733,7 +4644,7 @@ async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> No await app.ainvoke({"my_key": "my value"}, config, debug=True) # test state w/ nested subgraph state (right after interrupt) # first get_state without subgraph state - assert await app.aget_state(config) == StateSnapshot( + expected = StateSnapshot( values={"my_key": "hi my value"}, tasks=( PregelTask( @@ -4758,17 +4669,10 @@ async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> No "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, interrupts=(), ) + assert await app.aget_state(config) == expected # now, get_state with subgraphs state assert await app.aget_state(config, subgraphs=True) == StateSnapshot( values={"my_key": "hi my value"}, @@ -4814,16 +4718,7 @@ async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> No "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, + parent_config=None, interrupts=(), ), ), @@ -4843,129 +4738,19 @@ async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> No "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, interrupts=(), ) # get_state_history returns outer graph checkpoints - history = [c async for c in app.aget_state_history(config)] - expected_history = [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - (PULL, "inner"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - } - }, - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), - interrupts=(), - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), - ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - interrupts=(), - ), - ] - - assert history == expected_history + assert [c async for c in app.aget_state_history(config)] == [expected] # get_state_history for a subgraph returns its checkpoints - child_history = [c async for c in app.aget_state_history(history[0].tasks[0].state)] + child_history = [ + c + async for c in app.aget_state_history( + (await app.aget_state(config)).tasks[0].state + ) + ] expected_child_history = [ StateSnapshot( values={"my_key": "hi my value here", "my_other_key": "hi my value"}, @@ -4992,103 +4777,8 @@ async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> No "langgraph_checkpoint_ns": AnyStr("inner:"), }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - } - ), - tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), - interrupts=(), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("inner_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "loop", - "step": 0, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - tasks=( - PregelTask( - AnyStr(), - "inner_1", - (PULL, "inner_1"), - result={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - ), - ), - interrupts=(), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("inner:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "inner", - "langgraph_path": [PULL, "inner"], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:inner"], - "langgraph_checkpoint_ns": AnyStr("inner:"), - }, - created_at=AnyStr(), parent_config=None, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), + tasks=(PregelTask(AnyStr(), "inner_2", (PULL, "inner_2")),), interrupts=(), ), ] @@ -5159,40 +4849,6 @@ async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> No ), interrupts=(), ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=( - PregelTask( - AnyStr(), - "outer_2", - (PULL, "outer_2"), - result={"my_key": "hi my value here and there and back again"}, - ), - ), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 2, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), StateSnapshot( values={"my_key": "hi my value"}, tasks=( @@ -5206,7 +4862,7 @@ async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> No "checkpoint_ns": AnyStr(), } }, - result={"my_key": "hi my value here and there"}, + result=None, ), ), next=("inner",), @@ -5224,74 +4880,6 @@ async def test_nested_graph_state(async_checkpointer: BaseCheckpointSaver) -> No "thread_id": "1", }, created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=( - PregelTask( - AnyStr(), - "outer_1", - (PULL, "outer_1"), - result={"my_key": "hi my value"}, - ), - ), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={}, - tasks=( - PregelTask( - AnyStr(), - "__start__", - (PULL, "__start__"), - result={"my_key": "my value"}, - ), - ), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), parent_config=None, interrupts=(), ), @@ -5399,15 +4987,7 @@ async def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, interrupts=(), ) child_state = await app.aget_state(outer_state.tasks[0].state) @@ -5452,21 +5032,7 @@ async def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - } - } - ), + parent_config=None, interrupts=(), ) grandchild_state = await app.aget_state(child_state.tasks[0].state) @@ -5513,22 +5079,7 @@ async def test_doubly_nested_graph_state( ], }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - } - ), + parent_config=None, interrupts=(), ) # get state with subgraphs @@ -5594,24 +5145,7 @@ async def test_doubly_nested_graph_state( ], }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr( - re.compile(r"child:.+|child1:") - ): AnyStr(), - } - ), - } - } - ), + parent_config=None, interrupts=(), ), ), @@ -5641,18 +5175,7 @@ async def test_doubly_nested_graph_state( "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - } - ), + parent_config=None, interrupts=(), ), ), @@ -5672,15 +5195,7 @@ async def test_doubly_nested_graph_state( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, interrupts=(), ) # resume @@ -5733,161 +5248,71 @@ async def test_doubly_nested_graph_state( # get outer graph history outer_history = [c async for c in app.aget_state_history(config)] - assert ( - outer_history[0] - == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 3, + assert outer_history == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("parent_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 2, + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 3, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask(id=AnyStr(), name="parent_2", path=(PULL, "parent_2")), + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + interrupts=(), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + (PULL, "child"), + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child"), + } + }, ), - interrupts=(), ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - (PULL, "child"), - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child"), - } - }, - ), - ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 1, + next=("child",), + config={ + "configurable": { "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - interrupts=(), - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("parent_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "step": 0, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=( - PregelTask(id=AnyStr(), name="parent_1", path=(PULL, "parent_1")), - ), - interrupts=(), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "step": -1, - "thread_id": "1", - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")), - ), - interrupts=(), - ), - ][0] - ) + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "step": 1, + "thread_id": "1", + }, + created_at=AnyStr(), + parent_config=None, + interrupts=(), + ), + ] # get child graph history child_history = [ - c async for c in app.aget_state_history(outer_history[2].tasks[0].state) + c async for c in app.aget_state_history(outer_history[1].tasks[0].state) ] assert child_history == [ StateSnapshot( @@ -5953,16 +5378,7 @@ async def test_doubly_nested_graph_state( "langgraph_checkpoint_ns": AnyStr("child:"), }, created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, + parent_config=None, tasks=( PregelTask( id=AnyStr(), @@ -5974,43 +5390,7 @@ async def test_doubly_nested_graph_state( "checkpoint_ns": AnyStr("child:"), } }, - result={"my_key": "hi my value here and there"}, - ), - ), - interrupts=(), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - {"": AnyStr(), AnyStr("child:"): AnyStr()} - ), - } - }, - metadata={ - "source": "input", - "step": -1, - "parents": {"": AnyStr()}, - "thread_id": "1", - "langgraph_node": "child", - "langgraph_path": [PULL, AnyStr("child")], - "langgraph_step": 2, - "langgraph_triggers": ["branch:to:child"], - "langgraph_checkpoint_ns": AnyStr("child:"), - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, + result=None, ), ), interrupts=(), @@ -6116,154 +5496,17 @@ async def test_doubly_nested_graph_state( ], }, created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, + parent_config=None, tasks=( PregelTask( id=AnyStr(), name="grandchild_2", path=(PULL, "grandchild_2"), - result={"my_key": "hi my value here and there"}, + result=None, ), ), interrupts=(), ), - StateSnapshot( - values={"my_key": "hi my value"}, - next=("grandchild_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "loop", - "step": 0, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [ - "branch:to:child_1", - ], - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - tasks=( - PregelTask( - id=AnyStr(), - name="grandchild_1", - path=(PULL, "grandchild_1"), - result={"my_key": "hi my value here"}, - ), - ), - interrupts=(), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), - } - ), - } - }, - metadata={ - "source": "input", - "step": -1, - "parents": AnyDict( - { - "": AnyStr(), - AnyStr("child:"): AnyStr(), - } - ), - "thread_id": "1", - "langgraph_checkpoint_ns": AnyStr("child:"), - "langgraph_node": "child_1", - "langgraph_path": [ - PULL, - AnyStr("child_1"), - ], - "langgraph_step": 1, - "langgraph_triggers": [ - "branch:to:child_1", - ], - }, - created_at=AnyStr(), - parent_config=None, - tasks=( - PregelTask( - id=AnyStr(), - name="__start__", - path=(PULL, "__start__"), - result={"my_key": "hi my value"}, - ), - ), - interrupts=(), - ), - ] - - # replay grandchild checkpoint - assert [ - c async for c in app.astream(None, grandchild_history[2].config, subgraphs=True) - ] == [ - ( - (AnyStr("child:"), AnyStr("child_1:")), - {"grandchild_1": {"my_key": "hi my value here"}}, - ), - ((), {"__interrupt__": ()}), ] @@ -6504,15 +5747,7 @@ async def test_weather_subgraph( "thread_id": "1", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, tasks=( PregelTask( id=AnyStr(), @@ -6596,15 +5831,7 @@ async def test_weather_subgraph( "thread_id": "14", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, tasks=( PregelTask( id=AnyStr(), @@ -6643,21 +5870,7 @@ async def test_weather_subgraph( "langgraph_checkpoint_ns": AnyStr("weather_graph:"), }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "14", - "checkpoint_ns": AnyStr("weather_graph:"), - "checkpoint_id": AnyStr(), - "checkpoint_map": AnyDict( - { - "": AnyStr(), - AnyStr("weather_graph:"): AnyStr(), - } - ), - } - } - ), + parent_config=None, interrupts=(), tasks=( PregelTask( @@ -6697,15 +5910,7 @@ async def test_weather_subgraph( "thread_id": "14", }, created_at=AnyStr(), - parent_config=( - { - "configurable": { - "thread_id": "14", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - } - ), + parent_config=None, interrupts=(), tasks=( PregelTask( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 60bb05385..2c3c15a4c 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b9712f4ab..de357e162 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -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():