diff --git a/libs/langgraph/langgraph/channels/manager.py b/libs/langgraph/langgraph/channels/manager.py index c5fa373f3..8bee9e545 100644 --- a/libs/langgraph/langgraph/channels/manager.py +++ b/libs/langgraph/langgraph/channels/manager.py @@ -44,19 +44,22 @@ async def AsyncChannelsManager( def create_checkpoint( checkpoint: Checkpoint, - channels: Mapping[str, BaseChannel], + channels: Optional[Mapping[str, BaseChannel]], step: int, *, id: Optional[str] = None, ) -> Checkpoint: """Create a checkpoint for the given channels.""" ts = datetime.now(timezone.utc).isoformat() - values: dict[str, Any] = {} - for k, v in channels.items(): - try: - values[k] = v.checkpoint() - except EmptyChannelError: - pass + if channels is None: + values = checkpoint["channel_values"] + else: + values: dict[str, Any] = {} + for k, v in channels.items(): + try: + values[k] = v.checkpoint() + except EmptyChannelError: + pass return Checkpoint( v=1, ts=ts, diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index c990b4dc6..52ed515ab 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -504,7 +504,7 @@ class Pregel( def update_state( self, config: RunnableConfig, - values: dict[str, Any] | Any, + values: Optional[Union[dict[str, Any], Any]], as_node: Optional[str] = None, ) -> RunnableConfig: """Update the state of the graph with the given values, as if they came from @@ -517,8 +517,28 @@ class Pregel( # get last checkpoint saved = self.checkpointer.get_tuple(config) checkpoint = copy_checkpoint(saved.checkpoint) if saved else empty_checkpoint() + step = saved.metadata.get("step", -1) if saved else -1 + # merge configurable fields with previous checkpoint config + checkpoint_config = config + if saved: + checkpoint_config = { + "configurable": { + **config.get("configurable", {}), + **saved.config["configurable"], + } + } # find last node that updated the state, if not provided - if as_node is None and not any( + if values is None and as_node is None: + return self.checkpointer.put( + checkpoint_config, + create_checkpoint(checkpoint, None, step), + { + "source": "update", + "step": step, + "writes": {}, + }, + ) + elif as_node is None and not any( v for vv in checkpoint["versions_seen"].values() for v in vv.values() ): if ( @@ -577,24 +597,13 @@ class Pregel( apply_writes( checkpoint, channels, [task], self.checkpointer.get_next_version ) - step = saved.metadata.get("step", -2) + 1 if saved else -1 - - # merge configurable fields with previous checkpoint config - checkpoint_config = config - if saved: - checkpoint_config = { - "configurable": { - **config.get("configurable", {}), - **saved.config["configurable"], - } - } return self.checkpointer.put( checkpoint_config, - create_checkpoint(checkpoint, channels, step), + create_checkpoint(checkpoint, channels, step + 1), { "source": "update", - "step": step, + "step": step + 1, "writes": {as_node: values}, }, ) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 04a71bfa3..3179194dc 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -549,172 +549,380 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + SqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +def test_invoke_two_processes_in_out_interrupt( + checkpointer: BaseCheckpointSaver, mocker: MockerFixture +) -> None: + try: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - memory = MemorySaverAssertImmutable() - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=memory, - interrupt_after_nodes=["one"], - ) - - # start execution, stop at inbox - assert app.invoke(2, {"configurable": {"thread_id": 1}}) is None - - # inbox == 3 - checkpoint = memory.get({"configurable": {"thread_id": 1}}) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 3 - - # resume execution, finish - assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 4 - - # start execution again, stop at inbox - assert app.invoke(20, {"configurable": {"thread_id": 1}}) is None - - # inbox == 21 - checkpoint = memory.get({"configurable": {"thread_id": 1}}) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 21 - - # send a new value in, interrupting the previous execution - assert app.invoke(3, {"configurable": {"thread_id": 1}}) is None - assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 5 - - # start execution again, stopping at inbox - assert app.invoke(20, {"configurable": {"thread_id": 2}}) is None - - # inbox == 21 - snapshot = app.get_state({"configurable": {"thread_id": 2}}) - assert snapshot.values["inbox"] == 21 - assert snapshot.next == ("two",) - - # update the state, resume - app.update_state({"configurable": {"thread_id": 2}}, 25, as_node="one") - assert app.invoke(None, {"configurable": {"thread_id": 2}}) == 26 - - # no pending tasks - snapshot = app.get_state({"configurable": {"thread_id": 2}}) - assert snapshot.next == () - - # list history - thread1 = {"configurable": {"thread_id": 1}} - assert [c for c in app.get_state_history(thread1)] == [ - StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, - next=(), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), }, - metadata={"source": "loop", "step": 6, "writes": 5}, - created_at=AnyStr(), - parent_config=[*app.checkpointer.list(thread1)][1].config, - ), - StateSnapshot( - values={"inbox": 4, "output": 4, "input": 3}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "loop", "step": 5, "writes": None}, - created_at=AnyStr(), - parent_config=[*app.checkpointer.list(thread1)][2].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "input", "step": 4, "writes": 3}, - created_at=AnyStr(), - parent_config=[*app.checkpointer.list(thread1)][3].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "loop", "step": 3, "writes": None}, - created_at=AnyStr(), - parent_config=[*app.checkpointer.list(thread1)][4].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "input", "step": 2, "writes": 20}, - created_at=AnyStr(), - parent_config=[*app.checkpointer.list(thread1)][5].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, - next=(), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "loop", "step": 1, "writes": 4}, - created_at=AnyStr(), - parent_config=[*app.checkpointer.list(thread1)][6].config, - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "loop", "step": 0, "writes": None}, - created_at=AnyStr(), - parent_config=[*app.checkpointer.list(thread1)][7].config, - ), - StateSnapshot( - values={"input": 2}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "input", "step": -1, "writes": 2}, - created_at=AnyStr(), - parent_config=None, - ), - ] + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + interrupt_after_nodes=["one"], + ) + thread1 = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} + + # start execution, stop at inbox + assert app.invoke(2, thread1) is None + + # inbox == 3 + checkpoint = checkpointer.get(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 + + # resume execution, finish + assert app.invoke(None, thread1) == 4 + + # start execution again, stop at inbox + assert app.invoke(20, thread1) is None + + # inbox == 21 + checkpoint = checkpointer.get(thread1) + assert checkpoint is not None + 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 + + # start execution again, stopping at inbox + assert app.invoke(20, thread2) is None + + # inbox == 21 + snapshot = app.get_state(thread2) + assert snapshot.values["inbox"] == 21 + assert snapshot.next == ("two",) + + # update the state, resume + app.update_state(thread2, 25, as_node="one") + assert app.invoke(None, thread2) == 26 + + # no pending tasks + snapshot = app.get_state(thread2) + assert snapshot.next == () + + # list history + history = [c for c in app.get_state_history(thread1)] + assert history == [ + StateSnapshot( + values={"inbox": 4, "output": 5, "input": 3}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 6, "writes": 5}, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values={"inbox": 4, "output": 4, "input": 3}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 5, "writes": None}, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 3}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 4, "writes": 3}, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": None}, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 2, "writes": 20}, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 1, "writes": 4}, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=history[7].config, + ), + StateSnapshot( + values={"input": 2}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 2}, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # forking from any previous checkpoint w/out forking should do nothing + assert [ + c for c in app.stream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c for c in app.stream(None, history[1].config, stream_mode="updates") + ] == [] + assert [ + c for c in app.stream(None, history[2].config, stream_mode="updates") + ] == [] + + # forking and re-running from any prev checkpoint should re-run nodes + fork_config = app.update_state(history[0].config, None) + assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [] + + fork_config = app.update_state(history[1].config, None) + assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [ + {"two": {"output": 5}} + ] + + fork_config = app.update_state(history[2].config, None) + assert [c for c in app.stream(None, fork_config, stream_mode="updates")] == [ + {"one": {"inbox": 4}} + ] + finally: + if hasattr(checkpointer, "__exit__"): + checkpointer.__exit__(None, None, None) + + +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + SqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +def test_fork_always_re_runs_nodes( + checkpointer: BaseCheckpointSaver, mocker: MockerFixture +) -> None: + try: + add_one = mocker.Mock(side_effect=lambda _: 1) + + builder = StateGraph(Annotated[int, operator.add]) + builder.add_node("add_one", add_one) + builder.add_edge(START, "add_one") + builder.add_conditional_edges( + "add_one", lambda cnt: "add_one" if cnt < 6 else END + ) + graph = builder.compile(checkpointer=checkpointer) + + thread1 = {"configurable": {"thread_id": "1"}} + + # start execution, stop at inbox + assert [*graph.stream(1, thread1, stream_mode=["values", "updates"])] == [ + ("values", 1), + ("updates", {"add_one": 1}), + ("values", 2), + ("updates", {"add_one": 1}), + ("values", 3), + ("updates", {"add_one": 1}), + ("values", 4), + ("updates", {"add_one": 1}), + ("values", 5), + ("updates", {"add_one": 1}), + ("values", 6), + ] + + # list history + history = [c for c in graph.get_state_history(thread1)] + assert history == [ + StateSnapshot( + values=6, + next=(), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 5, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values=5, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 4, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values=4, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values=3, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 2, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values=2, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 1, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values=1, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values=0, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 1}, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # forking from any previous checkpoint w/out forking should do nothing + assert [ + c for c in graph.stream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c for c in graph.stream(None, history[1].config, stream_mode="updates") + ] == [] + + # forking and re-running from any prev checkpoint should re-run nodes + fork_config = graph.update_state(history[0].config, None) + assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [] + + fork_config = graph.update_state(history[1].config, None) + assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [ + {"add_one": 1} + ] + + fork_config = graph.update_state(history[2].config, None) + assert [c for c in graph.stream(None, fork_config, stream_mode="updates")] == [ + {"add_one": 1}, + {"add_one": 1}, + ] + finally: + if hasattr(checkpointer, "__exit__"): + checkpointer.__exit__(None, None, None) def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: @@ -6798,7 +7006,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: created_at=AnyStr(), metadata={ "source": "update", - "step": -1, + "step": 0, "writes": {START: {"my_key": "key", "market": "DE"}}, }, ) @@ -6815,7 +7023,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], metadata={ "source": "loop", - "step": 0, + "step": 1, "writes": {"prepare": {"my_key": " prepared"}}, }, parent_config=uconfig, @@ -6832,7 +7040,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], metadata={ "source": "loop", - "step": 2, + "step": 3, "writes": {"finish": {"my_key": " finished"}}, }, parent_config=[*tool_two.checkpointer.list(thread3, limit=2)][-1].config, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 0aeb1546a..598d2902e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -665,172 +665,364 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + AsyncSqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +async def test_invoke_two_processes_in_out_interrupt( + checkpointer: BaseCheckpointSaver, mocker: MockerFixture +) -> None: + try: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - memory = MemorySaverAssertImmutable() - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - checkpointer=memory, - interrupt_after_nodes=["one"], - ) - - # start execution, stop at inbox - assert await app.ainvoke(2, {"configurable": {"thread_id": 1}}) is None - - # inbox == 3 - checkpoint = await memory.aget({"configurable": {"thread_id": 1}}) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 3 - - # resume execution, finish - assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 4 - - # start execution again, stop at inbox - assert await app.ainvoke(20, {"configurable": {"thread_id": 1}}) is None - - # inbox == 21 - checkpoint = await memory.aget({"configurable": {"thread_id": 1}}) - assert checkpoint is not None - assert checkpoint["channel_values"]["inbox"] == 21 - - # send a new value in, interrupting the previous execution - assert await app.ainvoke(3, {"configurable": {"thread_id": 1}}) is None - assert await app.ainvoke(None, {"configurable": {"thread_id": 1}}) == 5 - - # start execution again, stopping at inbox - assert await app.ainvoke(20, {"configurable": {"thread_id": 2}}) is None - - # inbox == 21 - snapshot = await app.aget_state({"configurable": {"thread_id": 2}}) - assert snapshot.values["inbox"] == 21 - assert snapshot.next == ("two",) - - # update the state, resume - await app.aupdate_state({"configurable": {"thread_id": 2}}, 25, as_node="one") - assert await app.ainvoke(None, {"configurable": {"thread_id": 2}}) == 26 - - # no pending tasks - snapshot = await app.aget_state({"configurable": {"thread_id": 2}}) - assert snapshot.next == () - - # list history - thread1 = {"configurable": {"thread_id": 1}} - assert [c async for c in app.aget_state_history(thread1)] == [ - StateSnapshot( - values={"inbox": 4, "output": 5, "input": 3}, - next=(), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), }, - metadata={"source": "loop", "step": 6, "writes": 5}, - created_at=AnyStr(), - parent_config=[c async for c in app.checkpointer.alist(thread1)][1].config, - ), - StateSnapshot( - values={"inbox": 4, "output": 4, "input": 3}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "loop", "step": 5, "writes": None}, - created_at=AnyStr(), - parent_config=[c async for c in app.checkpointer.alist(thread1)][2].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 3}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "input", "step": 4, "writes": 3}, - created_at=AnyStr(), - parent_config=[c async for c in app.checkpointer.alist(thread1)][3].config, - ), - StateSnapshot( - values={"inbox": 21, "output": 4, "input": 20}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "loop", "step": 3, "writes": None}, - created_at=AnyStr(), - parent_config=[c async for c in app.checkpointer.alist(thread1)][4].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 20}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "input", "step": 2, "writes": 20}, - created_at=AnyStr(), - parent_config=[c async for c in app.checkpointer.alist(thread1)][5].config, - ), - StateSnapshot( - values={"inbox": 3, "output": 4, "input": 2}, - next=(), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "loop", "step": 1, "writes": 4}, - created_at=AnyStr(), - parent_config=[c async for c in app.checkpointer.alist(thread1)][6].config, - ), - StateSnapshot( - values={"inbox": 3, "input": 2}, - next=("two",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "loop", "step": 0, "writes": None}, - created_at=AnyStr(), - parent_config=[c async for c in app.checkpointer.alist(thread1)][7].config, - ), - StateSnapshot( - values={"input": 2}, - next=("one",), - config={ - "configurable": { - "thread_id": 1, - "thread_ts": AnyStr(), - } - }, - metadata={"source": "input", "step": -1, "writes": 2}, - created_at=AnyStr(), - parent_config=None, - ), - ] + input_channels="input", + output_channels="output", + checkpointer=checkpointer, + interrupt_after_nodes=["one"], + ) + thread1 = {"configurable": {"thread_id": "1"}} + thread2 = {"configurable": {"thread_id": "2"}} + + # start execution, stop at inbox + assert await app.ainvoke(2, thread1) is None + + # inbox == 3 + checkpoint = await checkpointer.aget(thread1) + assert checkpoint is not None + assert checkpoint["channel_values"]["inbox"] == 3 + + # resume execution, finish + assert await app.ainvoke(None, thread1) == 4 + + # start execution again, stop at inbox + assert await app.ainvoke(20, thread1) is None + + # inbox == 21 + checkpoint = await checkpointer.aget(thread1) + assert checkpoint is not None + 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 + + # start execution again, stopping at inbox + assert await app.ainvoke(20, thread2) is None + + # inbox == 21 + snapshot = await app.aget_state(thread2) + assert snapshot.values["inbox"] == 21 + assert snapshot.next == ("two",) + + # update the state, resume + await app.aupdate_state(thread2, 25, as_node="one") + assert await app.ainvoke(None, thread2) == 26 + + # no pending tasks + snapshot = await app.aget_state(thread2) + assert snapshot.next == () + + # list history + history = [c async for c in app.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values={"inbox": 4, "output": 5, "input": 3}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 6, "writes": 5}, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values={"inbox": 4, "output": 4, "input": 3}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 5, "writes": None}, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 3}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 4, "writes": 3}, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values={"inbox": 21, "output": 4, "input": 20}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": None}, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 20}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": 2, "writes": 20}, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values={"inbox": 3, "output": 4, "input": 2}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 1, "writes": 4}, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values={"inbox": 3, "input": 2}, + next=("two",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=history[7].config, + ), + StateSnapshot( + values={"input": 2}, + next=("one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 2}, + created_at=AnyStr(), + parent_config=None, + ), + ] + finally: + if hasattr(checkpointer, "__aexit__"): + await checkpointer.__aexit__(None, None, None) + + +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + AsyncSqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +async def test_fork_always_re_runs_nodes( + checkpointer: BaseCheckpointSaver, mocker: MockerFixture +) -> None: + try: + add_one = mocker.Mock(side_effect=lambda _: 1) + + builder = StateGraph(Annotated[int, operator.add]) + builder.add_node("add_one", add_one) + builder.add_edge(START, "add_one") + builder.add_conditional_edges( + "add_one", lambda cnt: "add_one" if cnt < 6 else END + ) + graph = builder.compile(checkpointer=checkpointer) + + thread1 = {"configurable": {"thread_id": "1"}} + + # start execution, stop at inbox + assert [ + c + async for c in graph.astream(1, thread1, stream_mode=["values", "updates"]) + ] == [ + ("values", 1), + ("updates", {"add_one": 1}), + ("values", 2), + ("updates", {"add_one": 1}), + ("values", 3), + ("updates", {"add_one": 1}), + ("values", 4), + ("updates", {"add_one": 1}), + ("values", 5), + ("updates", {"add_one": 1}), + ("values", 6), + ] + + # list history + history = [c async for c in graph.aget_state_history(thread1)] + assert history == [ + StateSnapshot( + values=6, + next=(), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 5, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[1].config, + ), + StateSnapshot( + values=5, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 4, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[2].config, + ), + StateSnapshot( + values=4, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 3, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[3].config, + ), + StateSnapshot( + values=3, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 2, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[4].config, + ), + StateSnapshot( + values=2, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 1, "writes": {"add_one": 1}}, + created_at=AnyStr(), + parent_config=history[5].config, + ), + StateSnapshot( + values=1, + next=("add_one",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "step": 0, "writes": None}, + created_at=AnyStr(), + parent_config=history[6].config, + ), + StateSnapshot( + values=0, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "input", "step": -1, "writes": 1}, + created_at=AnyStr(), + parent_config=None, + ), + ] + + # forking from any previous checkpoint w/out forking should do nothing + assert [ + c + async for c in graph.astream(None, history[0].config, stream_mode="updates") + ] == [] + assert [ + c + async for c in graph.astream(None, history[1].config, stream_mode="updates") + ] == [] + + # forking and re-running from any prev checkpoint should re-run nodes + fork_config = await graph.aupdate_state(history[0].config, None) + assert [ + c async for c in graph.astream(None, fork_config, stream_mode="updates") + ] == [] + + fork_config = await graph.aupdate_state(history[1].config, None) + assert [ + c async for c in graph.astream(None, fork_config, stream_mode="updates") + ] == [{"add_one": 1}] + + fork_config = await graph.aupdate_state(history[2].config, None) + assert [ + c async for c in graph.astream(None, fork_config, stream_mode="updates") + ] == [ + {"add_one": 1}, + {"add_one": 1}, + ] + finally: + if hasattr(checkpointer, "__aexit__"): + await checkpointer.__aexit__(None, None, None) async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: