WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-06 17:57:49 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(pregel): add bulk update state method

Browse Source 
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
This commit is contained in:
Tat Dat Duong
2025-03-19 14:21:43 +01:00
parent 24f7d7c439
commit e6726802f7
4 changed files with 292 additions and 156 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+270 -155
View File
@@ -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,
libs/langgraph/langgraph/pregel/protocol.py
+15 -1
View File
@@ -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,
libs/langgraph/langgraph/pregel/types.py
+2
View File
@@ -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",
libs/langgraph/langgraph/types.py
+5
View File
@@ -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