diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 9c39ec2b1..51e9690c8 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6806,1754 +6806,6 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert times_called == 1 -@pytest.mark.skip("TODO") -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_interrupts( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - - class InnerState(TypedDict): - my_key: str - my_other_key: str - - async def inner_1(state: InnerState): - return { - "my_key": state["my_key"] + " here", - "my_other_key": state["my_key"], - } - - async def inner_2(state: InnerState): - return { - "my_key": state["my_key"] + " and there", - "my_other_key": state["my_key"], - } - - inner = StateGraph(InnerState) - inner.add_node("inner_1", inner_1) - inner.add_node("inner_2", inner_2) - inner.add_edge("inner_1", "inner_2") - inner.set_entry_point("inner_1") - inner.set_finish_point("inner_2") - - class State(TypedDict): - my_key: str - - async def outer_1(state: State): - return {"my_key": "hi " + state["my_key"]} - - async def outer_2(state: State): - return {"my_key": state["my_key"] + " and back again"} - - graph = StateGraph(State) - graph.add_node("outer_1", outer_1) - graph.add_node( - "inner", - inner.compile(interrupt_before=["inner_2"]), - ) - graph.add_node("outer_2", outer_2) - graph.set_entry_point("outer_1") - graph.add_edge("outer_1", "inner") - graph.add_edge("inner", "outer_2") - graph.set_finish_point("outer_2") - - app = graph.compile(checkpointer=checkpointer) - - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { - "my_key": "hi my value", - } - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - "inner_2", - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert await app.ainvoke(None, config, debug=True) == { - "my_key": "hi my value here and there and back again", - } - assert [s async for s in app.aget_state_history(config)] == [ - 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", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - next=(), - tasks=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ - {"outer_1": {"my_key": "hi my value"}}, - ] - assert [c async for c in app.astream(None, config)] == [ - {"inner": {"my_key": "hi my value here and there"}}, - {"outer_2": {"my_key": "hi my value here and there and back again"}}, - ] - - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [ - c - async for c in app.astream({"my_key": "my value"}, config, stream_mode="values") - ] == [ - { - "my_key": "my value", - }, - { - "my_key": "hi my value", - }, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", - }, - { - "my_key": "hi my value here and there and back again", - }, - ] - - # test interrupts BEFORE the node w/ interrupts - app = graph.compile(checkpointer=checkpointer, interrupt_before=["inner"]) - config = {"configurable": {"thread_id": "4"}} - assert [ - c - async for c in app.astream({"my_key": "my value"}, config, stream_mode="values") - ] == [ - { - "my_key": "my value", - }, - { - "my_key": "hi my value", - }, - ] - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # while we're waiting for the node w/ interrupt inside to finish - assert [c async for c in app.astream(None, config, stream_mode="values")] == [] - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - "inner_2", - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", - }, - { - "my_key": "hi my value here and there and back again", - }, - ] - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # test interrupts AFTER the node w/ interrupts - app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"]) - config = {"configurable": {"thread_id": "5"}} - assert [ - c - async for c in app.astream({"my_key": "my value"}, config, stream_mode="values") - ] == [ - { - "my_key": "my value", - }, - { - "my_key": "hi my value", - }, - ] - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - name="inner_2", - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", - }, - ] - # interrupted after "inner" - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there and back again", - }, - ] - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # test restarting from checkpoint_id - config = {"configurable": {"thread_id": "6"}} - app = graph.compile(checkpointer=checkpointer) - await app.ainvoke({"my_key": "my value"}, config, debug=True) - - state_history = [c async for c in app.aget_state_history(config)] - assert state_history == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - "inner_2", - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - child_state_history = [ - c - async for c in app.aget_state_history( - {"configurable": {"thread_id": "6", "checkpoint_ns": "inner"}} - ) - ] - assert child_state_history == [ - StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ), - # there should be a single child checkpoint because we only keep - # one child checkpoint per parent checkpoint (in which child ran) - ] - - # check that child snapshot matches id of parent - child_snapshot = child_state_history[0] - assert ( - child_snapshot.config["configurable"]["checkpoint_id"] - == state_history[0].config["configurable"]["checkpoint_id"] - ) - # check resuming from interrupt w/ checkpoint_id - interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2] - before_interrupt_config = before_interrupt_state_snapshot.config - # going to get to interrupt again here, so the output is None - assert await app.ainvoke(None, before_interrupt_config, debug=True) == { - "my_key": "hi my value" - } - # one more "identical" snapshot than before, at top of list - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ), - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ), - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # going to restart from interrupt - interrupt_config = interrupt_state_snapshot.config - assert await app.ainvoke(None, interrupt_config, debug=True) == { - "my_key": "hi my value here and there and back again", - } - assert [s async for s in app.aget_state_history(config)] == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - -@pytest.mark.skip("TODO") @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_nested_graph_interrupts_parallel( request: pytest.FixtureRequest, checkpointer_name: str @@ -8620,11 +6872,13 @@ async def test_nested_graph_interrupts_parallel( # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) # test stream updates w/ nested interrupt config = {"configurable": {"thread_id": "2"}} - assert [c async for c in app.astream({"my_key": ""}, config)] == [ + assert [c async for c in app.astream({"my_key": ""}, config, subgraphs=True)] == [ # we got to parallel node first - {"outer_1": {"my_key": " and parallel"}}, + ((), {"outer_1": {"my_key": " and parallel"}}), + ((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}), ] assert [c async for c in app.astream(None, config)] == [ + {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, {"inner": {"my_key": "got here and there"}}, {"outer_2": {"my_key": " and back again"}}, ]