diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index f1d45c192..486b4cdf9 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -3683,10 +3683,14 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_subgraph_checkpoint_true( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) class InnerState(TypedDict): @@ -3718,7 +3722,12 @@ def test_subgraph_checkpoint_true( app = graph.compile(checkpointer=checkpointer) config = {"configurable": {"thread_id": "2"}} - assert [c for c in app.stream({"my_key": ""}, config, subgraphs=True)] == [ + assert [ + c + for c in app.stream( + {"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during + ) + ] == [ (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}), (("inner",), {"inner_2": {"my_key": " and there"}}), ((), {"inner": {"my_key": " got here and there"}}), @@ -3743,10 +3752,14 @@ def test_subgraph_checkpoint_true( ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_subgraph_checkpoint_true_interrupt( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) # Define subgraph @@ -3785,15 +3798,18 @@ def test_subgraph_checkpoint_true_interrupt( builder.add_edge(START, "node_1") builder.add_edge("node_1", "node_2") - checkpointer = MemorySaver() graph = builder.compile(checkpointer=checkpointer) config = {"configurable": {"thread_id": "1"}} - assert graph.invoke({"foo": "foo"}, config) == {"foo": "hi! foo"} + assert graph.invoke( + {"foo": "foo"}, config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foo"} assert graph.get_state(config, subgraphs=True).tasks[0].state.values == { "bar": "hi! foo" } - assert graph.invoke(Command(resume="baz"), config) == {"foo": "hi! foobaz"} + assert graph.invoke( + Command(resume="baz"), config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foobaz"} @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) @@ -3909,10 +3925,14 @@ def test_stream_buffering_single_node( ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_nested_graph_interrupts_parallel( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) class InnerState(TypedDict): @@ -3959,11 +3979,11 @@ def test_nested_graph_interrupts_parallel( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert app.invoke({"my_key": ""}, config, debug=True) == { + assert app.invoke({"my_key": ""}, config, checkpoint_during=checkpoint_during) == { "my_key": " and parallel", } - assert app.invoke(None, config, debug=True) == { + assert app.invoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "got here and there and parallel and back again", } @@ -3972,13 +3992,17 @@ 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 [*app.stream({"my_key": ""}, config, subgraphs=True)] == [ + assert [ + *app.stream( + {"my_key": ""}, config, subgraphs=True, checkpoint_during=checkpoint_during + ) + ] == [ # we got to parallel node first ((), {"outer_1": {"my_key": " and parallel"}}), ((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}), ((), {"__interrupt__": ()}), ] - assert [*app.stream(None, config)] == [ + assert [*app.stream(None, config, checkpoint_during=checkpoint_during)] == [ {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, {"inner": {"my_key": "got here and there"}}, {"outer_2": {"my_key": " and back again"}}, @@ -3986,11 +4010,22 @@ def test_nested_graph_interrupts_parallel( # test stream values w/ nested interrupt config = {"configurable": {"thread_id": "3"}} - assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -3999,15 +4034,28 @@ def test_nested_graph_interrupts_parallel( # test interrupts BEFORE the parallel node app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) config = {"configurable": {"thread_id": "4"}} - assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ - {"my_key": ""} - ] + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [{"my_key": ""}] # while we're waiting for the node w/ interrupt inside to finish - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -4016,24 +4064,43 @@ def test_nested_graph_interrupts_parallel( # test interrupts AFTER the parallel node app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) config = {"configurable": {"thread_id": "5"}} - assert [*app.stream({"my_key": ""}, config, stream_mode="values")] == [ + assert [ + *app.stream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_doubly_nested_graph_interrupts( - request: pytest.FixtureRequest, checkpointer_name: str + request: pytest.FixtureRequest, checkpointer_name: str, checkpoint_during: bool ) -> None: + if not checkpoint_during and "shallow" in checkpointer_name: + pytest.skip("Unsupported combo") + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) class State(TypedDict): @@ -4087,11 +4154,13 @@ def test_doubly_nested_graph_interrupts( # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert app.invoke({"my_key": "my value"}, config, debug=True) == { + assert app.invoke( + {"my_key": "my value"}, config, checkpoint_during=checkpoint_during + ) == { "my_key": "hi my value", } - assert app.invoke(None, config, debug=True) == { + assert app.invoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "hi my value here and there and back again", } @@ -4100,12 +4169,14 @@ def test_doubly_nested_graph_interrupts( config = { "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} } - assert [*app.stream({"my_key": "my value"}, config)] == [ + assert [ + *app.stream({"my_key": "my value"}, config, checkpoint_during=checkpoint_during) + ] == [ {"parent_1": {"my_key": "hi my value"}}, {"__interrupt__": ()}, ] assert nodes == ["parent_1", "grandchild_1"] - assert [*app.stream(None, config)] == [ + assert [*app.stream(None, config, checkpoint_during=checkpoint_during)] == [ {"child": {"my_key": "hi my value here and there"}}, {"parent_2": {"my_key": "hi my value here and there and back again"}}, ] @@ -4120,11 +4191,22 @@ def test_doubly_nested_graph_interrupts( # test stream values w/ nested interrupt config = {"configurable": {"thread_id": "3"}} - assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ + assert [ + *app.stream( + {"my_key": "my value"}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) + ] == [ {"my_key": "my value"}, {"my_key": "hi my value"}, ] - assert [*app.stream(None, config, stream_mode="values")] == [ + assert [ + *app.stream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "hi my value"}, {"my_key": "hi my value here and there"}, {"my_key": "hi my value here and there and back again"}, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 4cc115f6f..20e2b96b5 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -5433,6 +5433,131 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert times_called == 1 +@pytest.mark.parametrize("checkpoint_during", [True, False]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_subgraph_checkpoint_true( + checkpointer_name: str, checkpoint_during: bool +) -> None: + class InnerState(TypedDict): + my_key: Annotated[str, operator.add] + my_other_key: str + + def inner_1(state: InnerState): + return {"my_key": " got here", "my_other_key": state["my_key"]} + + def inner_2(state: InnerState): + return {"my_key": " and there"} + + 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 + + graph = StateGraph(State) + graph.add_node("inner", inner.compile(checkpointer=True)) + graph.add_edge(START, "inner") + graph.add_conditional_edges( + "inner", lambda s: "inner" if s["my_key"].count("there") < 2 else END + ) + + async with awith_checkpointer(checkpointer_name) as checkpointer: + app = graph.compile(checkpointer=checkpointer) + + config = {"configurable": {"thread_id": "2"}} + assert [ + c + async for c in app.astream( + {"my_key": ""}, + config, + subgraphs=True, + checkpoint_during=checkpoint_during, + ) + ] == [ + (("inner",), {"inner_1": {"my_key": " got here", "my_other_key": ""}}), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ((), {"inner": {"my_key": " got here and there"}}), + ( + ("inner",), + { + "inner_1": { + "my_key": " got here", + "my_other_key": " got here and there got here and there", + } + }, + ), + (("inner",), {"inner_2": {"my_key": " and there"}}), + ( + (), + { + "inner": { + "my_key": " got here and there got here and there got here and there" + } + }, + ), + ] + + +@pytest.mark.parametrize("checkpoint_during", [True, False]) +@pytest.mark.parametrize("checkpointer_name", REGULAR_CHECKPOINTERS_ASYNC) +async def test_subgraph_checkpoint_true_interrupt( + checkpointer_name: str, checkpoint_during: bool +) -> None: + # Define subgraph + class SubgraphState(TypedDict): + # note that none of these keys are shared with the parent graph state + bar: str + baz: str + + def subgraph_node_1(state: SubgraphState): + baz_value = interrupt("Provide baz value") + return {"baz": baz_value} + + def subgraph_node_2(state: SubgraphState): + return {"bar": state["bar"] + state["baz"]} + + subgraph_builder = StateGraph(SubgraphState) + subgraph_builder.add_node(subgraph_node_1) + subgraph_builder.add_node(subgraph_node_2) + subgraph_builder.add_edge(START, "subgraph_node_1") + subgraph_builder.add_edge("subgraph_node_1", "subgraph_node_2") + subgraph = subgraph_builder.compile(checkpointer=True) + + class ParentState(TypedDict): + foo: str + + def node_1(state: ParentState): + return {"foo": "hi! " + state["foo"]} + + async def node_2(state: ParentState): + response = await subgraph.ainvoke({"bar": state["foo"]}) + return {"foo": response["bar"]} + + builder = StateGraph(ParentState) + builder.add_node("node_1", node_1) + builder.add_node("node_2", node_2) + builder.add_edge(START, "node_1") + builder.add_edge("node_1", "node_2") + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} + + assert await graph.ainvoke( + {"foo": "foo"}, config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foo"} + assert (await graph.aget_state(config, subgraphs=True)).tasks[ + 0 + ].state.values == {"bar": "hi! foo"} + assert await graph.ainvoke( + Command(resume="baz"), config, checkpoint_during=checkpoint_during + ) == {"foo": "hi! foobaz"} + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_stream_subgraphs_during_execution(checkpointer_name: str) -> None: class InnerState(TypedDict): @@ -5541,8 +5666,11 @@ async def test_stream_buffering_single_node(checkpointer_name: str) -> None: ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: +async def test_nested_graph_interrupts_parallel( + checkpointer_name: str, checkpoint_during: bool +) -> None: class InnerState(TypedDict): my_key: Annotated[str, operator.add] my_other_key: str @@ -5591,11 +5719,13 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": ""}, config, debug=True) == { + assert await app.ainvoke( + {"my_key": ""}, config, checkpoint_during=checkpoint_during + ) == { "my_key": " and parallel", } - assert await app.ainvoke(None, config, debug=True) == { + assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "got here and there and parallel and back again", } @@ -5605,7 +5735,13 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # test stream updates w/ nested interrupt config = {"configurable": {"thread_id": "2"}} assert [ - c async for c in app.astream({"my_key": ""}, config, subgraphs=True) + c + async for c in app.astream( + {"my_key": ""}, + config, + subgraphs=True, + checkpoint_during=checkpoint_during, + ) ] == [ # we got to parallel node first ((), {"outer_1": {"my_key": " and parallel"}}), @@ -5615,7 +5751,12 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: ), ((), {"__interrupt__": ()}), ] - assert [c async for c in app.astream(None, config)] == [ + assert [ + c + async for c in app.astream( + None, config, checkpoint_during=checkpoint_during + ) + ] == [ {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, {"inner": {"my_key": "got here and there"}}, {"outer_2": {"my_key": " and back again"}}, @@ -5624,12 +5765,23 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: # test stream values w/ nested interrupt config = {"configurable": {"thread_id": "3"}} assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -5639,16 +5791,32 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: app = graph.compile(checkpointer=checkpointer, interrupt_before=["outer_1"]) config = {"configurable": {"thread_id": "4"}} assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) ] == [ {"my_key": ""}, ] # 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 [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, @@ -5658,23 +5826,42 @@ async def test_nested_graph_interrupts_parallel(checkpointer_name: str) -> None: app = graph.compile(checkpointer=checkpointer, interrupt_after=["outer_1"]) config = {"configurable": {"thread_id": "5"}} assert [ - c async for c in app.astream({"my_key": ""}, config, stream_mode="values") + c + async for c in app.astream( + {"my_key": ""}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, + ) ] == [ {"my_key": ""}, {"my_key": " and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": ""}, {"my_key": "got here and there and parallel"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "got here and there and parallel"}, {"my_key": "got here and there and parallel and back again"}, ] +@pytest.mark.parametrize("checkpoint_during", [True, False]) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: +async def test_doubly_nested_graph_interrupts( + checkpointer_name: str, checkpoint_during: bool +) -> None: class State(TypedDict): my_key: str @@ -5727,11 +5914,13 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: # test invoke w/ nested interrupt config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { + assert await app.ainvoke( + {"my_key": "my value"}, config, checkpoint_during=checkpoint_during + ) == { "my_key": "hi my value", } - assert await app.ainvoke(None, config, debug=True) == { + assert await app.ainvoke(None, config, checkpoint_during=checkpoint_during) == { "my_key": "hi my value here and there and back again", } @@ -5740,12 +5929,22 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: config = { "configurable": {"thread_id": "2", CONFIG_KEY_NODE_FINISHED: nodes.append} } - assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ + assert [ + c + async for c in app.astream( + {"my_key": "my value"}, config, checkpoint_during=checkpoint_during + ) + ] == [ {"parent_1": {"my_key": "hi my value"}}, {"__interrupt__": ()}, ] assert nodes == ["parent_1", "grandchild_1"] - assert [c async for c in app.astream(None, config)] == [ + assert [ + c + async for c in app.astream( + None, config, checkpoint_during=checkpoint_during + ) + ] == [ {"child": {"my_key": "hi my value here and there"}}, {"parent_2": {"my_key": "hi my value here and there and back again"}}, ] @@ -5763,13 +5962,21 @@ async def test_doubly_nested_graph_interrupts(checkpointer_name: str) -> None: assert [ c async for c in app.astream( - {"my_key": "my value"}, config, stream_mode="values" + {"my_key": "my value"}, + config, + stream_mode="values", + checkpoint_during=checkpoint_during, ) ] == [ {"my_key": "my value"}, {"my_key": "hi my value"}, ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ + assert [ + c + async for c in app.astream( + None, config, stream_mode="values", checkpoint_during=checkpoint_during + ) + ] == [ {"my_key": "hi my value"}, {"my_key": "hi my value here and there"}, {"my_key": "hi my value here and there and back again"},