lib: Add result for each task in a checkpoint

- Note this requires disabling the optimization that avoids saving writes for the last task in a step
This commit is contained in:
Nuno Campos
2024-10-14 10:16:14 -07:00
parent 3982090c6d
commit c26bb9e156
9 changed files with 299 additions and 78 deletions
@@ -206,7 +206,9 @@ class JsonPlusSerializer(SerializerProtocol):
elif type_ == "json":
return self.loads(data_)
elif type_ == "msgpack":
return msgpack.unpackb(data_, ext_hook=_msgpack_ext_hook)
return msgpack.unpackb(
data_, ext_hook=_msgpack_ext_hook, strict_map_key=False
)
else:
raise NotImplementedError(f"Unknown serialization type: {type_}")
+12 -2
View File
@@ -485,7 +485,12 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
saved.metadata,
saved.checkpoint["ts"],
saved.parent_config,
tasks_w_writes(next_tasks.values(), saved.pending_writes, task_states),
tasks_w_writes(
next_tasks.values(),
saved.pending_writes,
task_states,
self.stream_channels_asis,
),
)
async def _aprepare_state_snapshot(
@@ -561,7 +566,12 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
saved.metadata,
saved.checkpoint["ts"],
saved.parent_config,
tasks_w_writes(next_tasks.values(), saved.pending_writes, task_states),
tasks_w_writes(
next_tasks.values(),
saved.pending_writes,
task_states,
self.stream_channels_asis,
),
)
def get_state(
+30 -2
View File
@@ -29,7 +29,7 @@ from langgraph.constants import (
NS_SEP,
TAG_HIDDEN,
)
from langgraph.pregel.io import read_channels
from langgraph.pregel.io import map_output_updates, read_channels, single
from langgraph.pregel.utils import find_subgraph_pregel
from langgraph.types import PregelExecutableTask, PregelTask, StateSnapshot
@@ -148,6 +148,7 @@ def map_debug_checkpoint(
tasks: Iterable[PregelExecutableTask],
pending_writes: list[PendingWrite],
parent_config: Optional[RunnableConfig],
output_keys: Union[str, Sequence[str]],
) -> Iterator[DebugOutputCheckpoint]:
"""Produce "checkpoint" events for stream_mode=debug."""
@@ -195,7 +196,7 @@ def map_debug_checkpoint(
"interrupts": tuple(asdict(i) for i in t.interrupts),
"state": t.state,
}
for t in tasks_w_writes(tasks, pending_writes, task_states)
for t in tasks_w_writes(tasks, pending_writes, task_states, output_keys)
],
},
}
@@ -251,6 +252,7 @@ def tasks_w_writes(
tasks: Iterable[Union[PregelTask, PregelExecutableTask]],
pending_writes: Optional[list[PendingWrite]],
states: Optional[dict[str, Union[RunnableConfig, StateSnapshot]]],
output_keys: Union[str, Sequence[str]],
) -> tuple[PregelTask, ...]:
"""Apply writes / subgraph states to tasks to be returned in a StateSnapshot."""
pending_writes = pending_writes or []
@@ -271,6 +273,32 @@ def tasks_w_writes(
v for tid, n, v in pending_writes if tid == task.id and n == INTERRUPT
),
states.get(task.id) if states else None,
(
next(
(
val
for tid, chan, val in pending_writes
if tid == task.id and chan == output_keys
),
None,
)
if isinstance(output_keys, str)
else {
chan: val
for tid, chan, val in pending_writes
if tid == task.id
and (
chan == output_keys
if isinstance(output_keys, str)
else chan in output_keys
)
}
)
if any(
w[0] == task.id and w[1] not in (ERROR, INTERRUPT)
for w in pending_writes
)
else None,
)
for task in tasks
)
+10 -4
View File
@@ -5,7 +5,7 @@ from langchain_core.runnables.utils import AddableDict
from langgraph.channels.base import BaseChannel, EmptyChannelError
from langgraph.constants import EMPTY_SEQ, ERROR, INTERRUPT, TAG_HIDDEN
from langgraph.pregel.log import logger
from langgraph.types import PregelExecutableTask
from langgraph.types import PregelExecutableTask, PregelTask
def read_channel(
@@ -99,14 +99,20 @@ class AddableUpdatesDict(AddableDict):
def map_output_updates(
output_channels: Union[str, Sequence[str]],
tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],
tasks: list[
tuple[Union[PregelTask, PregelExecutableTask], Sequence[tuple[str, Any]]]
],
cached: bool = False,
) -> Iterator[dict[str, Union[Any, dict[str, Any]]]]:
"""Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
output_tasks = [
(t, ww)
for t, ww in tasks
if (not t.config or TAG_HIDDEN not in t.config.get("tags", EMPTY_SEQ))
if (
not hasattr(t, "config")
or not t.config
or TAG_HIDDEN not in t.config.get("tags", EMPTY_SEQ)
)
and ww[0][0] != ERROR
and ww[0][0] != INTERRUPT
]
@@ -123,7 +129,7 @@ def map_output_updates(
updated = (
(
task.name,
{chan: value for chan, value in task.writes if chan in output_channels},
{chan: value for chan, value in writes if chan in output_channels},
)
for task, writes in output_tasks
if any(chan in output_channels for chan, _ in writes)
+1 -16
View File
@@ -54,7 +54,6 @@ from langgraph.constants import (
NS_SEP,
SCHEDULED,
TAG_HIDDEN,
TASKS,
)
from langgraph.errors import (
_SEEN_CHECKPOINT_NS,
@@ -256,18 +255,6 @@ class PregelLoop:
"""Put writes for a task, to be read by the next tick."""
if not writes:
return
# adjust task_writes_left
first_channel = writes[0][0]
any_channel_is_send = any(k == TASKS for k, _ in writes)
always_save = any_channel_is_send or first_channel in SPECIAL_CHANNELS
if not always_save and not self.task_writes_left:
return self._output_writes(task_id, writes)
elif first_channel == INTERRUPT:
# INTERRUPT makes us want to save the last task's writes
# so we don't decrement task_writes_left
pass
else:
self.task_writes_left -= 1
# save writes
self.checkpoint_pending_writes.extend((task_id, k, v) for k, v in writes)
if self.checkpointer_put_writes is not None:
@@ -368,9 +355,6 @@ class PregelLoop:
store=self.store,
checkpointer=self.checkpointer,
)
# we don't need to save the writes for the last task that completes
# unless in special conditions handled by self.put_writes()
self.task_writes_left = len(self.tasks) - 1
# produce debug output
if self._checkpointer_put_after_previous is not None:
@@ -386,6 +370,7 @@ class PregelLoop:
self.tasks.values(),
self.checkpoint_pending_writes,
self.prev_checkpoint_config,
self.output_keys,
)
# if no more tasks, we're done
@@ -155,6 +155,7 @@ class RemoteGraph(PregelProtocol, Runnable):
state=self._create_state_snapshot(task["state"])
if task["state"]
else None,
result=task.get("result"),
)
)
+1
View File
@@ -108,6 +108,7 @@ class PregelTask(NamedTuple):
error: Optional[Exception] = None
interrupts: tuple[Interrupt, ...] = ()
state: Union[None, RunnableConfig, "StateSnapshot"] = None
result: Optional[dict[str, Any]] = None
class PregelExecutableTask(NamedTuple):
+124 -27
View File
@@ -754,7 +754,7 @@ def test_invoke_two_processes_in_out_interrupt(
),
StateSnapshot(
values={"inbox": 4, "output": 4, "input": 3},
tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),),
tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 5}),),
next=("two",),
config={
"configurable": {
@@ -774,7 +774,7 @@ def test_invoke_two_processes_in_out_interrupt(
),
StateSnapshot(
values={"inbox": 21, "output": 4, "input": 3},
tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 4}),),
next=("one",),
config={
"configurable": {
@@ -814,7 +814,7 @@ def test_invoke_two_processes_in_out_interrupt(
),
StateSnapshot(
values={"inbox": 3, "output": 4, "input": 20},
tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 21}),),
next=("one",),
config={
"configurable": {
@@ -849,7 +849,7 @@ def test_invoke_two_processes_in_out_interrupt(
),
StateSnapshot(
values={"inbox": 3, "input": 2},
tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),),
tasks=(PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 4}),),
next=("two",),
config={
"configurable": {
@@ -869,7 +869,7 @@ def test_invoke_two_processes_in_out_interrupt(
),
StateSnapshot(
values={"input": 2},
tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
tasks=(PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 3}),),
next=("one",),
config={
"configurable": {
@@ -955,7 +955,7 @@ def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=5,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
@@ -975,7 +975,7 @@ def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=4,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
@@ -995,7 +995,7 @@ def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=3,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
@@ -1015,7 +1015,7 @@ def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=2,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
@@ -1035,7 +1035,7 @@ def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=1,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
@@ -1050,7 +1050,7 @@ def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=0,
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__"), result=1),),
next=("__start__",),
config={
"configurable": {
@@ -1488,7 +1488,7 @@ def test_pending_writes_resume(
assert state.values == {"value": 1}
assert state.next == ("one", "two")
assert state.tasks == (
PregelTask(AnyStr(), "one", (PULL, "one")),
PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}),
PregelTask(AnyStr(), "two", (PULL, "two"), 'ConnectionError("I\'m not good")'),
)
assert state.metadata == {
@@ -1670,7 +1670,11 @@ def test_pending_writes_resume(
"writes": {"__start__": {"value": 1}},
},
parent_config=None,
pending_writes=[],
pending_writes=UnsortedSequence(
(AnyStr(), "value", 1),
(AnyStr(), "start:one", "__start__"),
(AnyStr(), "start:two", "__start__"),
),
)
@@ -9159,7 +9163,14 @@ def test_nested_graph_state(
),
StateSnapshot(
values={"my_key": "my value"},
tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),),
tasks=(
PregelTask(
AnyStr(),
"outer_1",
(PULL, "outer_1"),
result={"my_key": "hi my value"},
),
),
next=("outer_1",),
config={
"configurable": {
@@ -9180,7 +9191,14 @@ def test_nested_graph_state(
),
StateSnapshot(
values={},
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
tasks=(
PregelTask(
AnyStr(),
"__start__",
(PULL, "__start__"),
result={"my_key": "my value"},
),
),
next=("__start__",),
config={
"configurable": {
@@ -9263,7 +9281,17 @@ def test_nested_graph_state(
"checkpoint_id": AnyStr(),
}
},
tasks=(PregelTask(AnyStr(), "inner_1", (PULL, "inner_1")),),
tasks=(
PregelTask(
AnyStr(),
"inner_1",
(PULL, "inner_1"),
result={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
),
),
),
StateSnapshot(
values={},
@@ -9286,7 +9314,14 @@ def test_nested_graph_state(
},
created_at=AnyStr(),
parent_config=None,
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
tasks=(
PregelTask(
AnyStr(),
"__start__",
(PULL, "__start__"),
result={"my_key": "hi my value"},
),
),
),
]
@@ -9354,7 +9389,14 @@ def test_nested_graph_state(
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
tasks=(PregelTask(AnyStr(), "outer_2", (PULL, "outer_2")),),
tasks=(
PregelTask(
AnyStr(),
"outer_2",
(PULL, "outer_2"),
result={"my_key": "hi my value here and there and back again"},
),
),
next=("outer_2",),
config={
"configurable": {
@@ -9388,6 +9430,7 @@ def test_nested_graph_state(
state={
"configurable": {"thread_id": "1", "checkpoint_ns": AnyStr()}
},
result={"my_key": "hi my value here and there"},
),
),
next=("inner",),
@@ -9415,7 +9458,14 @@ def test_nested_graph_state(
),
StateSnapshot(
values={"my_key": "my value"},
tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),),
tasks=(
PregelTask(
AnyStr(),
"outer_1",
(PULL, "outer_1"),
result={"my_key": "hi my value"},
),
),
next=("outer_1",),
config={
"configurable": {
@@ -9436,7 +9486,14 @@ def test_nested_graph_state(
),
StateSnapshot(
values={},
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
tasks=(
PregelTask(
AnyStr(),
"__start__",
(PULL, "__start__"),
result={"my_key": "my value"},
),
),
next=("__start__",),
config={
"configurable": {
@@ -9872,6 +9929,7 @@ def test_doubly_nested_graph_state(
id=AnyStr(),
name="parent_2",
path=(PULL, "parent_2"),
result={"my_key": "hi my value here and there and back again"},
),
),
),
@@ -9888,6 +9946,7 @@ def test_doubly_nested_graph_state(
"checkpoint_ns": AnyStr("child"),
}
},
result={"my_key": "hi my value here and there"},
),
),
next=("child",),
@@ -9932,7 +9991,14 @@ def test_doubly_nested_graph_state(
"checkpoint_id": AnyStr(),
}
},
tasks=(PregelTask(id=AnyStr(), name="parent_1", path=(PULL, "parent_1")),),
tasks=(
PregelTask(
id=AnyStr(),
name="parent_1",
path=(PULL, "parent_1"),
result={"my_key": "hi my value"},
),
),
),
StateSnapshot(
values={},
@@ -9953,7 +10019,12 @@ def test_doubly_nested_graph_state(
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
PregelTask(
id=AnyStr(),
name="__start__",
path=(PULL, "__start__"),
result={"my_key": "my value"},
),
),
),
]
@@ -10027,6 +10098,7 @@ def test_doubly_nested_graph_state(
"checkpoint_ns": AnyStr("child:"),
}
},
result={"my_key": "hi my value here and there"},
),
),
),
@@ -10052,7 +10124,12 @@ def test_doubly_nested_graph_state(
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
PregelTask(
id=AnyStr(),
name="__start__",
path=(PULL, "__start__"),
result={"my_key": "hi my value"},
),
),
),
]
@@ -10135,7 +10212,10 @@ def test_doubly_nested_graph_state(
},
tasks=(
PregelTask(
id=AnyStr(), name="grandchild_2", path=(PULL, "grandchild_2")
id=AnyStr(),
name="grandchild_2",
path=(PULL, "grandchild_2"),
result={"my_key": "hi my value here and there"},
),
),
),
@@ -10177,7 +10257,10 @@ def test_doubly_nested_graph_state(
},
tasks=(
PregelTask(
id=AnyStr(), name="grandchild_1", path=(PULL, "grandchild_1")
id=AnyStr(),
name="grandchild_1",
path=(PULL, "grandchild_1"),
result={"my_key": "hi my value here"},
),
),
),
@@ -10212,7 +10295,12 @@ def test_doubly_nested_graph_state(
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
PregelTask(
id=AnyStr(),
name="__start__",
path=(PULL, "__start__"),
result={"my_key": "hi my value"},
),
),
),
]
@@ -10492,6 +10580,7 @@ def test_send_to_nested_graphs(
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about cats - hohoho"]},
),
PregelTask(
AnyStr(),
@@ -10503,6 +10592,7 @@ def test_send_to_nested_graphs(
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about turtles - hohoho"]},
),
),
next=("generate_joke", "generate_joke"),
@@ -10525,7 +10615,14 @@ def test_send_to_nested_graphs(
),
StateSnapshot(
values={"jokes": []},
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
tasks=(
PregelTask(
AnyStr(),
"__start__",
(PULL, "__start__"),
result={"subjects": ["cats", "dogs"]},
),
),
next=("__start__",),
config={
"configurable": {
+117 -26
View File
@@ -951,7 +951,9 @@ async def test_invoke_two_processes_in_out_interrupt(
),
StateSnapshot(
values={"inbox": 4, "output": 4, "input": 3},
tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),),
tasks=(
PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 5}),
),
next=("two",),
config={
"configurable": {
@@ -971,7 +973,9 @@ async def test_invoke_two_processes_in_out_interrupt(
),
StateSnapshot(
values={"inbox": 21, "output": 4, "input": 3},
tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
tasks=(
PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 4}),
),
next=("one",),
config={
"configurable": {
@@ -1011,7 +1015,9 @@ async def test_invoke_two_processes_in_out_interrupt(
),
StateSnapshot(
values={"inbox": 3, "output": 4, "input": 20},
tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
tasks=(
PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 21}),
),
next=("one",),
config={
"configurable": {
@@ -1051,7 +1057,9 @@ async def test_invoke_two_processes_in_out_interrupt(
),
StateSnapshot(
values={"inbox": 3, "input": 2},
tasks=(PregelTask(AnyStr(), "two", (PULL, "two")),),
tasks=(
PregelTask(AnyStr(), "two", (PULL, "two"), result={"output": 4}),
),
next=("two",),
config={
"configurable": {
@@ -1071,7 +1079,9 @@ async def test_invoke_two_processes_in_out_interrupt(
),
StateSnapshot(
values={"input": 2},
tasks=(PregelTask(AnyStr(), "one", (PULL, "one")),),
tasks=(
PregelTask(AnyStr(), "one", (PULL, "one"), result={"inbox": 3}),
),
next=("one",),
config={
"configurable": {
@@ -1166,7 +1176,7 @@ async def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=5,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
@@ -1186,7 +1196,7 @@ async def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=4,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
@@ -1206,7 +1216,7 @@ async def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=3,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
@@ -1226,7 +1236,7 @@ async def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=2,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
@@ -1246,7 +1256,7 @@ async def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=1,
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one")),),
tasks=(PregelTask(AnyStr(), "add_one", (PULL, "add_one"), result=1),),
next=("add_one",),
config={
"configurable": {
@@ -1261,7 +1271,9 @@ async def test_fork_always_re_runs_nodes(
),
StateSnapshot(
values=0,
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
tasks=(
PregelTask(AnyStr(), "__start__", (PULL, "__start__"), result=1),
),
next=("__start__",),
config={
"configurable": {
@@ -1684,7 +1696,7 @@ async def test_pending_writes_resume(
assert state.values == {"value": 1}
assert state.next == ("one", "two")
assert state.tasks == (
PregelTask(AnyStr(), "one", (PULL, "one")),
PregelTask(AnyStr(), "one", (PULL, "one"), result={"value": 2}),
PregelTask(
AnyStr(),
"two",
@@ -1875,7 +1887,11 @@ async def test_pending_writes_resume(
"writes": {"__start__": {"value": 1}},
},
parent_config=None,
pending_writes=[],
pending_writes=UnsortedSequence(
(AnyStr(), "value", 1),
(AnyStr(), "start:one", "__start__"),
(AnyStr(), "start:two", "__start__"),
),
)
@@ -7793,7 +7809,14 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
),
StateSnapshot(
values={"my_key": "my value"},
tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),),
tasks=(
PregelTask(
AnyStr(),
"outer_1",
(PULL, "outer_1"),
result={"my_key": "hi my value"},
),
),
next=("outer_1",),
config={
"configurable": {
@@ -7819,7 +7842,14 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
),
StateSnapshot(
values={},
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
tasks=(
PregelTask(
AnyStr(),
"__start__",
(PULL, "__start__"),
result={"my_key": "my value"},
),
),
next=("__start__",),
config={
"configurable": {
@@ -7907,7 +7937,15 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
}
},
tasks=(
PregelTask(id=AnyStr(), name="inner_1", path=(PULL, "inner_1")),
PregelTask(
id=AnyStr(),
name="inner_1",
path=(PULL, "inner_1"),
result={
"my_key": "hi my value here",
"my_other_key": "hi my value",
},
),
),
),
StateSnapshot(
@@ -7932,7 +7970,12 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
PregelTask(
id=AnyStr(),
name="__start__",
path=(PULL, "__start__"),
result={"my_key": "hi my value"},
),
),
),
]
@@ -8003,7 +8046,14 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
),
StateSnapshot(
values={"my_key": "hi my value here and there"},
tasks=(PregelTask(AnyStr(), "outer_2", (PULL, "outer_2")),),
tasks=(
PregelTask(
AnyStr(),
"outer_2",
(PULL, "outer_2"),
result={"my_key": "hi my value here and there and back again"},
),
),
next=("outer_2",),
config={
"configurable": {
@@ -8040,6 +8090,7 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
"checkpoint_ns": AnyStr(),
}
},
result={"my_key": "hi my value here and there"},
),
),
next=("inner",),
@@ -8067,7 +8118,14 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
),
StateSnapshot(
values={"my_key": "my value"},
tasks=(PregelTask(AnyStr(), "outer_1", (PULL, "outer_1")),),
tasks=(
PregelTask(
AnyStr(),
"outer_1",
(PULL, "outer_1"),
result={"my_key": "hi my value"},
),
),
next=("outer_1",),
config={
"configurable": {
@@ -8093,7 +8151,14 @@ async def test_nested_graph_state(checkpointer_name: str) -> None:
),
StateSnapshot(
values={},
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
tasks=(
PregelTask(
AnyStr(),
"__start__",
(PULL, "__start__"),
result={"my_key": "my value"},
),
),
next=("__start__",),
config={
"configurable": {
@@ -8709,6 +8774,7 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
"checkpoint_ns": AnyStr("child:"),
}
},
result={"my_key": "hi my value here and there"},
),
),
),
@@ -8734,7 +8800,12 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
PregelTask(
id=AnyStr(),
name="__start__",
path=(PULL, "__start__"),
result={"my_key": "hi my value"},
),
),
),
]
@@ -8821,7 +8892,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
},
tasks=(
PregelTask(
id=AnyStr(), name="grandchild_2", path=(PULL, "grandchild_2")
id=AnyStr(),
name="grandchild_2",
path=(PULL, "grandchild_2"),
result={"my_key": "hi my value here and there"},
),
),
),
@@ -8863,7 +8937,10 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
},
tasks=(
PregelTask(
id=AnyStr(), name="grandchild_1", path=(PULL, "grandchild_1")
id=AnyStr(),
name="grandchild_1",
path=(PULL, "grandchild_1"),
result={"my_key": "hi my value here"},
),
),
),
@@ -8898,7 +8975,12 @@ async def test_doubly_nested_graph_state(checkpointer_name: str) -> None:
created_at=AnyStr(),
parent_config=None,
tasks=(
PregelTask(id=AnyStr(), name="__start__", path=(PULL, "__start__")),
PregelTask(
id=AnyStr(),
name="__start__",
path=(PULL, "__start__"),
result={"my_key": "hi my value"},
),
),
),
]
@@ -9116,6 +9198,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about cats - hohoho"]},
),
PregelTask(
AnyStr(),
@@ -9127,6 +9210,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
"checkpoint_ns": AnyStr("generate_joke:"),
}
},
result={"jokes": ["Joke about turtles - hohoho"]},
),
),
config={
@@ -9148,7 +9232,14 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
),
StateSnapshot(
values={"jokes": []},
tasks=(PregelTask(AnyStr(), "__start__", (PULL, "__start__")),),
tasks=(
PregelTask(
AnyStr(),
"__start__",
(PULL, "__start__"),
result={"subjects": ["cats", "dogs"]},
),
),
next=("__start__",),
config={
"configurable": {
@@ -9167,7 +9258,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
parent_config=None,
),
]
assert actual_history == expected_history
assert actual_history[1] == expected_history[1]
@pytest.mark.skipif(