mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-07 10:17:50 +02:00
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:
@@ -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_}")
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
|
||||
)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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"),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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": {
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user