WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-17 21:25:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(langgraph): task masquerading with update state (#5189)

Browse Source
This commit is contained in:
Nuno Campos
2025-06-25 17:13:54 -07:00
committed by GitHub
parent 929bba8337 cee9ac0b7a
commit 1c9bfba23a
5 changed files with 906 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/pregel/__init__.py
+191 -62
View File
@@ -1354,7 +1354,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
Args:
config: The config to apply the updates to.
supersteps: A list of supersteps, each including a list of updates to apply sequentially to a graph state.
Each update is a tuple of the form `(values, as_node)`.
Each update is a tuple of the form `(values, as_node, task_id)` where task_id is optional.
Raises:
ValueError: If no checkpointer is set or no updates are provided.
@@ -1421,7 +1421,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
self.channels,
checkpoint,
)
values, as_node = updates[0]
values, as_node = updates[0][:2]
# no values as END, just clear all tasks
if values is None and as_node == END:
@@ -1443,9 +1443,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
step + 3,
for_execution=True,
store=self.store,
checkpointer=self.checkpointer
if isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
checkpointer=checkpointer,
manager=None,
)
# apply null writes
@@ -1455,10 +1453,10 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
if w[0] == NULL_TASK_ID
]:
apply_writes(
saved.checkpoint,
checkpoint,
channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
None,
checkpointer.get_next_version,
self.trigger_to_nodes,
)
# apply writes from tasks that already ran
@@ -1473,19 +1471,22 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
checkpoint,
channels,
next_tasks.values(),
None,
checkpointer.get_next_version,
self.trigger_to_nodes,
)
# save checkpoint
next_config = checkpointer.put(
checkpoint_config,
create_checkpoint(checkpoint, None, step),
create_checkpoint(checkpoint, channels, step),
{
"source": "update",
"step": step + 1,
"parents": saved.metadata.get("parents", {}) if saved else {},
},
{},
get_new_channel_versions(
checkpoint_previous_versions,
checkpoint["channel_versions"],
),
)
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
@@ -1544,28 +1545,87 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
f"Received no input writes for {self.input_channels}"
)
# no values, copy checkpoint
if values is None and as_node == "__copy__":
# copy checkpoint
if as_node == "__copy__":
if len(updates) > 1:
raise InvalidUpdateError(
"Cannot copy checkpoint with multiple updates"
)
if saved is None:
raise InvalidUpdateError("Cannot copy a non-existent checkpoint")
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
next_config = checkpointer.put(
saved.parent_config or saved.config if saved else checkpoint_config,
saved.parent_config
or patch_configurable(
saved.config, {CONFIG_KEY_CHECKPOINT_ID: None}
),
next_checkpoint,
{
"source": "fork",
"step": step + 1,
"parents": saved.metadata.get("parents", {}) if saved else {},
"parents": saved.metadata.get("parents", {}),
},
{},
)
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
# we want to both clone a checkpoint and update state in one go.
# reuse the same task ID if possible.
if isinstance(values, list) and len(values) > 0:
# figure out the task IDs for the next update checkpoint
next_tasks = prepare_next_tasks(
next_checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
managed,
next_config,
step + 2,
step + 4,
for_execution=True,
store=self.store,
checkpointer=checkpointer,
manager=None,
)
tasks_group_by = defaultdict(list)
user_group_by: dict[str, list[StateUpdate]] = defaultdict(list)
for task in next_tasks.values():
tasks_group_by[task.name].append(task.id)
for item in values:
if not isinstance(item, Sequence):
raise InvalidUpdateError(
f"Invalid update item: {item} when copying checkpoint"
)
values, as_node = item[:2]
user_group = user_group_by[as_node]
tasks_group = tasks_group_by[as_node]
target_idx = len(user_group)
task_id = (
tasks_group[target_idx]
if target_idx < len(tasks_group)
else None
)
user_group_by[as_node].append(
StateUpdate(values=values, as_node=as_node, task_id=task_id)
)
return perform_superstep(
patch_checkpoint_map(next_config, saved.metadata),
[item for lst in user_group_by.values() for item in lst],
)
return patch_checkpoint_map(next_config, saved.metadata)
# apply pending writes, if not on specific checkpoint
if (
CONFIG_KEY_CHECKPOINT_ID not in config[CONF]
@@ -1584,11 +1644,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
step + 3,
for_execution=True,
store=self.store,
checkpointer=(
self.checkpointer
if isinstance(self.checkpointer, BaseCheckpointSaver)
else None
),
checkpointer=checkpointer,
manager=None,
)
# apply null writes
@@ -1596,10 +1652,10 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
]:
apply_writes(
saved.checkpoint,
checkpoint,
channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
None,
checkpointer.get_next_version,
self.trigger_to_nodes,
)
# apply writes
@@ -1611,11 +1667,15 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
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, self.trigger_to_nodes
checkpoint,
channels,
tasks,
checkpointer.get_next_version,
self.trigger_to_nodes,
)
valid_updates: list[tuple[str, dict[str, Any] | None]] = []
valid_updates: list[tuple[str, dict[str, Any] | None, str | None]] = []
if len(updates) == 1:
values, as_node = updates[0]
values, as_node, task_id = updates[0]
# find last node that updated the state, if not provided
if as_node is None and len(self.nodes) == 1:
as_node = tuple(self.nodes)[0]
@@ -1646,9 +1706,9 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
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((as_node, values))
valid_updates.append((as_node, values, task_id))
else:
for values, as_node in updates:
for values, as_node, task_id in updates:
if as_node is None:
raise InvalidUpdateError(
"as_node is required when applying multiple updates"
@@ -1656,19 +1716,21 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
if as_node not in self.nodes:
raise InvalidUpdateError(f"Node {as_node} does not exist")
valid_updates.append((as_node, values))
valid_updates.append((as_node, values, task_id))
run_tasks: list[PregelTaskWrites] = []
run_task_ids: list[str] = []
for as_node, values in valid_updates:
for as_node, values, provided_task_id 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))
task_id = provided_task_id or str(
uuid5(UUID(checkpoint["id"]), INTERRUPT)
)
run_tasks.append(task)
run_task_ids.append(task_id)
run = RunnableSequence(*writers) if len(writers) > 1 else writers[0]
@@ -1681,6 +1743,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_TASK_ID: task_id,
CONFIG_KEY_READ: partial(
local_read,
_scratchpad(
@@ -1750,7 +1813,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
Args:
config: The config to apply the updates to.
supersteps: A list of supersteps, each including a list of updates to apply sequentially to a graph state.
Each update is a tuple of the form `(values, as_node)`.
Each update is a tuple of the form `(values, as_node, task_id)` where task_id is optional.
Raises:
ValueError: If no checkpointer is set or no updates are provided.
@@ -1817,7 +1880,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
self.channels,
checkpoint,
)
values, as_node = updates[0]
values, as_node = updates[0][:2]
# no values, just clear all tasks
if values is None and as_node == END:
if len(updates) > 1:
@@ -1837,9 +1900,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
step + 3,
for_execution=True,
store=self.store,
checkpointer=self.checkpointer
if isinstance(self.checkpointer, BaseCheckpointSaver)
else None,
checkpointer=checkpointer,
manager=None,
)
# apply null writes
@@ -1849,10 +1910,10 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
if w[0] == NULL_TASK_ID
]:
apply_writes(
saved.checkpoint,
checkpoint,
channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
None,
checkpointer.get_next_version,
self.trigger_to_nodes,
)
# apply writes from tasks that already ran
@@ -1867,19 +1928,21 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
checkpoint,
channels,
next_tasks.values(),
None,
checkpointer.get_next_version,
self.trigger_to_nodes,
)
# save checkpoint
next_config = await checkpointer.aput(
checkpoint_config,
create_checkpoint(checkpoint, None, step),
create_checkpoint(checkpoint, channels, step),
{
"source": "update",
"step": step + 1,
"parents": saved.metadata.get("parents", {}) if saved else {},
},
{},
get_new_channel_versions(
checkpoint_previous_versions, checkpoint["channel_versions"]
),
)
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
@@ -1939,24 +2002,83 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
)
# no values, copy checkpoint
if values is None and as_node == "__copy__":
if as_node == "__copy__":
if len(updates) > 1:
raise InvalidUpdateError(
"Cannot copy checkpoint with multiple updates"
)
if saved is None:
raise InvalidUpdateError("Cannot copy a non-existent checkpoint")
next_checkpoint = create_checkpoint(checkpoint, None, step)
# copy checkpoint
next_config = await checkpointer.aput(
saved.parent_config or saved.config if saved else checkpoint_config,
saved.parent_config
or patch_configurable(
saved.config, {CONFIG_KEY_CHECKPOINT_ID: None}
),
next_checkpoint,
{
"source": "fork",
"step": step + 1,
"parents": saved.metadata.get("parents", {}) if saved else {},
"parents": saved.metadata.get("parents", {}),
},
{},
)
# we want to both clone a checkpoint and update state in one go.
# reuse the same task ID if possible.
if isinstance(values, list) and len(values) > 0:
# figure out the task IDs for the next update checkpoint
next_tasks = prepare_next_tasks(
next_checkpoint,
saved.pending_writes or [],
self.nodes,
channels,
managed,
next_config,
step + 2,
step + 4,
for_execution=True,
store=self.store,
checkpointer=checkpointer,
manager=None,
)
tasks_group_by = defaultdict(list)
user_group_by: dict[str, list[StateUpdate]] = defaultdict(list)
for task in next_tasks.values():
tasks_group_by[task.name].append(task.id)
for item in values:
if not isinstance(item, Sequence):
raise InvalidUpdateError(
f"Invalid update item: {item} when copying checkpoint"
)
values, as_node = item[:2]
user_group = user_group_by[as_node]
tasks_group = tasks_group_by[as_node]
target_idx = len(user_group)
task_id = (
tasks_group[target_idx]
if target_idx < len(tasks_group)
else None
)
user_group_by[as_node].append(
StateUpdate(values=values, as_node=as_node, task_id=task_id)
)
return await aperform_superstep(
patch_checkpoint_map(next_config, saved.metadata),
[item for lst in user_group_by.values() for item in lst],
)
return patch_checkpoint_map(
next_config, saved.metadata if saved else None
)
@@ -1978,11 +2100,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
step + 3,
for_execution=True,
store=self.store,
checkpointer=(
self.checkpointer
if isinstance(self.checkpointer, BaseCheckpointSaver)
else None
),
checkpointer=checkpointer,
manager=None,
)
# apply null writes
@@ -1990,10 +2108,10 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
w[1:] for w in saved.pending_writes or [] if w[0] == NULL_TASK_ID
]:
apply_writes(
saved.checkpoint,
checkpoint,
channels,
[PregelTaskWrites((), INPUT, null_writes, [])],
None,
checkpointer.get_next_version,
self.trigger_to_nodes,
)
for tid, k, v in saved.pending_writes:
@@ -2004,11 +2122,15 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
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, self.trigger_to_nodes
checkpoint,
channels,
tasks,
checkpointer.get_next_version,
self.trigger_to_nodes,
)
valid_updates: list[tuple[str, dict[str, Any] | None]] = []
valid_updates: list[tuple[str, dict[str, Any] | None, str | None]] = []
if len(updates) == 1:
values, as_node = updates[0]
values, as_node, task_id = updates[0]
# find last node that updated the state, if not provided
if as_node is None and len(self.nodes) == 1:
as_node = tuple(self.nodes)[0]
@@ -2035,9 +2157,9 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
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((as_node, values))
valid_updates.append((as_node, values, task_id))
else:
for values, as_node in updates:
for values, as_node, task_id in updates:
if as_node is None:
raise InvalidUpdateError(
"as_node is required when applying multiple updates"
@@ -2045,19 +2167,21 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
if as_node not in self.nodes:
raise InvalidUpdateError(f"Node {as_node} does not exist")
valid_updates.append((as_node, values))
valid_updates.append((as_node, values, task_id))
run_tasks: list[PregelTaskWrites] = []
run_task_ids: list[str] = []
for as_node, values in valid_updates:
for as_node, values, provided_task_id 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))
task_id = provided_task_id or str(
uuid5(UUID(checkpoint["id"]), INTERRUPT)
)
run_tasks.append(task)
run_task_ids.append(task_id)
run = RunnableSequence(*writers) if len(writers) > 1 else writers[0]
@@ -2070,6 +2194,7 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: writes.extend,
CONFIG_KEY_TASK_ID: task_id,
CONFIG_KEY_READ: partial(
local_read,
_scratchpad(
@@ -2136,24 +2261,28 @@ class Pregel(PregelProtocol[StateT, InputT, OutputT], Generic[StateT, InputT, Ou
config: RunnableConfig,
values: dict[str, Any] | Any | None,
as_node: str | None = None,
task_id: str | None = 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, [[StateUpdate(values, as_node)]])
return self.bulk_update_state(config, [[StateUpdate(values, as_node, task_id)]])
async def aupdate_state(
self,
config: RunnableConfig,
values: dict[str, Any] | Any,
as_node: str | None = None,
task_id: str | None = None,
) -> RunnableConfig:
"""Asynchronously 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 await self.abulk_update_state(config, [[StateUpdate(values, as_node)]])
return await self.abulk_update_state(
config, [[StateUpdate(values, as_node, task_id)]]
)
def _defaults(
self,
libs/langgraph/langgraph/types.py
+1
View File
@@ -164,6 +164,7 @@ class Interrupt:
class StateUpdate(NamedTuple):
values: dict[str, Any] | None
as_node: str | None = None
task_id: str | None = None
class PregelTask(NamedTuple):
libs/langgraph/tests/test_large_cases.py
+8 -22
View File
@@ -4301,7 +4301,7 @@ def test_dynamic_interrupt(sync_checkpointer: BaseCheckpointSaver) -> None:
)
def test_copy_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None:
def test_partial_pending_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
@@ -4450,27 +4450,13 @@ def test_copy_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None:
)
# clear the interrupt and next tasks
tool_two.update_state(thread1, None, as_node="__copy__")
# interrupt is cleared, next task is kept
tool_two.update_state(thread1, None, as_node=END)
# interrupt and unresolved tasks are cleared, finished tasks are kept
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=(
"tool_one",
"tool_two",
),
tasks=(
PregelTask(
id=AnyStr(),
name="tool_one",
path=("__pregel_push", 0, False),
),
PregelTask(
AnyStr(),
"tool_two",
(PULL, "tool_two"),
interrupts=(),
),
),
values={"my_key": "value ⛰️ one", "market": "DE"},
next=(),
tasks=(),
config={
"configurable": {
"thread_id": "1",
@@ -4481,7 +4467,7 @@ def test_copy_checkpoint(sync_checkpointer: BaseCheckpointSaver) -> None:
created_at=AnyStr(),
metadata={
"parents": {},
"source": "fork",
"source": "update",
"step": 1,
},
parent_config=([*tool_two.checkpointer.list(thread1, limit=2)][-1].config),
libs/langgraph/tests/test_pregel.py
+344
View File
@@ -8047,3 +8047,347 @@ def test_timeout_with_parent_command(
graph.invoke({"value": "start"}, thread1)
assert exc_info.value.args[0].goto == "test_cmd"
assert exc_info.value.args[0].update == {"key": "value"}
def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> None:
"""Test forking and updating task results with state history."""
def checkpoint(values: dict[str, Any]):
return ("checkpoint", {"values": values})
def task(name: str, result: Any):
return ("task", {"name": name, "result": result})
def get_tree(history: list[StateSnapshot]) -> list:
"""Build a tree structure from state history for comparison."""
if not history:
return []
# Build a tree structure similar to renderForks
node_map: dict[str, dict] = {}
root_nodes: list[dict] = []
# Second pass: establish parent-child relationships
for item in reversed(history):
checkpoint_id = item.config["configurable"]["checkpoint_id"]
parent_checkpoint_id = (
item.parent_config["configurable"]["checkpoint_id"]
if item.parent_config
else None
)
node_map[checkpoint_id] = {"item": item, "children": []}
parent = node_map.get(parent_checkpoint_id)
(parent["children"] if parent else root_nodes).append(
node_map[checkpoint_id]
)
def node_to_tree(node: dict) -> list:
"""Convert a node to tree structure."""
result = [
checkpoint(node["item"].values),
] + [
task(task_info.name, task_info.result)
for task_info in node["item"].tasks
]
if len(node["children"]) > 1:
branches = [node_to_tree(child) for child in node["children"]]
return result + [branches]
elif len(node["children"]) == 1:
return result + node_to_tree(node["children"][0])
else:
return result
if len(root_nodes) == 1:
# Process all root nodes
return node_to_tree(root_nodes[0])
elif len(root_nodes) > 1:
# Multiple root nodes - treat as branches
branches = [node_to_tree(node) for node in root_nodes]
return branches
else:
return []
class State(TypedDict):
name: Annotated[str, lambda a, b: " > ".join([a, b]) if a else b]
# Define the graph with a sequence of nodes
def one(state: State) -> Command:
return Command(goto=[Send("two", {})], update={"name": "one"})
def two(state: State) -> State:
return {"name": "two"}
def three(state: State) -> State:
return {"name": "three"}
graph = (
StateGraph(State)
.add_node("one", one)
.add_node("two", two)
.add_node("three", three)
.add_edge(START, "one")
.add_edge("one", "two")
.add_edge("two", "three")
.compile(checkpointer=sync_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
history: list[StateSnapshot] = []
# Initial run
graph.invoke({"name": "start"}, config)
history = list(graph.get_state_history(config))
assert get_tree(history) == [
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
]
# Update the start state
graph.invoke(
None,
graph.update_state(
history[4].config,
values=[StateUpdate(values={"name": "start*"}, as_node="__start__")],
as_node="__copy__",
),
)
history = list(graph.get_state_history(config))
assert get_tree(history) == [
[
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": ""}),
task("__start__", {"name": "start*"}),
checkpoint({"name": "start*"}),
task("one", {"name": "one"}),
checkpoint({"name": "start* > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start* > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start* > one > two > two > three"}),
],
]
# Fork from task "one"
# Start from the checkpoint that has the task "one"
assert history[3].values == {"name": "start*"}
assert len(history[3].tasks) == 1
assert history[3].tasks[0].name == "one"
graph.invoke(
None,
graph.update_state(
history[3].config,
[StateUpdate(values={"name": "one*"}, as_node="one")],
"__copy__",
),
)
history = list(graph.get_state_history(config))
assert get_tree(history) == [
[
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": ""}),
task("__start__", {"name": "start*"}),
[
[
checkpoint({"name": "start*"}),
task("one", {"name": "one"}),
checkpoint({"name": "start* > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start* > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start* > one > two > two > three"}),
],
[
checkpoint({"name": "start*"}),
task("one", {"name": "one*"}),
checkpoint({"name": "start* > one*"}),
task("two", {"name": "two"}),
checkpoint({"name": "start* > one* > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start* > one* > two > three"}),
],
],
],
]
config = {"configurable": {"thread_id": "2"}}
# Initialize the thread once again
graph.invoke({"name": "start"}, config)
history = list(graph.get_state_history(config))
# Fork from task "two"
# Start from the checkpoint that has the task "two"
assert history[2].values == {"name": "start > one"}
graph.invoke(
None,
graph.update_state(
history[2].config,
[
StateUpdate(values={"name": "two"}, as_node="two"),
StateUpdate(values={"name": "two"}, as_node="two"),
],
"__copy__",
),
)
history = list(graph.get_state_history(config))
assert get_tree(history) == [
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
[
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
],
]
# Fork task three
assert history[1].values == {"name": "start > one > two > two"}
assert len(history[1].tasks) == 1
assert history[1].tasks[0].name == "three"
graph.invoke(
None,
graph.update_state(
history[1].config,
[StateUpdate(values={"name": "three*"}, as_node="three")],
"__copy__",
),
)
history = list(graph.get_state_history(config))
assert get_tree(history) == [
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
[
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
[
[
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three*"}),
checkpoint({"name": "start > one > two > two > three*"}),
],
],
],
],
]
# Regenerate task three
assert history[3].values == {"name": "start > one > two > two"}
assert len(history[3].tasks) == 1
assert history[3].tasks[0].name == "three"
graph.invoke(None, graph.update_state(history[3].config, None, "__copy__"))
history = list(graph.get_state_history(config))
assert get_tree(history) == [
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
[
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
[
[
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three*"}),
checkpoint({"name": "start > one > two > two > three*"}),
],
[
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
],
],
],
]
libs/langgraph/tests/test_pregel_async.py
+362 -27
View File
@@ -40,7 +40,7 @@ from langgraph.checkpoint.base import (
CheckpointTuple,
)
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, PUSH, START
from langgraph.constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL, START
from langgraph.errors import InvalidUpdateError, NodeInterrupt, ParentCommand
from langgraph.func import entrypoint, task
from langgraph.graph import END, StateGraph
@@ -860,7 +860,9 @@ async def test_dynamic_interrupt_subgraph(
@NEEDS_CONTEXTVARS
async def test_copy_checkpoint(async_checkpointer: BaseCheckpointSaver) -> None:
async def test_partial_pending_checkpoint(
async_checkpointer: BaseCheckpointSaver,
) -> None:
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
@@ -1017,36 +1019,21 @@ async def test_copy_checkpoint(async_checkpointer: BaseCheckpointSaver) -> None:
)
# clear the interrupt and next tasks
await tool_two.aupdate_state(thread1, None, as_node="__copy__")
# interrupt is cleared, next task is kept
tup = await tool_two.checkpointer.aget_tuple(thread1)
await tool_two.aupdate_state(thread1, None, as_node=END)
# interrupt and next tasks are cleared, finished tasks are kept
tup_upd = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_one", "tool_two"),
tasks=(
PregelTask(
AnyStr(),
"tool_one",
(PUSH, 0, False),
result=None,
),
PregelTask(
AnyStr(),
"tool_two",
(PULL, "tool_two"),
interrupts=(),
),
),
config=tup.config,
created_at=tup.checkpoint["ts"],
values={"my_key": "value ⛰️ one", "market": "DE"},
next=(),
tasks=(),
config=tup_upd.config,
created_at=tup_upd.checkpoint["ts"],
metadata={
"parents": {},
"source": "fork",
"source": "update",
"step": 1,
},
parent_config=(
[c async for c in tool_two.checkpointer.alist(thread1, limit=2)][-1].config
),
parent_config=tup.config,
interrupts=(),
)
@@ -8770,3 +8757,351 @@ async def test_timeout_with_parent_command(
await graph.ainvoke({"value": "start"}, thread1)
assert exc_info.value.args[0].goto == "test_cmd"
assert exc_info.value.args[0].update == {"key": "value"}
async def test_fork_and_update_task_results(
async_checkpointer: BaseCheckpointSaver,
) -> None:
"""Test forking and updating task results with state history."""
def checkpoint(values: dict[str, Any]):
return ("checkpoint", {"values": values})
def task(name: str, result: Any):
return ("task", {"name": name, "result": result})
def get_tree(history: list[StateSnapshot]) -> list:
"""Build a tree structure from state history for comparison."""
if not history:
return []
# Build a tree structure similar to renderForks
node_map: dict[str, dict] = {}
root_nodes: list[dict] = []
# Second pass: establish parent-child relationships
for item in reversed(history):
checkpoint_id = item.config["configurable"]["checkpoint_id"]
parent_checkpoint_id = (
item.parent_config["configurable"]["checkpoint_id"]
if item.parent_config
else None
)
node_map[checkpoint_id] = {"item": item, "children": []}
parent = node_map.get(parent_checkpoint_id)
(parent["children"] if parent else root_nodes).append(
node_map[checkpoint_id]
)
def node_to_tree(node: dict) -> list:
"""Convert a node to tree structure."""
result = [
checkpoint(node["item"].values),
] + [
task(task_info.name, task_info.result)
for task_info in node["item"].tasks
]
if len(node["children"]) > 1:
branches = [node_to_tree(child) for child in node["children"]]
return result + [branches]
elif len(node["children"]) == 1:
return result + node_to_tree(node["children"][0])
else:
return result
if len(root_nodes) == 1:
# Process all root nodes
return node_to_tree(root_nodes[0])
elif len(root_nodes) > 1:
# Multiple root nodes - treat as branches
branches = [node_to_tree(node) for node in root_nodes]
return branches
else:
return []
class State(TypedDict):
name: Annotated[str, lambda a, b: " > ".join([a, b]) if a else b]
# Define the graph with a sequence of nodes
def one(state: State) -> Command:
return Command(goto=[Send("two", {})], update={"name": "one"})
def two(state: State) -> State:
return {"name": "two"}
def three(state: State) -> State:
return {"name": "three"}
graph = (
StateGraph(State)
.add_node("one", one)
.add_node("two", two)
.add_node("three", three)
.add_edge(START, "one")
.add_edge("one", "two")
.add_edge("two", "three")
.compile(checkpointer=async_checkpointer)
)
config = {"configurable": {"thread_id": "1"}}
history: list[StateSnapshot] = []
# Initial run
await graph.ainvoke({"name": "start"}, config)
history = [c async for c in graph.aget_state_history(config)]
assert get_tree(history) == [
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
]
# Update the start state
await graph.ainvoke(
None,
await graph.aupdate_state(
history[4].config,
values=[StateUpdate(values={"name": "start*"}, as_node="__start__")],
as_node="__copy__",
),
)
history = [c async for c in graph.aget_state_history(config)]
assert get_tree(history) == [
[
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": ""}),
task("__start__", {"name": "start*"}),
checkpoint({"name": "start*"}),
task("one", {"name": "one"}),
checkpoint({"name": "start* > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start* > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start* > one > two > two > three"}),
],
]
# Fork from task "one"
# Start from the checkpoint that has the task "one"
assert history[3].values == {"name": "start*"}
assert len(history[3].tasks) == 1
assert history[3].tasks[0].name == "one"
await graph.ainvoke(
None,
await graph.aupdate_state(
history[3].config,
[StateUpdate(values={"name": "one*"}, as_node="one")],
"__copy__",
),
)
history = [c async for c in graph.aget_state_history(config)]
assert get_tree(history) == [
[
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": ""}),
task("__start__", {"name": "start*"}),
[
[
checkpoint({"name": "start*"}),
task("one", {"name": "one"}),
checkpoint({"name": "start* > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start* > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start* > one > two > two > three"}),
],
[
checkpoint({"name": "start*"}),
task("one", {"name": "one*"}),
checkpoint({"name": "start* > one*"}),
task("two", {"name": "two"}),
checkpoint({"name": "start* > one* > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start* > one* > two > three"}),
],
],
],
]
config = {"configurable": {"thread_id": "2"}}
# Initialize the thread once again
await graph.ainvoke({"name": "start"}, config)
history = [c async for c in graph.aget_state_history(config)]
# Fork from task "two"
# Start from the checkpoint that has the task "two"
assert history[2].values == {"name": "start > one"}
await graph.ainvoke(
None,
await graph.aupdate_state(
history[2].config,
[
StateUpdate(values={"name": "two"}, as_node="two"),
StateUpdate(values={"name": "two"}, as_node="two"),
],
"__copy__",
),
)
history = [c async for c in graph.aget_state_history(config)]
assert get_tree(history) == [
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
[
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
],
]
# Fork task three
assert history[1].values == {"name": "start > one > two > two"}
assert len(history[1].tasks) == 1
assert history[1].tasks[0].name == "three"
await graph.ainvoke(
None,
await graph.aupdate_state(
history[1].config,
[StateUpdate(values={"name": "three*"}, as_node="three")],
"__copy__",
),
)
history = [c async for c in graph.aget_state_history(config)]
assert get_tree(history) == [
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
[
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
[
[
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three*"}),
checkpoint({"name": "start > one > two > two > three*"}),
],
],
],
],
]
# Regenerate task three
assert history[3].values == {"name": "start > one > two > two"}
assert len(history[3].tasks) == 1
assert history[3].tasks[0].name == "three"
await graph.ainvoke(
None, await graph.aupdate_state(history[3].config, None, "__copy__")
)
history = [c async for c in graph.aget_state_history(config)]
assert get_tree(history) == [
checkpoint({"name": ""}),
task("__start__", {"name": "start"}),
checkpoint({"name": "start"}),
task("one", {"name": "one"}),
[
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": "start > one"}),
task("two", {"name": "two"}),
task("two", {"name": "two"}),
[
[
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
[
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three*"}),
checkpoint({"name": "start > one > two > two > three*"}),
],
[
checkpoint({"name": "start > one > two > two"}),
task("three", {"name": "three"}),
checkpoint({"name": "start > one > two > two > three"}),
],
],
],
],
]