From c043f148a6a22739f1fa5e36e20d0c4654727a61 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 10:48:34 -0800 Subject: [PATCH 1/2] lib: update_state(values=None) should clear all tasks --- libs/langgraph/langgraph/pregel/__init__.py | 121 ++++++++++++++------ libs/langgraph/langgraph/pregel/algo.py | 1 + libs/langgraph/tests/test_pregel.py | 18 +++ libs/langgraph/tests/test_pregel_async.py | 22 ++++ 4 files changed, 128 insertions(+), 34 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 5ff3c69b8..30e393ce4 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/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,46 @@ 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: + # 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: + 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 +946,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 +1086,46 @@ class Pregel(PregelProtocol): channels, managed, ): + # no values, just clear all tasks + if values is None and as_node is 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: + 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 +1154,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 diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index af71294ae..3b7bc12a3 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/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( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index e84a35214..9d49e8551 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index d04c7e48f..364be8bfa 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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( From c0513076a252e52b7840a50f216be95dc81281cd Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 11 Nov 2024 11:07:36 -0800 Subject: [PATCH 2/2] Lint --- libs/langgraph/langgraph/pregel/__init__.py | 90 +++++++++++---------- 1 file changed, 46 insertions(+), 44 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 30e393ce4..84f8af8b6 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -879,28 +879,29 @@ class Pregel(PregelProtocol): ) as (channels, managed): # no values, just clear all tasks if values is None and as_node is 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: - 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) + 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, @@ -1088,28 +1089,29 @@ class Pregel(PregelProtocol): ): # no values, just clear all tasks if values is None and as_node is 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: - 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) + 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,