diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index c1575dd7f..a30a66cba 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -4949,6 +4949,69 @@ def test_interrupt_loop(request: pytest.FixtureRequest, checkpointer_name: str): ] +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_interrupt_functional( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + @task + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + def bar(state: dict) -> dict: + return {"a": state["a"] + "bar", "b": state["b"]} + + @entrypoint(checkpointer=checkpointer) + def graph(inputs: dict) -> dict: + fut_foo = foo(inputs) + value = interrupt("Provide value for bar:") + bar_input = {**fut_foo.result(), "b": value} + fut_bar = bar(bar_input) + return fut_bar.result() + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + graph.invoke({"a": ""}, config) + # Resume with an answer + res = graph.invoke(Command(resume="bar"), config) + assert res == {"a": "foobar", "b": "bar"} + + +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_interrupt_task_functional( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) + + @task + def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + def bar(state: dict) -> dict: + value = interrupt("Provide value for bar:") + return {"a": state["a"] + value} + + @entrypoint(checkpointer=checkpointer) + def graph(inputs: dict) -> dict: + fut_foo = foo(inputs) + fut_bar = bar(fut_foo.result()) + return fut_bar.result() + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + graph.invoke({"a": ""}, config) + # Resume with an answer + res = graph.invoke(Command(resume="bar"), config) + assert res == {"a": "foobar"} + + def test_root_mixed_return() -> None: def my_node(state: list[str]): return [Command(update=["a"]), ["b"]] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 558a93709..6767a1510 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6281,6 +6281,63 @@ async def test_interrupt_loop(checkpointer_name: str): ] +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_interrupt_functional(checkpointer_name: str) -> None: + @task + async def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + async def bar(state: dict) -> dict: + return {"a": state["a"] + "bar", "b": state["b"]} + + async with awith_checkpointer(checkpointer_name) as checkpointer: + + @entrypoint(checkpointer=checkpointer) + async def graph(inputs: dict) -> dict: + foo_result = await foo(inputs) + value = interrupt("Provide value for bar:") + bar_input = {**foo_result, "b": value} + bar_result = await bar(bar_input) + return bar_result + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + await graph.ainvoke({"a": ""}, config) + # Resume with an answer + res = await graph.ainvoke(Command(resume="bar"), config) + assert res == {"a": "foobar", "b": "bar"} + + +@NEEDS_CONTEXTVARS +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) +async def test_interrupt_task_functional(checkpointer_name: str) -> None: + @task + async def foo(state: dict) -> dict: + return {"a": state["a"] + "foo"} + + @task + async def bar(state: dict) -> dict: + value = interrupt("Provide value for bar:") + return {"a": state["a"] + value} + + async with awith_checkpointer(checkpointer_name) as checkpointer: + + @entrypoint(checkpointer=checkpointer) + async def graph(inputs: dict) -> dict: + foo_result = await foo(inputs) + bar_result = await bar(foo_result) + return bar_result + + config = {"configurable": {"thread_id": "1"}} + # First run, interrupted at bar + await graph.ainvoke({"a": ""}, config) + # Resume with an answer + res = await graph.ainvoke(Command(resume="bar"), config) + assert res == {"a": "foobar"} + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_command_with_static_breakpoints(checkpointer_name: str) -> None: """Test that we can use Command to resume and update with static breakpoints."""