From 70eeb2a67059cca1f6c2e2aa88dcc77463fa4323 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Tue, 10 Dec 2024 11:52:49 -0500 Subject: [PATCH 1/2] x --- libs/langgraph/tests/test_pregel.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0103ce6b2..684dfa9fb 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -14942,8 +14942,8 @@ def test_command_with_static_breakpoints() -> None: # Start the graph and interrupt at the first node graph.invoke({"foo": "abc"}, config) - result = graph.invoke(Command(resume="node1"), config) - assert result == {"foo": "abc|node-1|node-2"} + result = graph.invoke(Command(update={"foo": "def"}), config) + assert result == {"foo": "def|node-1|node-2"} @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) From 081b2cbdcf9ab503f3fb5181d8856526114fb900 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 10 Dec 2024 11:00:09 -0800 Subject: [PATCH 2/2] Fix --- libs/langgraph/langgraph/pregel/loop.py | 24 ++++++++-------- libs/langgraph/tests/test_pregel.py | 16 +++++------ libs/langgraph/tests/test_pregel_async.py | 35 +++++++++++++++++++++++ 3 files changed, 54 insertions(+), 21 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index e96276259..a8e945edd 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -422,18 +422,6 @@ class PregelLoop(LoopProtocol): self.status = "out_of_steps" return False - # apply NULL writes - if null_writes := [ - w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID - ]: - mv_writes = apply_writes( - self.checkpoint, - self.channels, - [PregelTaskWrites((), INPUT, null_writes, [])], - self.checkpointer_get_next_version, - ) - for key, values in mv_writes.items(): - self._update_mv(key, values) # prepare next tasks self.tasks = prepare_next_tasks( self.checkpoint, @@ -552,6 +540,18 @@ class PregelLoop(LoopProtocol): # save writes for tid, ws in writes.items(): self.put_writes(tid, ws) + # apply NULL writes + if null_writes := [ + w[1:] for w in self.checkpoint_pending_writes if w[0] == NULL_TASK_ID + ]: + mv_writes = apply_writes( + self.checkpoint, + self.channels, + [PregelTaskWrites((), INPUT, null_writes, [])], + self.checkpointer_get_next_version, + ) + for key, values in mv_writes.items(): + self._update_mv(key, values) # proceed past previous checkpoint if is_resuming: self.checkpoint["versions_seen"].setdefault(INTERRUPT, {}) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 684dfa9fb..20aaec74c 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -14906,9 +14906,14 @@ def test_dict_mixed_return() -> None: assert graph.invoke({"foo": ""}) == {"foo": "ab"} -def test_command_with_static_breakpoints() -> None: +@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) +def test_command_with_static_breakpoints( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: """Test that we can use Command to resume and update with static breakpoints.""" + checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + class State(TypedDict): """The graph state.""" @@ -14930,15 +14935,8 @@ def test_command_with_static_breakpoints() -> None: builder.add_edge(START, "node1") builder.add_edge("node1", "node2") - # A checkpointer must be enabled for interrupts to work! - checkpointer = MemorySaver() graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"]) - - config = { - "configurable": { - "thread_id": uuid.uuid4(), - } - } + config = {"configurable": {"thread_id": str(uuid.uuid4())}} # Start the graph and interrupt at the first node graph.invoke({"foo": "abc"}, config) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 0064117bf..5cc7f3312 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -13201,6 +13201,41 @@ async def test_interrupt_loop(checkpointer_name: str): ] +@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.""" + + class State(TypedDict): + """The graph state.""" + + foo: str + + def node1(state: State): + return { + "foo": state["foo"] + "|node-1", + } + + def node2(state: State): + return { + "foo": state["foo"] + "|node-2", + } + + builder = StateGraph(State) + builder.add_node("node1", node1) + builder.add_node("node2", node2) + builder.add_edge(START, "node1") + builder.add_edge("node1", "node2") + + async with awith_checkpointer(checkpointer_name) as checkpointer: + graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node1"]) + config = {"configurable": {"thread_id": str(uuid.uuid4())}} + + # Start the graph and interrupt at the first node + await graph.ainvoke({"foo": "abc"}, config) + result = await graph.ainvoke(Command(update={"foo": "def"}), config) + assert result == {"foo": "def|node-1|node-2"} + + @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_multistep_plan(checkpointer_name: str): from langchain_core.messages import AnyMessage