diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 52e9d6da3..5465b2802 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -110,6 +110,8 @@ class JsonPlusSerializer(SerializerProtocol): return self._encode_constructor_args( obj.__class__, method="fromhex", args=[obj.hex()] ) + elif isinstance(obj, BaseException): + return self._encode_constructor_args(obj.__class__, args=obj.args) else: raise TypeError( f"Object of type {obj.__class__.__name__} is not JSON serializable" @@ -131,9 +133,16 @@ class JsonPlusSerializer(SerializerProtocol): # Instantiate class if value["method"] is not None: method = getattr(cls, value["method"]) - return method(*value["args"], **value["kwargs"]) else: - return cls(*value["args"], **value["kwargs"]) + method = cls + if value["args"] and value["kwargs"]: + return method(*value["args"], **value["kwargs"]) + elif value["args"]: + return method(*value["args"]) + elif value["kwargs"]: + return method(**value["kwargs"]) + else: + return method() except (ImportError, AttributeError): return None diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index ae447d62e..af4ba4d8d 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -6,9 +6,11 @@ CONFIG_KEY_READ = "__pregel_read" CONFIG_KEY_CHECKPOINTER = "__pregel_checkpointer" CONFIG_KEY_RESUMING = "__pregel_resuming" INTERRUPT = "__interrupt__" +ERROR = "__error__" TASKS = "__pregel_tasks" RESERVED = { INTERRUPT, + ERROR, TASKS, CONFIG_KEY_SEND, CONFIG_KEY_READ, diff --git a/libs/langgraph/langgraph/managed/base.py b/libs/langgraph/langgraph/managed/base.py index 0e1476895..0455ed58b 100644 --- a/libs/langgraph/langgraph/managed/base.py +++ b/libs/langgraph/langgraph/managed/base.py @@ -3,7 +3,6 @@ from abc import ABC, abstractmethod from contextlib import AsyncExitStack, ExitStack, asynccontextmanager, contextmanager from inspect import isclass from typing import ( - TYPE_CHECKING, Any, AsyncGenerator, Generator, @@ -17,9 +16,6 @@ from typing import ( from langchain_core.runnables import RunnableConfig from typing_extensions import Self, TypeGuard -if TYPE_CHECKING: - from langgraph.pregel.types import PregelTaskDescription - V = TypeVar("V") @@ -60,7 +56,7 @@ class ManagedValue(ABC, Generic[V]): pass @abstractmethod - def __call__(self, step: int, task: "PregelTaskDescription") -> V: + def __call__(self, step: int) -> V: ... diff --git a/libs/langgraph/langgraph/managed/is_last_step.py b/libs/langgraph/langgraph/managed/is_last_step.py index 5860a9cbe..3c6daf5f4 100644 --- a/libs/langgraph/langgraph/managed/is_last_step.py +++ b/libs/langgraph/langgraph/managed/is_last_step.py @@ -1,11 +1,10 @@ from typing import Annotated from langgraph.managed.base import ManagedValue -from langgraph.pregel.types import PregelExecutableTask class IsLastStepManager(ManagedValue[bool]): - def __call__(self, step: int, task: PregelExecutableTask) -> bool: + def __call__(self, step: int) -> bool: return step == self.config["recursion_limit"] - 1 diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 939cbc256..7b2fdb5f9 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -68,6 +68,7 @@ from langgraph.constants import ( CONFIG_KEY_READ, CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, + ERROR, INTERRUPT, ) from langgraph.errors import GraphRecursionError, InvalidUpdateError @@ -87,6 +88,7 @@ from langgraph.pregel.debug import ( print_step_checkpoint, print_step_tasks, print_step_writes, + tasks_w_writes, ) from langgraph.pregel.io import ( map_output_updates, @@ -376,7 +378,7 @@ class Pregel( channels, managed, config, - -1, + saved.metadata.get("step", -1) + 1 if saved else -1, for_execution=False, ) return StateSnapshot( @@ -386,6 +388,7 @@ class Pregel( saved.metadata if saved else None, saved.checkpoint["ts"] if saved else None, saved.parent_config if saved else None, + tasks_w_writes(next_tasks, saved.pending_writes), ) async def aget_state(self, config: RunnableConfig) -> StateSnapshot: @@ -413,7 +416,7 @@ class Pregel( channels, managed, config, - -1, + saved.metadata.get("step", -1) + 1 if saved else -1, for_execution=False, ) return StateSnapshot( @@ -423,6 +426,7 @@ class Pregel( saved.metadata if saved else None, saved.checkpoint["ts"] if saved else None, saved.parent_config if saved else None, + tasks_w_writes(next_tasks, saved.pending_writes), ) def get_state_history( @@ -441,9 +445,13 @@ class Pregel( and signature(self.checkpointer.list).parameters.get("filter") is None ): raise ValueError("Checkpointer does not support filtering") - for config, checkpoint, metadata, parent_config, _ in self.checkpointer.list( - config, before=before, limit=limit, filter=filter - ): + for ( + config, + checkpoint, + metadata, + parent_config, + pending_writes, + ) in self.checkpointer.list(config, before=before, limit=limit, filter=filter): with ChannelsManager( { k: LastValue(None) if isinstance(c, Context) else c @@ -460,7 +468,7 @@ class Pregel( channels, managed, config, - -1, + metadata.get("step", -1) + 1, for_execution=False, ) yield StateSnapshot( @@ -470,6 +478,7 @@ class Pregel( metadata, checkpoint["ts"], parent_config, + tasks_w_writes(next_tasks, pending_writes), ) async def aget_state_history( @@ -493,7 +502,7 @@ class Pregel( checkpoint, metadata, parent_config, - _, + pending_writes, ) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter): async with AsyncChannelsManager( { @@ -511,7 +520,7 @@ class Pregel( channels, managed, config, - -1, + metadata.get("step", -1) + 1, for_execution=False, ) yield StateSnapshot( @@ -521,6 +530,7 @@ class Pregel( metadata, checkpoint["ts"], parent_config, + tasks_w_writes(next_tasks, pending_writes), ) def update_state( @@ -935,7 +945,7 @@ class Pregel( manager=run_manager, ): # debug flag - if self.debug: + if debug: print_step_checkpoint( loop.checkpoint_metadata, loop.channels, @@ -986,10 +996,9 @@ class Pregel( break # timed out for fut in done: task = futures.pop(fut) - if fut.exception() is not None: - # we got an exception, break out of while loop - # exception will be handled in panic_or_proceed - futures.clear() + if exc := _exception(fut): + # save error to checkpointer + loop.put_writes(task.id, [(ERROR, exc)]) else: # save task writes to checkpointer loop.put_writes(task.id, task.writes) @@ -1013,6 +1022,8 @@ class Pregel( else: # remove references to loop vars del fut, task + if _should_stop_others(done): + break # panic on failure or timeout _panic_or_proceed(done, inflight, loop.step) @@ -1179,7 +1190,7 @@ class Pregel( manager=run_manager, ): # debug flag - if self.debug: + if debug: print_step_checkpoint( loop.checkpoint_metadata, loop.channels, @@ -1231,10 +1242,9 @@ class Pregel( break # timed out for fut in done: task = futures.pop(fut) - if fut.exception() is not None: - # we got an exception, break out of while loop - # exception will be handled in panic_or_proceed - futures.clear() + if exc := _exception(fut): + # save error to checkpointer + loop.put_writes(task.id, [(ERROR, exc)]) else: # save task writes to checkpointer loop.put_writes(task.id, task.writes) @@ -1260,6 +1270,8 @@ class Pregel( else: # remove references to loop vars del fut, task + if _should_stop_others(done): + break # panic on failure or timeout _panic_or_proceed(done, inflight, loop.step, asyncio.TimeoutError) @@ -1403,6 +1415,31 @@ class Pregel( return chunks +def _should_stop_others( + done: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]], +) -> bool: + for fut in done: + if fut.cancelled(): + return True + if fut.exception() is not None: + # TODO don't stop others if exception is interrupt + return True + else: + return False + + +def _exception( + fut: Union[concurrent.futures.Future[Any], asyncio.Task[Any]], +) -> Optional[BaseException]: + if fut.cancelled(): + if isinstance(fut, asyncio.Task): + return asyncio.CancelledError() + else: + return concurrent.futures.CancelledError() + else: + return fut.exception() + + def _panic_or_proceed( done: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]], inflight: Union[set[concurrent.futures.Future[Any]], set[asyncio.Task[Any]]], diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 3e8b7e48b..fd725627a 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -49,7 +49,7 @@ from langgraph.managed.base import ManagedValueMapping, is_managed_value from langgraph.pregel.io import read_channel, read_channels from langgraph.pregel.log import logger from langgraph.pregel.read import PregelNode -from langgraph.pregel.types import All, PregelExecutableTask, PregelTaskDescription +from langgraph.pregel.types import All, PregelExecutableTask, PregelTask class WritesProtocol(Protocol): @@ -225,7 +225,7 @@ def prepare_next_tasks( is_resuming: bool = False, checkpointer: Literal[None] = None, manager: Literal[None] = None, -) -> list[PregelTaskDescription]: +) -> list[PregelTask]: ... @@ -258,9 +258,9 @@ def prepare_next_tasks( is_resuming: bool = False, checkpointer: Optional[BaseCheckpointSaver] = None, manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, -) -> Union[list[PregelTaskDescription], list[PregelExecutableTask]]: +) -> Union[list[PregelTask], list[PregelExecutableTask]]: parent_ns = config.get("configurable", {}).get("checkpoint_ns", "") - tasks: Union[list[PregelTaskDescription], list[PregelExecutableTask]] = [] + tasks: Union[list[PregelTask], list[PregelExecutableTask]] = [] # Consume pending packets for packet in checkpoint["pending_sends"]: if not isinstance(packet, Send): @@ -269,24 +269,25 @@ def prepare_next_tasks( if packet.node not in processes: logger.warn(f"Ignoring unknown node name {packet.node} in pending sends") continue + # create task id + triggers = [TASKS] + metadata = { + "langgraph_step": step, + "langgraph_node": packet.node, + "langgraph_triggers": triggers, + "langgraph_task_idx": len(tasks), + } + checkpoint_ns = ( + f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}" + if parent_ns + else packet.node + ) + task_id = str( + uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata))) + ) if for_execution: proc = processes[packet.node] if node := proc.get_node(): - triggers = [TASKS] - metadata = { - "langgraph_step": step, - "langgraph_node": packet.node, - "langgraph_triggers": triggers, - "langgraph_task_idx": len(tasks), - } - checkpoint_ns = ( - f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}" - if parent_ns - else packet.node - ) - task_id = str( - uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata))) - ) writes = deque() tasks.append( PregelExecutableTask( @@ -328,7 +329,7 @@ def prepare_next_tasks( ) ) else: - tasks.append(PregelTaskDescription(packet.node)) + tasks.append(PregelTask(task_id, packet.node)) # Check if any processes should be run in next step # If so, prepare the values to be passed to them version_type = type(next(iter(checkpoint["channel_versions"].values()), None)) @@ -356,26 +357,27 @@ def prepare_next_tasks( except StopIteration: continue + # create task id + metadata = { + "langgraph_step": step, + "langgraph_node": name, + "langgraph_triggers": triggers, + "langgraph_task_idx": len(tasks), + } + checkpoint_ns = ( + f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{name}" + if parent_ns + else name + ) + task_id = str( + uuid5( + UUID(checkpoint["id"]), + json.dumps((checkpoint_ns, metadata)), + ) + ) + if for_execution: if node := proc.get_node(): - metadata = { - "langgraph_step": step, - "langgraph_node": name, - "langgraph_triggers": triggers, - "langgraph_task_idx": len(tasks), - } - checkpoint_ns = ( - f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{name}" - if parent_ns - else name - ) - task_id = str( - uuid5( - UUID(checkpoint["id"]), - json.dumps((checkpoint_ns, metadata)), - ) - ) - writes = deque() tasks.append( PregelExecutableTask( @@ -424,7 +426,7 @@ def prepare_next_tasks( ) ) else: - tasks.append(PregelTaskDescription(name)) + tasks.append(PregelTask(task_id, name)) return tasks @@ -454,9 +456,7 @@ def _proc_input( managed_values = {} for key, chan in proc.channels.items(): if is_managed_value(chan): - managed_values[key] = managed[key]( - step, PregelTaskDescription(name) - ) + managed_values[key] = managed[key](step) val.update(managed_values) except EmptyChannelError: diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index 3dbb46675..931f3a180 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -10,9 +10,9 @@ from langchain_core.utils.input import get_bolded_text, get_colored_text from langgraph.channels.base import BaseChannel from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata, PendingWrite -from langgraph.constants import TAG_HIDDEN +from langgraph.constants import ERROR, TAG_HIDDEN from langgraph.pregel.io import read_channels -from langgraph.pregel.types import PregelExecutableTask +from langgraph.pregel.types import PregelExecutableTask, PregelTask class TaskPayload(TypedDict): @@ -28,10 +28,18 @@ class TaskResultPayload(TypedDict): result: list[tuple[str, Any]] +class CheckpointTask(TypedDict): + id: str + name: str + error: Optional[str] + + class CheckpointPayload(TypedDict): config: Optional[RunnableConfig] metadata: CheckpointMetadata values: dict[str, Any] + next: list[str] + tasks: list[CheckpointTask] class DebugOutputBase(TypedDict): @@ -131,6 +139,19 @@ def map_debug_checkpoint( "values": read_channels(channels, stream_channels), "metadata": metadata, "next": [t.name for t in tasks], + "tasks": [ + { + "id": t.id, + "name": t.name, + "error": t.error, + } + if t.error + else { + "id": t.id, + "name": t.name, + } + for t in tasks_w_writes(tasks, pending_writes) + ], }, } @@ -176,3 +197,25 @@ def print_step_checkpoint( + get_bolded_text(f"State at the end of step {step}:\n") + pformat(read_channels(channels, whitelist), depth=3) ) + + +def tasks_w_writes( + tasks: list[PregelExecutableTask], + pending_writes: Optional[list[PendingWrite]], +) -> tuple[PregelTask, ...]: + return tuple( + PregelTask( + task.id, + task.name, + next( + ( + exc + for tid, n, exc in pending_writes or [] + if tid == task.id + if n == ERROR + ), + None, + ), + ) + for task in tasks + ) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index cd4bff66a..f43aa25a7 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -39,7 +39,13 @@ from langgraph.checkpoint.base import ( create_checkpoint, empty_checkpoint, ) -from langgraph.constants import CONFIG_KEY_READ, CONFIG_KEY_RESUMING, INPUT, INTERRUPT +from langgraph.constants import ( + CONFIG_KEY_READ, + CONFIG_KEY_RESUMING, + ERROR, + INPUT, + INTERRUPT, +) from langgraph.errors import EmptyInputError, GraphInterrupt from langgraph.managed.base import ( AsyncManagedValuesManager, @@ -252,6 +258,8 @@ class PregelLoop: # if there are pending writes from a previous loop, apply them if self.checkpoint_pending_writes: for tid, k, v in self.checkpoint_pending_writes: + if k == ERROR: # TODO same for INTERRUPT + continue if task := next((t for t in self.tasks if t.id == tid), None): task.writes.append((k, v)) diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 7784cf023..8343329a4 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -56,8 +56,10 @@ class RetryPolicy(NamedTuple): """List of exception classes that should trigger a retry, or a callable that returns True for exceptions that should trigger a retry.""" -class PregelTaskDescription(NamedTuple): +class PregelTask(NamedTuple): + id: str name: str + error: Optional[Exception] = None class PregelExecutableTask(NamedTuple): @@ -72,18 +74,22 @@ class PregelExecutableTask(NamedTuple): class StateSnapshot(NamedTuple): + """Snapshot of the state of the graph at the beginning of a step.""" + values: Union[dict[str, Any], Any] """Current values of channels""" - next: tuple[str] - """Nodes to execute in the next step, if any""" + next: tuple[str, ...] + """The name of the node to execute in each task for this step.""" config: RunnableConfig """Config used to fetch this snapshot""" metadata: Optional[CheckpointMetadata] """Metadata associated with this snapshot""" created_at: Optional[str] """Timestamp of snapshot creation""" - parent_config: Optional[RunnableConfig] = None + parent_config: Optional[RunnableConfig] """Config used to fetch the parent snapshot, if any""" + tasks: tuple[PregelTask, ...] + """Tasks to execute in this step. If already attempted, may contain an error.""" All = Literal["*"] diff --git a/libs/langgraph/tests/any_str.py b/libs/langgraph/tests/any_str.py index 0cc419e8b..9aa030d3b 100644 --- a/libs/langgraph/tests/any_str.py +++ b/libs/langgraph/tests/any_str.py @@ -4,3 +4,21 @@ class AnyStr(str): def __eq__(self, other: object) -> bool: return isinstance(other, str) + + def __hash__(self) -> int: + return hash(str(self)) + + +class ExceptionLike: + def __init__(self, exc: Exception) -> None: + self.exc = exc + + def __eq__(self, value: object) -> bool: + return ( + isinstance(value, Exception) + and self.exc.__class__ == value.__class__ + and str(self.exc) == str(value) + ) + + def __hash__(self) -> int: + return hash((self.exc.__class__, str(self.exc))) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 7bb733669..68bbc65ce 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1,5 +1,4 @@ import json -import logging import operator import time import warnings @@ -43,6 +42,7 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import ( + BaseCheckpointSaver, Checkpoint, CheckpointMetadata, CheckpointTuple, @@ -51,7 +51,7 @@ from langgraph.checkpoint.memory import MemorySaver from langgraph.checkpoint.serde.base import SerializerProtocol from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer from langgraph.checkpoint.sqlite import SqliteSaver -from langgraph.constants import Send +from langgraph.constants import ERROR, Send from langgraph.errors import InvalidUpdateError from langgraph.graph import END, Graph from langgraph.graph.graph import START @@ -63,7 +63,8 @@ from langgraph.prebuilt.chat_agent_executor import ( from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy -from tests.any_str import AnyStr +from langgraph.pregel.types import PregelTask +from tests.any_str import AnyStr, ExceptionLike from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, MemorySaverAssertImmutable, @@ -644,6 +645,7 @@ def test_invoke_two_processes_in_out_interrupt( assert history == [ StateSnapshot( values={"inbox": 4, "output": 5, "input": 3}, + tasks=(), next=(), config={ "configurable": { @@ -658,6 +660,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 4, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "two"),), next=("two",), config={ "configurable": { @@ -672,6 +675,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 21, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "one"),), next=("one",), config={ "configurable": { @@ -686,6 +690,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 21, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "two"),), next=("two",), config={ "configurable": { @@ -700,6 +705,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 3, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "one"),), next=("one",), config={ "configurable": { @@ -714,6 +720,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 3, "output": 4, "input": 2}, + tasks=(), next=(), config={ "configurable": { @@ -728,6 +735,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 3, "input": 2}, + tasks=(PregelTask(AnyStr(), "two"),), next=("two",), config={ "configurable": { @@ -742,6 +750,7 @@ def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"input": 2}, + tasks=(PregelTask(AnyStr(), "one"),), next=("one",), config={ "configurable": { @@ -814,6 +823,7 @@ def test_fork_always_re_runs_nodes( assert history == [ StateSnapshot( values=6, + tasks=(), next=(), config={ "configurable": { @@ -828,6 +838,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=5, + tasks=(PregelTask(AnyStr(), "add_one"),), next=("add_one",), config={ "configurable": { @@ -842,6 +853,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=4, + tasks=(PregelTask(AnyStr(), "add_one"),), next=("add_one",), config={ "configurable": { @@ -856,6 +868,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=3, + tasks=(PregelTask(AnyStr(), "add_one"),), next=("add_one",), config={ "configurable": { @@ -870,6 +883,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=2, + tasks=(PregelTask(AnyStr(), "add_one"),), next=("add_one",), config={ "configurable": { @@ -884,6 +898,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=1, + tasks=(PregelTask(AnyStr(), "add_one"),), next=("add_one",), config={ "configurable": { @@ -898,6 +913,7 @@ def test_fork_always_re_runs_nodes( ), StateSnapshot( values=0, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -1192,8 +1208,7 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non assert app.invoke(2) == [3, 3] -def test_invoke_checkpoint(mocker: MockerFixture, caplog) -> None: - caplog.set_level(logging.INFO) +def test_invoke_checkpoint(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"]) errored_once = False @@ -1239,16 +1254,19 @@ def test_invoke_checkpoint(mocker: MockerFixture, caplog) -> None: # total is now 2, so output is 2+3=5 assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5 assert errored_once, "errored and retried" - checkpoint = memory.get({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 + checkpoint_tup = memory.get_tuple({"configurable": {"thread_id": "1"}}) + assert checkpoint_tup is not None + assert checkpoint_tup.checkpoint["channel_values"].get("total") == 7 # total is now 2+5=7, so output would be 7+4=11, but raises ValueError with pytest.raises(ValueError): app.invoke(4, {"configurable": {"thread_id": "1"}}) - # checkpoint is not updated - checkpoint = memory.get({"configurable": {"thread_id": "1"}}) - assert checkpoint is not None - assert checkpoint["channel_values"].get("total") == 7 + # checkpoint is not updated, error is recorded + checkpoint_tup = memory.get_tuple({"configurable": {"thread_id": "1"}}) + assert checkpoint_tup is not None + assert checkpoint_tup.checkpoint["channel_values"].get("total") == 7 + assert checkpoint_tup.pending_writes == [ + (AnyStr(), ERROR, ExceptionLike(ValueError("Input is too large"))) + ] # on a new thread, total starts out as 0, so output is 0+5=5 assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5 checkpoint = memory.get({"configurable": {"thread_id": "1"}}) @@ -1266,7 +1284,9 @@ def test_invoke_checkpoint(mocker: MockerFixture, caplog) -> None: def test_pending_writes_resume( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) class State(TypedDict): value: Annotated[int, operator.add] @@ -1310,16 +1330,28 @@ def test_pending_writes_resume( assert state is not None assert state.values == {"value": 1} assert state.next == ("one", "two") + assert state.tasks == ( + PregelTask(AnyStr(), "one"), + PregelTask(AnyStr(), "two", ExceptionLike(ConnectionError("I'm not good"))), + ) assert state.metadata == {"source": "loop", "step": 0, "writes": None} # should contain pending write of "one" checkpoint = checkpointer.get_tuple(thread1) assert checkpoint is not None - assert checkpoint.pending_writes == [ + # should contain error from "two" + expected_writes = [ (AnyStr(), "one", "one"), (AnyStr(), "value", 2), + (AnyStr(), ERROR, ExceptionLike(ConnectionError("I'm not good"))), ] - # both pending writes come from same task - assert checkpoint.pending_writes[0][0] == checkpoint.pending_writes[1][0] + assert len(checkpoint.pending_writes) == 3 + assert all(w in expected_writes for w in checkpoint.pending_writes) + # both non-error pending writes come from same task + non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] + assert non_error_writes[0][0] == non_error_writes[1][0] + # error write is from the other task + error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) + assert error_write[0] != non_error_writes[0][0] # resume execution with pytest.raises(ConnectionError, match="I'm not good"): @@ -1332,7 +1364,7 @@ def test_pending_writes_resume( # confirm no new checkpoints saved state_two = graph.get_state(thread1) - assert state_two == state + assert state_two.metadata == state.metadata # resume execution, without exception two.rtn = {"value": 3} @@ -1951,6 +1983,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), }, }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], config=app_w_interrupt.checkpointer.get_tuple(config).config, @@ -2000,6 +2033,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "input": "what is weather in sf", }, }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2099,6 +2133,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), }, }, + tasks=(), next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2159,6 +2194,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), }, }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2202,6 +2238,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: "input": "what is weather in sf", }, }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2301,6 +2338,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), }, }, + tasks=(), next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2361,6 +2399,7 @@ def test_conditional_graph(snapshot: SnapshotAssertion) -> None: ), }, }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2795,6 +2834,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2835,6 +2875,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2909,6 +2950,7 @@ def test_conditional_state_graph( ) ], }, + tasks=(), next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2956,6 +2998,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -2995,6 +3038,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3067,6 +3111,7 @@ def test_conditional_state_graph( ) ], }, + tasks=(), next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3102,6 +3147,7 @@ def test_conditional_state_graph( values={ "intermediate_steps": [], }, + tasks=(PregelTask(AnyStr(), "agent"),), next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3126,6 +3172,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3178,6 +3225,7 @@ def test_conditional_state_graph( ) ], }, + tasks=(PregelTask(AnyStr(), "agent"),), next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3241,6 +3289,7 @@ def test_conditional_state_graph( ), "intermediate_steps": [], }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -3293,6 +3342,7 @@ def test_conditional_state_graph( ) ], }, + tasks=(PregelTask(AnyStr(), "agent"),), next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4056,6 +4106,7 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: ), ] }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(app_w_interrupt.checkpointer.get_tuple(config)).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], @@ -4106,6 +4157,7 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: ), ] }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], @@ -4204,6 +4256,7 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: ), ] }, + tasks=(PregelTask(AnyStr(), "tools"), PregelTask(AnyStr(), "tools")), next=("tools", "tools"), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], @@ -4267,6 +4320,7 @@ def test_state_graph_packets(serde: SerializerProtocol) -> None: AIMessage(content="answer", id="ai2"), ] }, + tasks=(), next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(app_w_interrupt.checkpointer.get_tuple(config)).checkpoint["ts"], @@ -4544,6 +4598,7 @@ def test_message_graph( id="ai1", ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4588,6 +4643,7 @@ def test_message_graph( ], ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=next_config, created_at=AnyStr(), @@ -4669,6 +4725,7 @@ def test_message_graph( id="ai2", ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4720,6 +4777,7 @@ def test_message_graph( ), AIMessage(content="answer", id="ai2"), ], + tasks=(), next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4769,6 +4827,7 @@ def test_message_graph( id="ai1", ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4813,6 +4872,7 @@ def test_message_graph( ], ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4894,6 +4954,7 @@ def test_message_graph( id="ai2", ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4945,6 +5006,7 @@ def test_message_graph( ), AIMessage(content="answer", id="ai2"), ], + tasks=(), next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -4984,6 +5046,7 @@ def test_message_graph( AIMessage(content="answer", id="ai2"), _AnyIdAIMessage(content="an extra message"), ], + tasks=(PregelTask(AnyStr(), "agent"),), next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5254,6 +5317,7 @@ def test_root_graph( id="ai1", ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5298,6 +5362,7 @@ def test_root_graph( ], ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=next_config, created_at=AnyStr(), @@ -5379,6 +5444,7 @@ def test_root_graph( id="ai2", ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5430,6 +5496,7 @@ def test_root_graph( ), AIMessage(content="answer", id="ai2"), ], + tasks=(), next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5479,6 +5546,7 @@ def test_root_graph( id="ai1", ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5523,6 +5591,7 @@ def test_root_graph( ], ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5604,6 +5673,7 @@ def test_root_graph( id="ai2", ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5655,6 +5725,7 @@ def test_root_graph( ), AIMessage(content="answer", id="ai2"), ], + tasks=(), next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5694,6 +5765,7 @@ def test_root_graph( AIMessage(content="answer", id="ai2"), _AnyIdAIMessage(content="an extra message"), ], + tasks=(PregelTask(AnyStr(), "agent"),), next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -5765,6 +5837,7 @@ def test_root_graph( _AnyIdAIMessage(content="an extra message"), ] }, + tasks=(PregelTask(AnyStr(), "agent"),), next=("agent",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -6098,6 +6171,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: ] assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6111,6 +6185,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value ⛰️ slow", "market": "DE"}, + tasks=(), next=(), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6130,6 +6205,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread2) == StateSnapshot( values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], @@ -6143,6 +6219,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread2) == StateSnapshot( values={"my_key": "value fast", "market": "US"}, + tasks=(), next=(), config=tool_two.checkpointer.get_tuple(thread2).config, created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], @@ -6162,6 +6239,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread3) == StateSnapshot( values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread3).config, created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], @@ -6172,6 +6250,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: tool_two.update_state(thread3, {"my_key": "key"}) # appends to my_key assert tool_two.get_state(thread3) == StateSnapshot( values={"my_key": "valuekey", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread3).config, created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], @@ -6189,6 +6268,7 @@ def test_start_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread3) == StateSnapshot( values={"my_key": "valuekey fast", "market": "US"}, + tasks=(), next=(), config=tool_two.checkpointer.get_tuple(thread3).config, created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], @@ -6263,6 +6343,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "writes": {"my_key": "value", "market": "DE"}, }, "next": ["__start__"], + "tasks": [{"id": AnyStr(), "name": "__start__"}], }, }, { @@ -6291,6 +6372,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "writes": None, }, "next": ["prepare"], + "tasks": [{"id": AnyStr(), "name": "prepare"}], }, }, { @@ -6340,6 +6422,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "writes": {"prepare": {"my_key": " prepared"}}, }, "next": ["tool_two_slow"], + "tasks": [{"id": AnyStr(), "name": "tool_two_slow"}], }, }, { @@ -6389,6 +6472,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "writes": {"tool_two_slow": {"my_key": " slow"}}, }, "next": ["finish"], + "tasks": [{"id": AnyStr(), "name": "finish"}], }, }, { @@ -6438,6 +6522,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "writes": {"finish": {"my_key": " finished"}}, }, "next": [], + "tasks": [], }, }, ] @@ -6458,6 +6543,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6475,6 +6561,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), next=(), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6494,6 +6581,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread2) == StateSnapshot( values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], @@ -6511,6 +6599,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread2) == StateSnapshot( values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), next=(), config=tool_two.checkpointer.get_tuple(thread2).config, created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], @@ -6539,6 +6628,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "my_key": "value prepared slow", "market": "DE", }, + tasks=(PregelTask(AnyStr(), "finish"),), next=("finish",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6557,6 +6647,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "my_key": "value prepared slower", "market": "DE", }, + tasks=(PregelTask(AnyStr(), "finish"),), next=("finish",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6585,6 +6676,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6602,6 +6694,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread1) == StateSnapshot( values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), next=(), config=tool_two.checkpointer.get_tuple(thread1).config, created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"], @@ -6621,6 +6714,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread2) == StateSnapshot( values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=tool_two.checkpointer.get_tuple(thread2).config, created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], @@ -6638,6 +6732,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread2) == StateSnapshot( values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), next=(), config=tool_two.checkpointer.get_tuple(thread2).config, created_at=tool_two.checkpointer.get_tuple(thread2).checkpoint["ts"], @@ -6655,6 +6750,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: # check current state assert tool_two.get_state(thread3) == StateSnapshot( values={"my_key": "key", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "prepare"),), next=("prepare",), config=uconfig, created_at=AnyStr(), @@ -6663,6 +6759,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: "step": 0, "writes": {START: {"my_key": "key", "market": "DE"}}, }, + parent_config=None, ) # run from this point assert tool_two.invoke(None, thread3) == { @@ -6672,6 +6769,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: # get state after first node assert tool_two.get_state(thread3) == StateSnapshot( values={"my_key": "key prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=tool_two.checkpointer.get_tuple(thread3).config, created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], @@ -6689,6 +6787,7 @@ def test_branch_then(snapshot: SnapshotAssertion) -> None: } assert tool_two.get_state(thread3) == StateSnapshot( values={"my_key": "key prepared slow finished", "market": "DE"}, + tasks=(), next=(), config=tool_two.checkpointer.get_tuple(thread3).config, created_at=tool_two.checkpointer.get_tuple(thread3).checkpoint["ts"], @@ -6805,6 +6904,7 @@ def test_in_one_fan_out_state_graph_waiting_edge(snapshot: SnapshotAssertion) -> "query": "analyzed: query: what is weather in sf", "docs": ["doc1", "doc2", "doc3", "doc4", "doc5"], }, + tasks=(PregelTask(AnyStr(), "qa"),), next=("qa",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=app_w_interrupt.checkpointer.get_tuple(config).checkpoint["ts"], @@ -7761,6 +7861,7 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -7785,6 +7886,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -7805,6 +7907,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -7828,6 +7931,7 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, + tasks=(), next=(), config={ "configurable": { @@ -7854,6 +7958,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), next=("outer_2",), config={ "configurable": { @@ -7878,6 +7983,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -7902,6 +8008,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -7922,6 +8029,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -7983,6 +8091,7 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8007,6 +8116,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -8027,6 +8137,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -8049,6 +8160,7 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8073,6 +8185,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -8093,6 +8206,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -8121,6 +8235,7 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, + tasks=(), next=(), config={ "configurable": { @@ -8147,6 +8262,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), next=("outer_2",), config={ "configurable": { @@ -8171,6 +8287,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8195,6 +8312,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -8215,6 +8333,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -8248,6 +8367,7 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8272,6 +8392,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -8292,6 +8413,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -8317,6 +8439,7 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), next=("outer_2",), config={ "configurable": { @@ -8341,6 +8464,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8365,6 +8489,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -8385,6 +8510,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -8410,6 +8536,7 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, + tasks=(), next=(), config={ "configurable": { @@ -8436,6 +8563,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), next=("outer_2",), config={ "configurable": { @@ -8460,6 +8588,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8484,6 +8613,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -8504,6 +8634,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -8532,6 +8663,7 @@ def test_nested_graph_interrupts( assert state_history == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8556,6 +8688,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -8576,6 +8709,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -8602,6 +8736,7 @@ def test_nested_graph_interrupts( assert child_state_history == [ StateSnapshot( values={"my_key": "hi my value here"}, + tasks=(), next=(), config={ "configurable": { @@ -8650,6 +8785,7 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8674,6 +8810,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8698,6 +8835,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -8718,6 +8856,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -8743,6 +8882,7 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, + tasks=(), next=(), config={ "configurable": { @@ -8769,6 +8909,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), next=("outer_2",), config={ "configurable": { @@ -8793,6 +8934,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8817,6 +8959,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8841,6 +8984,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -8861,6 +9005,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index ef49aed62..b50ecd36a 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -39,6 +39,7 @@ from langgraph.channels.last_value import LastValue from langgraph.channels.topic import Topic from langgraph.channels.untracked_value import UntrackedValue from langgraph.checkpoint.base import ( + BaseCheckpointSaver, ChannelVersions, Checkpoint, CheckpointMetadata, @@ -46,7 +47,7 @@ from langgraph.checkpoint.base import ( ) from langgraph.checkpoint.memory import MemorySaver from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver -from langgraph.constants import Send +from langgraph.constants import ERROR, Send from langgraph.errors import InvalidUpdateError from langgraph.graph import END, Graph, StateGraph from langgraph.graph.graph import START @@ -58,7 +59,8 @@ from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot from langgraph.pregel.retry import RetryPolicy -from tests.any_str import AnyStr +from langgraph.pregel.types import PregelTask +from tests.any_str import AnyStr, ExceptionLike from tests.memory_assert import ( MemorySaverAssertCheckpointMetadata, MemorySaverAssertImmutable, @@ -732,6 +734,7 @@ async def test_invoke_two_processes_in_out_interrupt( assert history == [ StateSnapshot( values={"inbox": 4, "output": 5, "input": 3}, + tasks=(), next=(), config={ "configurable": { @@ -746,6 +749,7 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 4, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "two"),), next=("two",), config={ "configurable": { @@ -760,6 +764,7 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 21, "output": 4, "input": 3}, + tasks=(PregelTask(AnyStr(), "one"),), next=("one",), config={ "configurable": { @@ -774,6 +779,7 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 21, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "two"),), next=("two",), config={ "configurable": { @@ -788,6 +794,7 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 3, "output": 4, "input": 20}, + tasks=(PregelTask(AnyStr(), "one"),), next=("one",), config={ "configurable": { @@ -802,6 +809,7 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 3, "output": 4, "input": 2}, + tasks=(), next=(), config={ "configurable": { @@ -816,6 +824,7 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"inbox": 3, "input": 2}, + tasks=(PregelTask(AnyStr(), "two"),), next=("two",), config={ "configurable": { @@ -830,6 +839,7 @@ async def test_invoke_two_processes_in_out_interrupt( ), StateSnapshot( values={"input": 2}, + tasks=(PregelTask(AnyStr(), "one"),), next=("one",), config={ "configurable": { @@ -913,6 +923,7 @@ async def test_fork_always_re_runs_nodes( StateSnapshot( values=6, next=(), + tasks=(), config={ "configurable": { "thread_id": "1", @@ -926,6 +937,7 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=5, + tasks=(PregelTask(AnyStr(), "add_one"),), next=("add_one",), config={ "configurable": { @@ -940,6 +952,7 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=4, + tasks=(PregelTask(AnyStr(), "add_one"),), next=("add_one",), config={ "configurable": { @@ -954,6 +967,7 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=3, + tasks=(PregelTask(AnyStr(), "add_one"),), next=("add_one",), config={ "configurable": { @@ -968,6 +982,7 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=2, + tasks=(PregelTask(AnyStr(), "add_one"),), next=("add_one",), config={ "configurable": { @@ -982,6 +997,7 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=1, + tasks=(PregelTask(AnyStr(), "add_one"),), next=("add_one",), config={ "configurable": { @@ -996,6 +1012,7 @@ async def test_fork_always_re_runs_nodes( ), StateSnapshot( values=0, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -1374,7 +1391,9 @@ async def test_invoke_checkpoint(mocker: MockerFixture) -> None: async def test_pending_writes_resume( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") + checkpointer: BaseCheckpointSaver = request.getfixturevalue( + f"checkpointer_{checkpointer_name}" + ) class State(TypedDict): value: Annotated[int, operator.add] @@ -1418,16 +1437,28 @@ async def test_pending_writes_resume( assert state is not None assert state.values == {"value": 1} assert state.next == ("one", "two") + assert state.tasks == ( + PregelTask(AnyStr(), "one"), + PregelTask(AnyStr(), "two", ExceptionLike(ValueError("I'm not good"))), + ) assert state.metadata == {"source": "loop", "step": 0, "writes": None} # should contain pending write of "one" checkpoint = await checkpointer.aget_tuple(thread1) assert checkpoint is not None - assert checkpoint.pending_writes == [ + # should contain error from "two" + expected_writes = [ (AnyStr(), "one", "one"), (AnyStr(), "value", 2), + (AnyStr(), ERROR, ExceptionLike(ValueError("I'm not good"))), ] - # both pending writes come from same task - assert checkpoint.pending_writes[0][0] == checkpoint.pending_writes[1][0] + assert len(checkpoint.pending_writes) == 3 + assert all(w in expected_writes for w in checkpoint.pending_writes) + # both non-error pending writes come from same task + non_error_writes = [w for w in checkpoint.pending_writes if w[1] != ERROR] + assert non_error_writes[0][0] == non_error_writes[1][0] + # error write is from the other task + error_write = next(w for w in checkpoint.pending_writes if w[1] == ERROR) + assert error_write[0] != non_error_writes[0][0] # resume execution with pytest.raises(ValueError, match="I'm not good"): @@ -1440,7 +1471,7 @@ async def test_pending_writes_resume( # confirm no new checkpoints saved state_two = await graph.aget_state(thread1) - assert state_two == state + assert state_two.metadata == state.metadata # resume execution, without exception two.rtn = {"value": 3} @@ -2115,6 +2146,7 @@ async def test_conditional_graph() -> None: ), }, }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2162,6 +2194,7 @@ async def test_conditional_graph() -> None: "input": "what is weather in sf", }, }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2265,6 +2298,7 @@ async def test_conditional_graph() -> None: ), }, }, + tasks=(), next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2332,6 +2366,7 @@ async def test_conditional_graph() -> None: ), }, }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2379,6 +2414,7 @@ async def test_conditional_graph() -> None: "input": "what is weather in sf", }, }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2482,6 +2518,7 @@ async def test_conditional_graph() -> None: ), }, }, + tasks=(), next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2549,6 +2586,7 @@ async def test_conditional_graph() -> None: ), }, }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2936,6 +2974,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: ), "intermediate_steps": [], }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -2980,6 +3019,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: ), "intermediate_steps": [], }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3058,6 +3098,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: ) ], }, + tasks=(), next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3111,6 +3152,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: ), "intermediate_steps": [], }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3154,6 +3196,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: ), "intermediate_steps": [], }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3230,6 +3273,7 @@ async def test_conditional_graph_state(mocker: MockerFixture) -> None: ) ], }, + tasks=(), next=(), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3793,6 +3837,7 @@ async def test_state_graph_packets() -> None: ), ] }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3845,6 +3890,7 @@ async def test_state_graph_packets() -> None: ), ] }, + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -3946,6 +3992,7 @@ async def test_state_graph_packets() -> None: ), ] }, + tasks=(PregelTask(AnyStr(), "tools"), PregelTask(AnyStr(), "tools")), next=("tools", "tools"), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -4010,6 +4057,7 @@ async def test_state_graph_packets() -> None: AIMessage(content="answer", id="ai2"), ] }, + tasks=(), next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -4230,6 +4278,7 @@ async def test_message_graph() -> None: id="ai1", ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=(await app_w_interrupt.checkpointer.aget_tuple(config)).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -4273,6 +4322,7 @@ async def test_message_graph() -> None: id="ai1", ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -4344,6 +4394,7 @@ async def test_message_graph() -> None: id="ai2", ), ], + tasks=(PregelTask(AnyStr(), "tools"),), next=("tools",), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -4396,6 +4447,7 @@ async def test_message_graph() -> None: ), AIMessage(content="answer", id="ai2"), ], + tasks=(), next=(), config=app_w_interrupt.checkpointer.get_tuple(config).config, created_at=(await app_w_interrupt.checkpointer.aget_tuple(config)).checkpoint[ @@ -4694,6 +4746,7 @@ async def test_start_branch_then() -> None: ] assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ @@ -4711,6 +4764,7 @@ async def test_start_branch_then() -> None: } assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value slow", "market": "DE"}, + tasks=(), next=(), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ @@ -4734,6 +4788,7 @@ async def test_start_branch_then() -> None: } assert await tool_two.aget_state(thread2) == StateSnapshot( values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ @@ -4751,6 +4806,7 @@ async def test_start_branch_then() -> None: } assert await tool_two.aget_state(thread2) == StateSnapshot( values={"my_key": "value fast", "market": "US"}, + tasks=(), next=(), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ @@ -4774,6 +4830,7 @@ async def test_start_branch_then() -> None: } assert await tool_two.aget_state(thread3) == StateSnapshot( values={"my_key": "value", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread3)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ @@ -4788,6 +4845,7 @@ async def test_start_branch_then() -> None: await tool_two.aupdate_state(thread3, {"my_key": "key"}) # appends to my_key assert await tool_two.aget_state(thread3) == StateSnapshot( values={"my_key": "valuekey", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread3)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ @@ -4809,6 +4867,7 @@ async def test_start_branch_then() -> None: } assert await tool_two.aget_state(thread3) == StateSnapshot( values={"my_key": "valuekey fast", "market": "US"}, + tasks=(), next=(), config=(await tool_two.checkpointer.aget_tuple(thread3)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ @@ -4886,6 +4945,12 @@ async def test_branch_then() -> None: "writes": {"my_key": "value", "market": "DE"}, }, "next": ["__start__"], + "tasks": [ + { + "id": AnyStr(), + "name": "__start__", + } + ], }, }, { @@ -4914,6 +4979,12 @@ async def test_branch_then() -> None: "writes": None, }, "next": ["prepare"], + "tasks": [ + { + "id": AnyStr(), + "name": "prepare", + } + ], }, }, { @@ -4963,6 +5034,12 @@ async def test_branch_then() -> None: "writes": {"prepare": {"my_key": " prepared"}}, }, "next": ["tool_two_slow"], + "tasks": [ + { + "id": AnyStr(), + "name": "tool_two_slow", + } + ], }, }, { @@ -5012,6 +5089,12 @@ async def test_branch_then() -> None: "writes": {"tool_two_slow": {"my_key": " slow"}}, }, "next": ["finish"], + "tasks": [ + { + "id": AnyStr(), + "name": "finish", + } + ], }, }, { @@ -5061,6 +5144,7 @@ async def test_branch_then() -> None: "writes": {"finish": {"my_key": " finished"}}, }, "next": [], + "tasks": [], }, }, ] @@ -5081,6 +5165,7 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ @@ -5102,6 +5187,7 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), next=(), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ @@ -5125,6 +5211,7 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread2) == StateSnapshot( values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ @@ -5146,6 +5233,7 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread2) == StateSnapshot( values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), next=(), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ @@ -5178,6 +5266,7 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ @@ -5199,6 +5288,7 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread1) == StateSnapshot( values={"my_key": "value prepared slow finished", "market": "DE"}, + tasks=(), next=(), config=(await tool_two.checkpointer.aget_tuple(thread1)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread1)).checkpoint[ @@ -5222,6 +5312,7 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread2) == StateSnapshot( values={"my_key": "value prepared", "market": "US"}, + tasks=(PregelTask(AnyStr(), "tool_two_fast"),), next=("tool_two_fast",), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ @@ -5243,6 +5334,7 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread2) == StateSnapshot( values={"my_key": "value prepared fast finished", "market": "US"}, + tasks=(), next=(), config=(await tool_two.checkpointer.aget_tuple(thread2)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread2)).checkpoint[ @@ -5266,6 +5358,7 @@ async def test_branch_then() -> None: # check current state assert await tool_two.aget_state(thread3) == StateSnapshot( values={"my_key": "key", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "prepare"),), next=("prepare",), config=uconfig, created_at=AnyStr(), @@ -5274,6 +5367,7 @@ async def test_branch_then() -> None: "step": 0, "writes": {START: {"my_key": "key", "market": "DE"}}, }, + parent_config=None, ) # run from this point assert await tool_two.ainvoke(None, thread3) == { @@ -5283,6 +5377,7 @@ async def test_branch_then() -> None: # get state after first node assert await tool_two.aget_state(thread3) == StateSnapshot( values={"my_key": "key prepared", "market": "DE"}, + tasks=(PregelTask(AnyStr(), "tool_two_slow"),), next=("tool_two_slow",), config=(await tool_two.checkpointer.aget_tuple(thread3)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ @@ -5302,6 +5397,7 @@ async def test_branch_then() -> None: } assert await tool_two.aget_state(thread3) == StateSnapshot( values={"my_key": "key prepared slow finished", "market": "DE"}, + tasks=(), next=(), config=(await tool_two.checkpointer.aget_tuple(thread3)).config, created_at=(await tool_two.checkpointer.aget_tuple(thread3)).checkpoint[ @@ -5643,6 +5739,7 @@ async def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class( "answer": "doc1,doc2,doc3,doc4", "docs": ["doc1", "doc2", "doc3", "doc4"], }, + tasks=(), next=(), config={ "configurable": { @@ -6254,6 +6351,7 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -6278,6 +6376,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -6298,6 +6397,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -6321,6 +6421,7 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, + tasks=(), next=(), config={ "configurable": { @@ -6347,6 +6448,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), next=("outer_2",), config={ "configurable": { @@ -6371,6 +6473,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -6395,6 +6498,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -6415,6 +6519,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -6481,6 +6586,7 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -6505,6 +6611,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -6525,6 +6632,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -6547,6 +6655,7 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -6571,6 +6680,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -6591,6 +6701,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -6619,6 +6730,7 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, + tasks=(), next=(), config={ "configurable": { @@ -6645,6 +6757,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), next=("outer_2",), config={ "configurable": { @@ -6669,6 +6782,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -6693,6 +6807,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -6713,6 +6828,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -6748,6 +6864,7 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -6772,6 +6889,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -6792,6 +6910,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -6817,6 +6936,7 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), next=("outer_2",), config={ "configurable": { @@ -6841,6 +6961,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -6865,6 +6986,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -6885,6 +7007,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -6910,6 +7033,7 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, + tasks=(), next=(), config={ "configurable": { @@ -6936,6 +7060,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), next=("outer_2",), config={ "configurable": { @@ -6960,6 +7085,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -6984,6 +7110,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -7004,6 +7131,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -7031,6 +7159,7 @@ async def test_nested_graph_interrupts( assert state_history == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -7055,6 +7184,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -7075,6 +7205,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -7102,6 +7233,7 @@ async def test_nested_graph_interrupts( assert child_state_history == [ StateSnapshot( values={"my_key": "hi my value here"}, + tasks=(), next=(), config={ "configurable": { @@ -7150,6 +7282,7 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -7174,6 +7307,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -7198,6 +7332,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -7218,6 +7353,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -7243,6 +7379,7 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, + tasks=(), next=(), config={ "configurable": { @@ -7269,6 +7406,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), next=("outer_2",), config={ "configurable": { @@ -7293,6 +7431,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -7317,6 +7456,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -7341,6 +7481,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -7361,6 +7502,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": {