Merge pull request #2388 from langchain-ai/nc/11nov/update-none-clear-all-tasks

lib: update_state(values=None) should clear all tasks

libs/langgraph/langgraph/pregel/__init__.py

@@ -478,6 +478,11 @@ class Pregel(PregelProtocol):
                 checkpointer=self.checkpointer or None,
                 manager=None,
             )
+            print(
+                saved.checkpoint["versions_seen"],
+                saved.checkpoint["pending_sends"],
+                # next_tasks,
+            )
             # get the subgraphs
             subgraphs = dict(self.get_subgraphs())
             parent_ns = saved.config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "")
@@ -872,6 +877,47 @@ class Pregel(PregelProtocol):
             checkpoint,
             LoopProtocol(config=config, step=step + 1, stop=step + 2),
         ) as (channels, managed):
+            # no values, just clear all tasks
+            if values is None and as_node is None:
+                if saved is not None:
+                    # tasks for this checkpoint
+                    next_tasks = prepare_next_tasks(
+                        checkpoint,
+                        self.nodes,
+                        channels,
+                        managed,
+                        saved.config,
+                        saved.metadata.get("step", -1) + 1,
+                        for_execution=True,
+                        store=self.store,
+                        checkpointer=self.checkpointer or None,
+                        manager=None,
+                    )
+                    # apply writes from tasks that already ran
+                    for tid, k, v in saved.pending_writes or []:
+                        if k in (ERROR, INTERRUPT, SCHEDULED):
+                            continue
+                        if tid not in next_tasks:
+                            continue
+                        next_tasks[tid].writes.append((k, v))
+                    # clear all current tasks
+                    apply_writes(checkpoint, channels, next_tasks.values(), None)
+                # save checkpoint
+                next_config = checkpointer.put(
+                    checkpoint_config,
+                    create_checkpoint(checkpoint, None, step),
+                    {
+                        **checkpoint_metadata,
+                        "source": "update",
+                        "step": step + 1,
+                        "writes": {},
+                        "parents": saved.metadata.get("parents", {}) if saved else {},
+                    },
+                    {},
+                )
+                return patch_checkpoint_map(
+                    next_config, saved.metadata if saved else None
+                )
             # apply pending writes, if not on specific checkpoint
             if (
                 CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
@@ -901,23 +947,7 @@ class Pregel(PregelProtocol):
                 if tasks := [t for t in next_tasks.values() if t.writes]:
                     apply_writes(checkpoint, channels, tasks, None)
             # find last node that updated the state, if not provided
-            if values is None and as_node is None:
-                next_config = checkpointer.put(
-                    checkpoint_config,
-                    create_checkpoint(checkpoint, None, step),
-                    {
-                        **checkpoint_metadata,
-                        "source": "update",
-                        "step": step + 1,
-                        "writes": {},
-                        "parents": saved.metadata.get("parents", {}) if saved else {},
-                    },
-                    {},
-                )
-                return patch_checkpoint_map(
-                    next_config, saved.metadata if saved else None
-                )
-            elif as_node is None and not any(
+            if as_node is None and not any(
                 v for vv in checkpoint["versions_seen"].values() for v in vv.values()
             ):
                 if (
@@ -1057,6 +1087,47 @@ class Pregel(PregelProtocol):
             channels,
             managed,
         ):
+            # no values, just clear all tasks
+            if values is None and as_node is None:
+                if saved is not None:
+                    # tasks for this checkpoint
+                    next_tasks = prepare_next_tasks(
+                        checkpoint,
+                        self.nodes,
+                        channels,
+                        managed,
+                        saved.config,
+                        saved.metadata.get("step", -1) + 1,
+                        for_execution=True,
+                        store=self.store,
+                        checkpointer=self.checkpointer or None,
+                        manager=None,
+                    )
+                    # apply writes from tasks that already ran
+                    for tid, k, v in saved.pending_writes or []:
+                        if k in (ERROR, INTERRUPT, SCHEDULED):
+                            continue
+                        if tid not in next_tasks:
+                            continue
+                        next_tasks[tid].writes.append((k, v))
+                    # clear all current tasks
+                    apply_writes(checkpoint, channels, next_tasks.values(), None)
+                # save checkpoint
+                next_config = await checkpointer.aput(
+                    checkpoint_config,
+                    create_checkpoint(checkpoint, None, step),
+                    {
+                        **checkpoint_metadata,
+                        "source": "update",
+                        "step": step + 1,
+                        "writes": {},
+                        "parents": saved.metadata.get("parents", {}) if saved else {},
+                    },
+                    {},
+                )
+                return patch_checkpoint_map(
+                    next_config, saved.metadata if saved else None
+                )
             # apply pending writes, if not on specific checkpoint
             if (
                 CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
@@ -1085,23 +1156,7 @@ class Pregel(PregelProtocol):
                 if tasks := [t for t in next_tasks.values() if t.writes]:
                     apply_writes(checkpoint, channels, tasks, None)
             # find last node that updated the state, if not provided
-            if values is None and as_node is None:
-                next_config = await checkpointer.aput(
-                    checkpoint_config,
-                    create_checkpoint(checkpoint, None, step),
-                    {
-                        **checkpoint_metadata,
-                        "source": "update",
-                        "step": step + 1,
-                        "writes": {},
-                        "parents": saved.metadata.get("parents", {}) if saved else {},
-                    },
-                    {},
-                )
-                return patch_checkpoint_map(
-                    next_config, saved.metadata if saved else None
-                )
-            elif as_node is None and not saved:
+            if as_node is None and not saved:
                 if (
                     isinstance(self.input_channels, str)
                     and self.input_channels in self.nodes

libs/langgraph/langgraph/pregel/algo.py

@@ -197,6 +197,7 @@ def apply_writes(
     # sort tasks on path
     tasks = sorted(tasks, key=lambda t: t.path)
 
+    print("versions_seen", checkpoint["versions_seen"], [task.name for task in tasks])
     # update seen versions
     for task in tasks:
         checkpoint["versions_seen"].setdefault(task.name, {}).update(

libs/langgraph/tests/test_pregel.py

@@ -7968,6 +7968,24 @@ def test_dynamic_interrupt(
         },
         parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config,
     )
+    # clear the interrupt and next tasks
+    tool_two.update_state(thread1, None)
+    # interrupt is cleared, task will still run next
+    assert tool_two.get_state(thread1) == StateSnapshot(
+        values={"my_key": "value ⛰️", "market": "DE"},
+        next=(),
+        tasks=(),
+        config=tool_two.checkpointer.get_tuple(thread1).config,
+        created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
+        metadata={
+            "parents": {},
+            "source": "update",
+            "step": 1,
+            "writes": {},
+            "thread_id": "1",
+        },
+        parent_config=[*tool_two.checkpointer.list(thread1, limit=2)][-1].config,
+    )
 
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)

libs/langgraph/tests/test_pregel_async.py

@@ -342,6 +342,28 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None:
             ][-1].config,
         )
 
+        # clear the interrupt and next tasks
+        await tool_two.aupdate_state(thread1, None)
+        # interrupt is cleared, task will still run next
+        tup = await tool_two.checkpointer.aget_tuple(thread1)
+        assert await tool_two.aget_state(thread1) == StateSnapshot(
+            values={"my_key": "value ⛰️", "market": "DE"},
+            next=(),
+            tasks=(),
+            config=tup.config,
+            created_at=tup.checkpoint["ts"],
+            metadata={
+                "parents": {},
+                "source": "update",
+                "step": 1,
+                "writes": {},
+                "thread_id": "1",
+            },
+            parent_config=[
+                c async for c in tool_two.checkpointer.alist(thread1, limit=2)
+            ][-1].config,
+        )
+
 
 @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
 async def test_node_not_cancelled_on_other_node_interrupted(