From 8b1597a38543534d46e7228938e50670cf0f3083 Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 16 Jan 2025 10:13:41 -0500 Subject: [PATCH 1/4] tests: add a test for interrupt() w/ functional API --- libs/langgraph/tests/test_pregel.py | 31 +++++++++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 27 ++++++++++++++++++++ 2 files changed, 58 insertions(+) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 64b5076de..a50de8466 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -4902,6 +4902,37 @@ 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: + 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..1662ac478 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6281,6 +6281,33 @@ async def test_interrupt_loop(checkpointer_name: str): ] +@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: + 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.""" From c4460e5dd2e69626ce5a0a58f604aad64c07b15c Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 16 Jan 2025 11:16:09 -0500 Subject: [PATCH 2/4] add another test --- libs/langgraph/tests/test_pregel.py | 32 +++++++++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 30 ++++++++++++++++++++- 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a50de8466..83d0f792b 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -4910,6 +4910,38 @@ def test_interrupt_functional( 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"} diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 1662ac478..541781c74 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6287,6 +6287,34 @@ async def test_interrupt_functional(checkpointer_name: str) -> None: 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"} + + +@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:") @@ -6305,7 +6333,7 @@ async def test_interrupt_functional(checkpointer_name: str) -> None: await graph.ainvoke({"a": ""}, config) # Resume with an answer res = await graph.ainvoke(Command(resume="bar"), config) - assert res == {"a": "foobar"} + assert res == {"a": "foobar"} @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) From 8c88e203bc541c50e27a296657aef9941fd5b5f1 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 16 Jan 2025 09:51:46 -0800 Subject: [PATCH 3/4] Fix --- libs/langgraph/langgraph/pregel/loop.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 1f2fe91a3..e92f0c267 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -345,11 +345,11 @@ class PregelLoop(LoopProtocol): (PUSH, task.path, write_idx, task.id, call), None, checkpoint=self.checkpoint, - pending_writes=[(task.id, *w) for w in task.writes], + pending_writes=self.checkpoint_pending_writes, processes=self.nodes, channels=self.channels, managed=self.managed, - config=self.config, + config=task.config, step=self.step, for_execution=True, store=self.store, From a16def5140d45be78d0ae70651b71a9c37c7f54e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 16 Jan 2025 09:55:04 -0800 Subject: [PATCH 4/4] Lint --- libs/langgraph/tests/test_pregel_async.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 541781c74..6767a1510 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6281,6 +6281,7 @@ 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 @@ -6309,6 +6310,7 @@ async def test_interrupt_functional(checkpointer_name: str) -> None: 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