From e6726802f7fcc4cc59cec958c4e4bab62315761a Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Fri, 7 Mar 2025 19:21:35 +0100 Subject: [PATCH] feat(pregel): add bulk update state method This method is useful for recreating a thread from a list of checkpoint writes. A new method is needed to clone a checkpoint that has been created from multiple writes (functional API, map-reduce) Port of https://github.com/langchain-ai/langgraphjs/pull/969 --- libs/langgraph/langgraph/pregel/__init__.py | 425 +++++++++++++------- libs/langgraph/langgraph/pregel/protocol.py | 16 +- libs/langgraph/langgraph/pregel/types.py | 2 + libs/langgraph/langgraph/types.py | 5 + 4 files changed, 292 insertions(+), 156 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index f6479422c..e84c0b208 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -106,6 +106,7 @@ from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore from langgraph.types import ( All, + BulkUpdate, Checkpointer, LoopProtocol, StateSnapshot, @@ -1162,22 +1163,22 @@ class Pregel(PregelProtocol): checkpoint_tuple.config, checkpoint_tuple ) - def update_state( + def bulk_update_state( self, config: RunnableConfig, - values: Optional[Union[dict[str, Any], Any]], - as_node: Optional[str] = None, + updates: list[BulkUpdate], ) -> RunnableConfig: - """Update the state of the graph with the given values, as if they came from - node `as_node`. If `as_node` is not provided, it will be set to the last node - that updated the state, if not ambiguous. - """ + """Apply updates to the graph state in bulk. Requires a checkpointer to be set.""" + checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get( CONFIG_KEY_CHECKPOINTER, self.checkpointer ) if not checkpointer: raise ValueError("No checkpointer set") + if len(updates) == 0: + raise ValueError("No updates provided") + # delegate to subgraph if ( checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") @@ -1186,10 +1187,9 @@ class Pregel(PregelProtocol): recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name for _, pregel in self.get_subgraphs(namespace=recast, recurse=True): - return pregel.update_state( + return pregel.bulk_update_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), - values, - as_node, + updates, ) else: raise ValueError(f"Subgraph {recast} not found") @@ -1216,8 +1216,15 @@ class Pregel(PregelProtocol): checkpoint, LoopProtocol(config=config, step=step + 1, stop=step + 2), ) as (channels, managed): + values, as_node = updates[0] + # no values as END, just clear all tasks if values is None and as_node == END: + if len(updates) > 1: + raise ValueError( + "Cannot apply multiple updates when clearing state" + ) + if saved is not None: # tasks for this checkpoint next_tasks = prepare_next_tasks( @@ -1276,6 +1283,11 @@ class Pregel(PregelProtocol): ) # no values, empty checkpoint if values is None and as_node is None: + if len(updates) > 1: + raise ValueError( + "Cannot create empty checkpoint with multiple updates" + ) + next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint next_config = checkpointer.put( @@ -1295,6 +1307,9 @@ class Pregel(PregelProtocol): ) # no values, copy checkpoint if values is None and as_node == "__copy__": + if len(updates) > 1: + raise ValueError("Cannot copy checkpoint with multiple updates") + next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint next_config = checkpointer.put( @@ -1354,79 +1369,106 @@ class Pregel(PregelProtocol): next_tasks[tid].writes.append((k, v)) 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 as_node is None and not any( - v for vv in checkpoint["versions_seen"].values() for v in vv.values() - ): - if ( - isinstance(self.input_channels, str) - and self.input_channels in self.nodes + valid_updates: list[tuple[Optional[dict[str, Any]], str]] = [] + if len(updates) == 1: + values, as_node = updates[0] + + # find last node that updated the state, if not provided + if as_node is None and not any( + v + for vv in checkpoint["versions_seen"].values() + for v in vv.values() ): - as_node = self.input_channels - elif as_node is None: - last_seen_by_node = sorted( - (v, n) - for n, seen in checkpoint["versions_seen"].items() - if n in self.nodes - for v in seen.values() + if ( + isinstance(self.input_channels, str) + and self.input_channels in self.nodes + ): + as_node = self.input_channels + elif as_node is None: + last_seen_by_node = sorted( + (v, n) + for n, seen in checkpoint["versions_seen"].items() + if n in self.nodes + for v in seen.values() + ) + # if two nodes updated the state at the same time, it's ambiguous + if last_seen_by_node: + if len(last_seen_by_node) == 1: + as_node = last_seen_by_node[0][1] + elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: + as_node = last_seen_by_node[-1][1] + if as_node is None: + raise InvalidUpdateError("Ambiguous update, specify as_node") + if as_node not in self.nodes: + raise InvalidUpdateError(f"Node {as_node} does not exist") + + valid_updates.append((values, as_node)) + else: + for values, as_node in updates: + if as_node is None: + raise InvalidUpdateError( + "as_node is required when applying multiple updates" + ) + + if as_node not in self.nodes: + raise InvalidUpdateError(f"Node {as_node} does not exist") + + valid_updates.append((values, as_node)) + + tasks: list[tuple[str, PregelTaskWrites]] = [] + for values, as_node in valid_updates: + # create task to run all writers of the chosen node + writers = self.nodes[as_node].flat_writers + if not writers: + raise InvalidUpdateError(f"Node {as_node} has no writers") + writes: deque[tuple[str, Any]] = deque() + task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) + task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) + tasks.append((task_id, task)) + + for _, task in tasks: + run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] + # execute task + run.invoke( + values, + patch_config( + config, + run_name=self.name + "UpdateState", + configurable={ + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + local_write, + writes.extend, + self.nodes.keys(), + ), + CONFIG_KEY_READ: partial( + local_read, + step + 1, + checkpoint, + channels, + managed, + task, + config, + ), + }, + ), ) - # if two nodes updated the state at the same time, it's ambiguous - if last_seen_by_node: - if len(last_seen_by_node) == 1: - as_node = last_seen_by_node[0][1] - elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: - as_node = last_seen_by_node[-1][1] - if as_node is None: - raise InvalidUpdateError("Ambiguous update, specify as_node") - if as_node not in self.nodes: - raise InvalidUpdateError(f"Node {as_node} does not exist") - # create task to run all writers of the chosen node - writers = self.nodes[as_node].flat_writers - if not writers: - raise InvalidUpdateError(f"Node {as_node} has no writers") - writes: deque[tuple[str, Any]] = deque() - task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) - task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) - run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] - # execute task - run.invoke( - values, - patch_config( - config, - run_name=self.name + "UpdateState", - configurable={ - # deque.extend is thread-safe - CONFIG_KEY_SEND: partial( - local_write, - writes.extend, - self.nodes.keys(), - ), - CONFIG_KEY_READ: partial( - local_read, - step + 1, - checkpoint, - channels, - managed, - task, - config, - ), - }, - ), - ) + # save task writes - # channel writes are saved to current checkpoint - # push writes are saved to next checkpoint - channel_writes, push_writes = ( - [w for w in task.writes if w[0] != PUSH], - [w for w in task.writes if w[0] == PUSH], - ) - if saved and channel_writes: - checkpointer.put_writes(checkpoint_config, channel_writes, task_id) + for _, task in tasks: + channel_writes = [w for w in task.writes if w[0] != PUSH] + + # channel writes are saved to current checkpoint + if saved and channel_writes: + checkpointer.put_writes(checkpoint_config, channel_writes, task_id) + # apply to checkpoint and save mv_writes = apply_writes( - checkpoint, channels, [task], checkpointer.get_next_version + checkpoint, channels, tasks, checkpointer.get_next_version ) + assert not mv_writes, "Can't write to SharedValues from update_state" + checkpoint = create_checkpoint(checkpoint, channels, step + 1) next_config = checkpointer.put( checkpoint_config, @@ -1435,23 +1477,27 @@ class Pregel(PregelProtocol): **checkpoint_metadata, "source": "update", "step": step + 1, - "writes": {as_node: values}, + "writes": {k: v for k, v in valid_updates}, "parents": saved.metadata.get("parents", {}) if saved else {}, }, get_new_channel_versions( checkpoint_previous_versions, checkpoint["channel_versions"] ), ) - if push_writes: - checkpointer.put_writes(next_config, push_writes, task_id) + + for _, task in tasks: + if push_writes := [w for w in task.writes if w[0] == PUSH]: + checkpointer.put_writes(next_config, push_writes, task_id) + return patch_checkpoint_map(next_config, saved.metadata if saved else None) - async def aupdate_state( + async def abulk_update_state( self, config: RunnableConfig, - values: dict[str, Any] | Any, - as_node: Optional[str] = None, + updates: list[BulkUpdate], ) -> RunnableConfig: + """Apply updates to the graph state in bulk. Requires a checkpointer to be set.""" + """Update the state of the graph asynchronously with the given values, as if they came from node `as_node`. If `as_node` is not provided, it will be set to the last node that updated the state, if not ambiguous. @@ -1462,6 +1508,9 @@ class Pregel(PregelProtocol): if not checkpointer: raise ValueError("No checkpointer set") + if len(updates) == 0: + raise ValueError("No updates provided") + # delegate to subgraph if ( checkpoint_ns := config[CONF].get(CONFIG_KEY_CHECKPOINT_NS, "") @@ -1470,10 +1519,9 @@ class Pregel(PregelProtocol): recast = recast_checkpoint_ns(checkpoint_ns) # find the subgraph with the matching name async for _, pregel in self.aget_subgraphs(namespace=recast, recurse=True): - return await pregel.aupdate_state( + return await pregel.abulk_update_state( patch_configurable(config, {CONFIG_KEY_CHECKPOINTER: checkpointer}), - values, - as_node, + updates, ) else: raise ValueError(f"Subgraph {recast} not found") @@ -1503,8 +1551,15 @@ class Pregel(PregelProtocol): channels, managed, ): + values, as_node = updates[0] + # no values, just clear all tasks if values is None and as_node == END: + if len(updates) > 1: + raise ValueError( + "Cannot apply multiple updates when clearing state" + ) + if saved is not None: # tasks for this checkpoint next_tasks = prepare_next_tasks( @@ -1563,6 +1618,11 @@ class Pregel(PregelProtocol): ) # no values, empty checkpoint if values is None and as_node is None: + if len(updates) > 1: + raise ValueError( + "Cannot create empty checkpoint with multiple updates" + ) + next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint next_config = await checkpointer.aput( @@ -1582,6 +1642,9 @@ class Pregel(PregelProtocol): ) # no values, copy checkpoint if values is None and as_node == "__copy__": + if len(updates) > 1: + raise ValueError("Cannot copy checkpoint with multiple updates") + next_checkpoint = create_checkpoint(checkpoint, None, step) # copy checkpoint next_config = await checkpointer.aput( @@ -1640,74 +1703,99 @@ class Pregel(PregelProtocol): next_tasks[tid].writes.append((k, v)) 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 as_node is None and not saved: - if ( - isinstance(self.input_channels, str) - and self.input_channels in self.nodes - ): - as_node = self.input_channels - elif as_node is None: - last_seen_by_node = sorted( - (v, n) - for n, seen in checkpoint["versions_seen"].items() - if n in self.nodes - for v in seen.values() + valid_updates: list[tuple[Optional[dict[str, Any]], str]] = [] + + if len(updates) == 1: + values, as_node = updates[0] + + # find last node that updated the state, if not provided + if as_node is None and not saved: + if ( + isinstance(self.input_channels, str) + and self.input_channels in self.nodes + ): + as_node = self.input_channels + elif as_node is None: + last_seen_by_node = sorted( + (v, n) + for n, seen in checkpoint["versions_seen"].items() + if n in self.nodes + for v in seen.values() + ) + # if two nodes updated the state at the same time, it's ambiguous + if last_seen_by_node: + if len(last_seen_by_node) == 1: + as_node = last_seen_by_node[0][1] + elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: + as_node = last_seen_by_node[-1][1] + if as_node is None: + raise InvalidUpdateError("Ambiguous update, specify as_node") + + if as_node not in self.nodes: + raise InvalidUpdateError(f"Node {as_node} does not exist") + + valid_updates.append((values, as_node)) + else: + for values, as_node in updates: + if as_node is None: + raise InvalidUpdateError( + "as_node is required when applying multiple updates" + ) + + if as_node not in self.nodes: + raise InvalidUpdateError(f"Node {as_node} does not exist") + + valid_updates.append((values, as_node)) + + tasks: list[tuple[str, PregelTaskWrites]] = [] + + for values, as_node in valid_updates: + # create task to run all writers of the chosen node + writers = self.nodes[as_node].flat_writers + if not writers: + raise InvalidUpdateError(f"Node {as_node} has no writers") + writes: deque[tuple[str, Any]] = deque() + task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) + task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) + tasks.append((task_id, task)) + + for _, task in tasks: + run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] + # execute task + await run.ainvoke( + values, + patch_config( + config, + run_name=self.name + "UpdateState", + configurable={ + # deque.extend is thread-safe + CONFIG_KEY_SEND: partial( + local_write, + writes.extend, + self.nodes.keys(), + ), + CONFIG_KEY_READ: partial( + local_read, + step + 1, + checkpoint, + channels, + managed, + task, + config, + ), + }, + ), ) - # if two nodes updated the state at the same time, it's ambiguous - if last_seen_by_node: - if len(last_seen_by_node) == 1: - as_node = last_seen_by_node[0][1] - elif last_seen_by_node[-1][0] != last_seen_by_node[-2][0]: - as_node = last_seen_by_node[-1][1] - if as_node is None: - raise InvalidUpdateError("Ambiguous update, specify as_node") - if as_node not in self.nodes: - raise InvalidUpdateError(f"Node {as_node} does not exist") - # create task to run all writers of the chosen node - writers = self.nodes[as_node].flat_writers - if not writers: - raise InvalidUpdateError(f"Node {as_node} has no writers") - writes: deque[tuple[str, Any]] = deque() - task = PregelTaskWrites((), as_node, writes, [INTERRUPT]) - task_id = str(uuid5(UUID(checkpoint["id"]), INTERRUPT)) - run = RunnableSequence(*writers) if len(writers) > 1 else writers[0] - # execute task - await run.ainvoke( - values, - patch_config( - config, - run_name=self.name + "UpdateState", - configurable={ - # deque.extend is thread-safe - CONFIG_KEY_SEND: partial( - local_write, - writes.extend, - self.nodes.keys(), - ), - CONFIG_KEY_READ: partial( - local_read, - step + 1, - checkpoint, - channels, - managed, - task, - config, - ), - }, - ), - ) + # save task writes - # channel writes are saved to current checkpoint - # push writes are saved to next checkpoint - channel_writes, push_writes = ( - [w for w in task.writes if w[0] != PUSH], - [w for w in task.writes if w[0] == PUSH], - ) - if saved and channel_writes: - await checkpointer.aput_writes( - checkpoint_config, channel_writes, task_id - ) + for _, task in tasks: + # channel writes are saved to current checkpoint + channel_writes = [w for w in task.writes if w[0] != PUSH] + if saved and channel_writes: + await checkpointer.aput_writes( + checkpoint_config, channel_writes, task_id + ) + # apply to checkpoint and save mv_writes = apply_writes( checkpoint, channels, [task], checkpointer.get_next_version @@ -1722,18 +1810,45 @@ class Pregel(PregelProtocol): **checkpoint_metadata, "source": "update", "step": step + 1, - "writes": {as_node: values}, + "writes": {k: v for k, v in valid_updates}, "parents": saved.metadata.get("parents", {}) if saved else {}, }, get_new_channel_versions( checkpoint_previous_versions, checkpoint["channel_versions"] ), ) - # save push writes - if push_writes: - await checkpointer.aput_writes(next_config, push_writes, task_id) + + for _, task in tasks: + # save push writes + if push_writes := [w for w in task.writes if w[0] == PUSH]: + await checkpointer.aput_writes(next_config, push_writes, task_id) + return patch_checkpoint_map(next_config, saved.metadata if saved else None) + def update_state( + self, + config: RunnableConfig, + values: Optional[Union[dict[str, Any], Any]], + as_node: Optional[str] = None, + ) -> RunnableConfig: + """Update the state of the graph with the given values, as if they came from + node `as_node`. If `as_node` is not provided, it will be set to the last node + that updated the state, if not ambiguous. + """ + return self.bulk_update_state(config, [(values, as_node)]) + + async def aupdate_state( + self, + config: RunnableConfig, + values: dict[str, Any] | Any, + as_node: Optional[str] = None, + ) -> RunnableConfig: + """Update the state of the graph asynchronously with the given values, as if they came from + node `as_node`. If `as_node` is not provided, it will be set to the last node + that updated the state, if not ambiguous. + """ + return await self.bulk_update_state(config, [(values, as_node)]) + def _defaults( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/pregel/protocol.py b/libs/langgraph/langgraph/pregel/protocol.py index ac046e949..43dd55652 100644 --- a/libs/langgraph/langgraph/pregel/protocol.py +++ b/libs/langgraph/langgraph/pregel/protocol.py @@ -12,7 +12,7 @@ from langchain_core.runnables import Runnable, RunnableConfig from langchain_core.runnables.graph import Graph as DrawableGraph from typing_extensions import Self -from langgraph.pregel.types import All, StateSnapshot, StreamMode +from langgraph.pregel.types import All, BulkUpdate, StateSnapshot, StreamMode class PregelProtocol( @@ -69,6 +69,20 @@ class PregelProtocol( limit: Optional[int] = None, ) -> AsyncIterator[StateSnapshot]: ... + @abstractmethod + def bulk_update_state( + self, + config: RunnableConfig, + updates: Sequence[BulkUpdate], + ) -> RunnableConfig: ... + + @abstractmethod + async def abulk_update_state( + self, + config: RunnableConfig, + updates: Sequence[BulkUpdate], + ) -> RunnableConfig: ... + @abstractmethod def update_state( self, diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 7a72b88c9..251d59034 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -2,6 +2,7 @@ from langgraph.types import ( All, + BulkUpdate, CachePolicy, PregelExecutableTask, PregelTask, @@ -14,6 +15,7 @@ from langgraph.types import ( __all__ = [ "All", + "BulkUpdate", "CachePolicy", "PregelExecutableTask", "PregelTask", diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 4c8cf6da4..0d530ddb5 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -133,6 +133,11 @@ class Interrupt: when: Literal["during"] = dataclasses.field(default="during", repr=False) +class BulkUpdate(NamedTuple): + values: dict[str, Any] | None + as_node: Optional[str] + + class PregelTask(NamedTuple): id: str name: str