diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index cee4e16dc..0ce8449c1 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -59,11 +59,13 @@ from langgraph.checkpoint.base import ( empty_checkpoint, ) from langgraph.constants import ( + CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_READ, CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, CONFIG_KEY_STREAM, + CONFIG_KEY_TASK_ID, ERROR, INTERRUPT, NS_END, @@ -77,6 +79,7 @@ from langgraph.pregel.algo import ( local_write, prepare_next_tasks, ) +from langgraph.pregel.config import patch_configurable from langgraph.pregel.debug import ( print_step_checkpoint, print_step_tasks, @@ -442,7 +445,19 @@ class Pregel( return StateSnapshot( read_channels(channels, self.stream_channels_asis), tuple(t.name for t in next_tasks), - saved.config, + patch_configurable( + saved.config, + { + CONFIG_KEY_CHECKPOINT_MAP: { + **saved.metadata["parents"], + saved.config["configurable"][ + "checkpoint_ns" + ]: saved.checkpoint["id"], + } + }, + ) + if saved.metadata.get("parents") + else saved.config, saved.metadata, saved.checkpoint["ts"], saved.parent_config, @@ -518,7 +533,19 @@ class Pregel( return StateSnapshot( read_channels(channels, self.stream_channels_asis), tuple(t.name for t in next_tasks), - saved.config, + patch_configurable( + saved.config, + { + CONFIG_KEY_CHECKPOINT_MAP: { + **saved.metadata["parents"], + saved.config["configurable"][ + "checkpoint_ns" + ]: saved.checkpoint["id"], + } + }, + ) + if saved.metadata.get("parents") + else saved.config, saved.metadata, saved.checkpoint["ts"], saved.parent_config, @@ -546,12 +573,9 @@ class Pregel( for name, pregel in self.get_subgraphs(recurse=True): if name == recast_checkpoint_ns: return pregel.get_state( - { - "configurable": { - **config["configurable"], - CONFIG_KEY_CHECKPOINTER: checkpointer, - } - }, + patch_configurable( + config, {CONFIG_KEY_CHECKPOINTER: checkpointer} + ), subgraphs=subgraphs, ) else: @@ -584,12 +608,9 @@ class Pregel( async for name, pregel in self.aget_subgraphs(recurse=True): if name == recast_checkpoint_ns: return await pregel.aget_state( - { - "configurable": { - **config["configurable"], - CONFIG_KEY_CHECKPOINTER: checkpointer, - } - }, + patch_configurable( + config, {CONFIG_KEY_CHECKPOINTER: checkpointer} + ), subgraphs=subgraphs, ) else: @@ -627,12 +648,9 @@ class Pregel( for name, pregel in self.get_subgraphs(recurse=True): if name == recast_checkpoint_ns: yield from pregel.get_state_history( - { - "configurable": { - **config["configurable"], - CONFIG_KEY_CHECKPOINTER: checkpointer, - } - }, + patch_configurable( + config, {CONFIG_KEY_CHECKPOINTER: checkpointer} + ), filter=filter, before=before, limit=limit, @@ -678,12 +696,9 @@ class Pregel( async for name, pregel in self.aget_subgraphs(recurse=True): if name == recast_checkpoint_ns: async for state in pregel.aget_state_history( - { - "configurable": { - **config["configurable"], - CONFIG_KEY_CHECKPOINTER: checkpointer, - } - }, + patch_configurable( + config, {CONFIG_KEY_CHECKPOINTER: checkpointer} + ), filter=filter, before=before, limit=limit, @@ -717,9 +732,32 @@ class Pregel( node `as_node`. If `as_node` is not provided, it will be set to the last node that updated the state, if not ambiguous. """ - if not self.checkpointer: + checkpointer: Optional[BaseCheckpointSaver] = config["configurable"].get( + CONFIG_KEY_CHECKPOINTER, self.checkpointer + ) + if not checkpointer: raise ValueError("No checkpointer set") + if ( + checkpoint_ns := config["configurable"].get("checkpoint_ns", "") + ) and CONFIG_KEY_CHECKPOINTER not in config["configurable"]: + # remove task_ids from checkpoint_ns + recast_checkpoint_ns = NS_SEP.join( + part.split(NS_END)[0] for part in checkpoint_ns.split(NS_SEP) + ) + # find the subgraph with the matching name + for name, pregel in self.get_subgraphs(recurse=True): + if name == recast_checkpoint_ns: + return pregel.update_state( + patch_configurable( + config, {CONFIG_KEY_CHECKPOINTER: checkpointer} + ), + values, + as_node, + ) + else: + raise ValueError(f"Subgraph {recast_checkpoint_ns} not found") + # get last checkpoint config = merge_configs(self.config, config) if self.config else config saved = self.checkpointer.get_tuple(config) @@ -729,21 +767,12 @@ class Pregel( ) step = saved.metadata.get("step", -1) if saved else -1 # merge configurable fields with previous checkpoint config - checkpoint_config = { - **config, - "configurable": { - **config["configurable"], - # TODO: add proper support for updating nested subgraph state - "checkpoint_ns": "", - }, - } + checkpoint_config = patch_configurable( + config, + {"checkpoint_ns": config["configurable"].get("checkpoint_ns", "")}, + ) if saved: - checkpoint_config = { - "configurable": { - **config.get("configurable", {}), - **saved.config["configurable"], - } - } + checkpoint_config = patch_configurable(config, saved.config["configurable"]) # find last node that updated the state, if not provided if values is None and as_node is None: return self.checkpointer.put( @@ -798,6 +827,7 @@ class Pregel( None, [INTERRUPT], None, + None, str(uuid5(UUID(checkpoint["id"]), INTERRUPT)), ) # execute task @@ -935,6 +965,7 @@ class Pregel( None, [INTERRUPT], None, + None, str(uuid5(UUID(checkpoint["id"]), INTERRUPT)), ) # execute task @@ -1016,7 +1047,7 @@ class Pregel( stream_mode = stream_mode if stream_mode is not None else self.stream_mode if not isinstance(stream_mode, list): stream_mode = [stream_mode] - if CONFIG_KEY_READ in config.get("configurable", {}): + if CONFIG_KEY_TASK_ID in config.get("configurable", {}): # if being called as a node in another graph, always use values mode stream_mode = ["values"] if CONFIG_KEY_CHECKPOINTER in config.get("configurable", {}): diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 936c609a4..e4e4b45eb 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -50,7 +50,7 @@ from langgraph.pregel.io import read_channel, read_channels from langgraph.pregel.log import logger from langgraph.pregel.manager import ChannelsManager from langgraph.pregel.read import PregelNode -from langgraph.pregel.types import All, PregelExecutableTask, PregelTask +from langgraph.pregel.types import EXACT_MATCH, All, PregelExecutableTask, PregelTask class WritesProtocol(Protocol): @@ -303,6 +303,7 @@ def prepare_next_tasks( if node := proc.get_node(): managed.replace_runtime_placeholders(step, packet.arg) writes = deque() + task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" tasks.append( PregelExecutableTask( packet.node, @@ -351,11 +352,15 @@ def prepare_next_tasks( }, CONFIG_KEY_RESUMING: is_resuming, "checkpoint_id": None, - "checkpoint_ns": f"{checkpoint_ns}:{task_id}", + "checkpoint_ns": task_checkpoint_ns, }, ), triggers, proc.retry_policy, + None + if task_checkpoint_ns + in configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}) + else EXACT_MATCH, task_id, ) ) @@ -406,6 +411,7 @@ def prepare_next_tasks( if for_execution: if node := proc.get_node(): writes = deque() + task_checkpoint_ns = f"{checkpoint_ns}:{task_id}" tasks.append( PregelExecutableTask( name, @@ -455,11 +461,15 @@ def prepare_next_tasks( parent_ns: checkpoint["id"], }, CONFIG_KEY_RESUMING: is_resuming, - "checkpoint_ns": f"{checkpoint_ns}:{task_id}", + "checkpoint_ns": task_checkpoint_ns, }, ), triggers, proc.retry_policy, + None + if task_checkpoint_ns + in configurable.get(CONFIG_KEY_CHECKPOINT_MAP, {}) + else EXACT_MATCH, task_id, ) ) diff --git a/libs/langgraph/langgraph/pregel/config.py b/libs/langgraph/langgraph/pregel/config.py new file mode 100644 index 000000000..45b2ab1a9 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/config.py @@ -0,0 +1,12 @@ +from typing import Any, Optional + +from langchain_core.runnables import RunnableConfig + + +def patch_configurable( + config: Optional[RunnableConfig], patch: dict[str, Any] +) -> RunnableConfig: + if config is None: + return {"configurable": patch} + else: + return {**config, "configurable": {**config["configurable"], **patch}} diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index fe32fe600..d408971a9 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -78,11 +78,11 @@ def map_debug_tasks( step: int, tasks: list[PregelExecutableTask] ) -> Iterator[DebugOutputTask]: ts = datetime.now(timezone.utc).isoformat() - for name, input, _, _, config, triggers, _, _ in tasks: - if config is not None and TAG_HIDDEN in config.get("tags", []): + for task in tasks: + if task.config is not None and TAG_HIDDEN in task.config.get("tags", []): continue - metadata = config["metadata"].copy() + metadata = task.config["metadata"].copy() metadata.pop("checkpoint_id", None) yield { @@ -90,10 +90,12 @@ def map_debug_tasks( "timestamp": ts, "step": step, "payload": { - "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))), - "name": name, - "input": input, - "triggers": triggers, + "id": str( + uuid5(TASK_NAMESPACE, json.dumps((task.name, step, metadata))) + ), + "name": task.name, + "input": task.input, + "triggers": task.triggers, }, } @@ -107,11 +109,11 @@ def map_debug_task_results( [stream_keys] if isinstance(stream_keys, str) else stream_keys ) ts = datetime.now(timezone.utc).isoformat() - for (name, _, _, _, config, _, _, _), writes in tasks: - if config is not None and TAG_HIDDEN in config.get("tags", []): + for task, writes in tasks: + if task.config is not None and TAG_HIDDEN in task.config.get("tags", []): continue - metadata = config["metadata"].copy() + metadata = task.config["metadata"].copy() metadata.pop("checkpoint_id", None) # TODO: make task IDs deterministic in tests and reuse task IDs for payload ID @@ -120,8 +122,10 @@ def map_debug_task_results( "timestamp": ts, "step": step, "payload": { - "id": str(uuid5(TASK_NAMESPACE, json.dumps((name, step, metadata)))), - "name": name, + "id": str( + uuid5(TASK_NAMESPACE, json.dumps((task.name, step, metadata))) + ), + "name": task.name, "error": next((w[1] for w in writes if w[0] == ERROR), None), "result": [w for w in writes if w[0] in stream_channels_list], "interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT], diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index be92c9018..6245a7ab0 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -40,9 +40,9 @@ from langgraph.checkpoint.base import ( ) from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_MAP, - CONFIG_KEY_READ, CONFIG_KEY_RESUMING, CONFIG_KEY_STREAM, + CONFIG_KEY_TASK_ID, ERROR, INPUT, INTERRUPT, @@ -60,6 +60,7 @@ from langgraph.pregel.algo import ( prepare_next_tasks, should_interrupt, ) +from langgraph.pregel.config import patch_configurable from langgraph.pregel.debug import ( map_debug_checkpoint, map_debug_task_results, @@ -178,11 +179,30 @@ class PregelLoop: self.specs = specs self.output_keys = output_keys self.stream_keys = stream_keys - self.is_nested = CONFIG_KEY_READ in self.config.get("configurable", {}) + self.is_nested = CONFIG_KEY_TASK_ID in self.config.get("configurable", {}) if CONFIG_KEY_STREAM in config["configurable"]: self.stream = DuplexStream( self.stream, config["configurable"][CONFIG_KEY_STREAM] ) + if not self.is_nested and config["configurable"].get("checkpoint_ns"): + self.config = patch_configurable( + config, {"checkpoint_ns": "", "checkpoint_id": None} + ) + if ( + CONFIG_KEY_CHECKPOINT_MAP in self.config["configurable"] + and self.config["configurable"].get("checkpoint_ns") + in self.config["configurable"][CONFIG_KEY_CHECKPOINT_MAP] + ): + self.checkpoint_config = patch_configurable( + self.config, + { + "checkpoint_id": config["configurable"][CONFIG_KEY_CHECKPOINT_MAP][ + self.config["configurable"]["checkpoint_ns"] + ] + }, + ) + else: + self.checkpoint_config = config def put_writes(self, task_id: str, writes: Sequence[tuple[str, Any]]) -> None: """Put writes for a task, to be read by the next tick.""" @@ -324,7 +344,14 @@ class PregelLoop: for tid, k, v in self.checkpoint_pending_writes: if k in (ERROR, INTERRUPT): continue - if task := next((t for t in self.tasks if t.id == tid), None): + if task := next( + ( + t + for t in self.tasks + if t.id == tid and t.cache_policy is not None + ), + None, + ): task.writes.append((k, v)) # print output for any tasks we applied previous writes to for task in self.tasks: @@ -524,7 +551,9 @@ class SyncPregelLoop(PregelLoop, ContextManager): def __enter__(self) -> Self: saved = ( - self.checkpointer.get_tuple(self.config) if self.checkpointer else None + self.checkpointer.get_tuple(self.checkpoint_config) + if self.checkpointer + else None ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, []) self.checkpoint_config = { **self.config, @@ -616,7 +645,7 @@ class AsyncPregelLoop(PregelLoop, AsyncContextManager): async def __aenter__(self) -> Self: saved = ( - await self.checkpointer.aget_tuple(self.config) + await self.checkpointer.aget_tuple(self.checkpoint_config) if self.checkpointer else None ) or CheckpointTuple(self.config, empty_checkpoint(), {"step": -2}, None, []) diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index d1ac00fd1..8a5a37b37 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -57,6 +57,15 @@ 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 CachePolicy(NamedTuple): + """Configuration for caching nodes.""" + + pass + + +EXACT_MATCH = CachePolicy() + + class PregelTask(NamedTuple): id: str name: str @@ -73,6 +82,7 @@ class PregelExecutableTask(NamedTuple): config: RunnableConfig triggers: list[str] retry_policy: Optional[RetryPolicy] + cache_policy: Optional[CachePolicy] id: str diff --git a/libs/langgraph/tests/any_str.py b/libs/langgraph/tests/any_str.py index 32383df34..9a1977a8c 100644 --- a/libs/langgraph/tests/any_str.py +++ b/libs/langgraph/tests/any_str.py @@ -1,16 +1,21 @@ -from typing import Any, Sequence +import re +from typing import Any, Sequence, Union class AnyStr(str): - def __init__(self, prefix: str = "") -> None: + def __init__(self, prefix: Union[str, re.Pattern] = "") -> None: super().__init__() self.prefix = prefix def __eq__(self, other: object) -> bool: - return isinstance(other, str) and other.startswith(self.prefix) + return isinstance(other, str) and ( + other.startswith(self.prefix) + if isinstance(self.prefix, str) + else self.prefix.match(other) + ) def __hash__(self) -> int: - return hash(str(self)) + return hash((str(self), self.prefix)) class AnyDict(dict): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 4aefed25a..d4801ffc3 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1,5 +1,6 @@ import json import operator +import re import time import warnings from collections import Counter @@ -8318,1815 +8319,6 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] -@pytest.mark.skip("TODO") -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_nested_graph_interrupts( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - - class InnerState(TypedDict): - my_key: str - my_other_key: str - - def inner_1(state: InnerState): - return { - "my_key": state["my_key"] + " here", - "my_other_key": state["my_key"], - } - - def inner_2(state: InnerState): - return { - "my_key": state["my_key"] + " and there", - "my_other_key": state["my_key"], - } - - inner = StateGraph(InnerState) - inner.add_node("inner_1", inner_1) - inner.add_node("inner_2", inner_2) - inner.add_edge("inner_1", "inner_2") - inner.set_entry_point("inner_1") - inner.set_finish_point("inner_2") - - class State(TypedDict): - my_key: str - - def outer_1(state: State): - return {"my_key": "hi " + state["my_key"]} - - def outer_2(state: State): - return {"my_key": state["my_key"] + " and back again"} - - graph = StateGraph(State) - graph.add_node("outer_1", outer_1) - graph.add_node( - "inner", - inner.compile(interrupt_before=["inner_2"]), - ) - graph.add_node("outer_2", outer_2) - graph.set_entry_point("outer_1") - graph.add_edge("outer_1", "inner") - graph.add_edge("inner", "outer_2") - graph.set_finish_point("outer_2") - - app = graph.compile(checkpointer=checkpointer) - - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert app.invoke({"my_key": "my value"}, config) == { - "my_key": "hi my value", - } - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - "inner_2", - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert app.invoke(None, config, debug=True) == { - "my_key": "hi my value here and there and back again", - } - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [*app.stream({"my_key": "my value"}, config)] == [ - {"outer_1": {"my_key": "hi my value"}}, - ] - assert [*app.stream(None, config)] == [ - {"inner": {"my_key": "hi my value here and there"}}, - {"outer_2": {"my_key": "hi my value here and there and back again"}}, - ] - - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ - { - "my_key": "my value", - }, - { - "my_key": "hi my value", - }, - ] - assert [*app.stream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", - }, - { - "my_key": "hi my value here and there and back again", - }, - ] - - # test interrupts BEFORE the node w/ interrupts - app = graph.compile(checkpointer=checkpointer, interrupt_before=["inner"]) - config = {"configurable": {"thread_id": "4"}} - assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ - { - "my_key": "my value", - }, - { - "my_key": "hi my value", - }, - ] - history = list(app.get_state_history(config)) - assert history[0] == StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=None, # no state because we haven't entered this node yet - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - assert history[1] == StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ) - assert history[2] == StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ) - - # while we're waiting for the node w/ interrupt inside to finish - assert [*app.stream(None, config, stream_mode="values")] == [] - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - "inner_2", - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [*app.stream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", - }, - { - "my_key": "hi my value here and there and back again", - }, - ] - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # test interrupts AFTER the node w/ interrupts - app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"]) - config = {"configurable": {"thread_id": "5"}} - assert [*app.stream({"my_key": "my value"}, config, stream_mode="values")] == [ - { - "my_key": "my value", - }, - { - "my_key": "hi my value", - }, - ] - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - name="inner_2", - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [*app.stream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", - }, - ] - # interrupted after "inner" - 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": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [*app.stream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there and back again", - }, - ] - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # test restarting from checkpoint_id - config = {"configurable": {"thread_id": "6"}} - app = graph.compile(checkpointer=checkpointer) - assert app.invoke({"my_key": "my value"}, config, debug=True) == { - "my_key": "hi my value" - } - state_history = [c for c in app.get_state_history(config)] - assert state_history == [ - StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - child_state_history = [ - c - for c in app.checkpointer.list( - { - "configurable": { - "thread_id": "6", - "checkpoint_ns": f"inner:{state_history[0].tasks[0].id}", - } - } - ) - ] - assert child_state_history == [ - CheckpointTuple( - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner:478a8877-4528-5ecf-9ba3-773c53e2db7f", - "checkpoint_id": "1ef64d0c-0c7a-63ee-8001-4c403fb2fdd8", - } - }, - checkpoint={ - "v": 1, - "ts": "2024-08-28T00:02:02.205896+00:00", - "id": "1ef64d0c-0c7a-63ee-8001-4c403fb2fdd8", - "channel_values": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - "inner_1": "inner_1", - }, - "channel_versions": { - "__start__": 2, - "my_key": 3, - "start:inner_1": 3, - "inner_1": 3, - "my_other_key": 3, - }, - "versions_seen": { - "__input__": {}, - "__start__": {"__start__": 1}, - "inner_1": {"start:inner_1": 2}, - }, - "pending_sends": [], - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner:478a8877-4528-5ecf-9ba3-773c53e2db7f", - "checkpoint_id": "1ef64d0c-0c78-6f26-8000-2c99f0dae586", - } - }, - pending_writes=[], - ), - CheckpointTuple( - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner:478a8877-4528-5ecf-9ba3-773c53e2db7f", - "checkpoint_id": "1ef64d0c-0c78-6f26-8000-2c99f0dae586", - } - }, - checkpoint={ - "v": 1, - "ts": "2024-08-28T00:02:02.205364+00:00", - "id": "1ef64d0c-0c78-6f26-8000-2c99f0dae586", - "channel_values": { - "my_key": "hi my value", - "start:inner_1": "__start__", - }, - "channel_versions": {"__start__": 2, "my_key": 2, "start:inner_1": 2}, - "versions_seen": {"__input__": {}, "__start__": {"__start__": 1}}, - "pending_sends": [], - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner:478a8877-4528-5ecf-9ba3-773c53e2db7f", - "checkpoint_id": "1ef64d0c-0c78-63dc-bfff-5d20aa222eff", - } - }, - pending_writes=[ - ("1defebd2-5caa-5a51-87f8-e38b7a5c16f9", "inner_1", "inner_1"), - ("1defebd2-5caa-5a51-87f8-e38b7a5c16f9", "my_key", "hi my value here"), - ("1defebd2-5caa-5a51-87f8-e38b7a5c16f9", "my_other_key", "hi my value"), - ], - ), - CheckpointTuple( - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner:478a8877-4528-5ecf-9ba3-773c53e2db7f", - "checkpoint_id": "1ef64d0c-0c78-63dc-bfff-5d20aa222eff", - } - }, - checkpoint={ - "v": 1, - "ts": "2024-08-28T00:02:02.205076+00:00", - "id": "1ef64d0c-0c78-63dc-bfff-5d20aa222eff", - "channel_values": {"__start__": {"my_key": "hi my value"}}, - "channel_versions": {"__start__": 1}, - "versions_seen": {"__input__": {}}, - "pending_sends": [], - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "hi my value"}, - "step": -1, - }, - parent_config=None, - pending_writes=[ - ("494a4aca-177b-5e8a-986c-987ff91978f5", "my_key", "hi my value"), - ("494a4aca-177b-5e8a-986c-987ff91978f5", "start:inner_1", "__start__"), - ], - ), - ] - - # check that child snapshot matches id of parent - child_snapshot = child_state_history[0] - assert ( - child_snapshot.config["configurable"]["checkpoint_id"] - == state_history[0].config["configurable"]["checkpoint_id"] - ) - # check resuming from interrupt w/ checkpoint_id - interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2] - before_interrupt_config = before_interrupt_state_snapshot.config - # going to get to interrupt again here, so the output is None - assert app.invoke(None, before_interrupt_config, debug=True) == { - "my_key": "hi my value" - } - # one more "identical" snapshot than before, at top of list - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # going to restart from interrupt - interrupt_config = interrupt_state_snapshot.config - assert app.invoke(None, interrupt_config, debug=True) == { - "my_key": "hi my value here and there and back again", - } - assert list(app.get_state_history(config)) == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - state=StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": AnyStr(), - "checkpoint_id": AnyStr(), - } - }, - ), - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - -@pytest.mark.skip("TODO") @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_nested_graph_interrupts_parallel( request: pytest.FixtureRequest, checkpointer_name: str @@ -10164,10 +8356,7 @@ def test_nested_graph_interrupts_parallel( return {"my_key": " and back again"} graph = StateGraph(State) - graph.add_node( - "inner", - inner.compile(interrupt_before=["inner_2"]), - ) + graph.add_node("inner", inner.compile(interrupt_before=["inner_2"])) graph.add_node("outer_1", outer_1) graph.add_node("outer_2", outer_2) @@ -10193,11 +8382,13 @@ def test_nested_graph_interrupts_parallel( # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) # test stream updates w/ nested interrupt config = {"configurable": {"thread_id": "2"}} - assert [*app.stream({"my_key": ""}, config)] == [ + assert [*app.stream({"my_key": ""}, config, subgraphs=True)] == [ # we got to parallel node first - {"outer_1": {"my_key": " and parallel"}}, + ((), {"outer_1": {"my_key": " and parallel"}}), + ((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}), ] assert [*app.stream(None, config)] == [ + {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, {"inner": {"my_key": "got here and there"}}, {"outer_2": {"my_key": " and back again"}}, ] @@ -10251,7 +8442,6 @@ def test_nested_graph_interrupts_parallel( ] -@pytest.mark.skip("TODO") @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) def test_doubly_nested_graph_interrupts( request: pytest.FixtureRequest, checkpointer_name: str @@ -10461,6 +8651,9 @@ def test_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -10603,6 +8796,9 @@ def test_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -10634,6 +8830,9 @@ def test_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -10660,6 +8859,9 @@ def test_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -11003,6 +9205,13 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, metadata={ @@ -11052,6 +9261,15 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr( + re.compile(r"child:.+|child1:") + ): AnyStr(), + } + ), } }, metadata={ @@ -11084,6 +9302,9 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -11316,6 +9537,9 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -11342,6 +9566,9 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -11379,6 +9606,9 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -11403,6 +9633,13 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, metadata={ @@ -11434,6 +9671,13 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, metadata={ @@ -11465,6 +9709,13 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, metadata={ @@ -11496,6 +9747,13 @@ def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, metadata={ @@ -11515,6 +9773,19 @@ def test_doubly_nested_graph_state( ), ] + # replay grandchild checkpoint + assert [ + c for c in app.stream(None, grandchild_history[2].config, subgraphs=True) + ] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + { + "grandchild_1": {"my_key": "hi my value here"}, + "__metadata__": {"cached": True}, + }, + ) + ] + @pytest.mark.skip("TODO") @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b12209977..f0c48efdc 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1,6 +1,7 @@ import asyncio import json import operator +import re from collections import Counter from contextlib import asynccontextmanager, contextmanager from typing import ( @@ -6805,1754 +6806,6 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert times_called == 1 -@pytest.mark.skip("TODO") -@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) -async def test_nested_graph_interrupts( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: - checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) - - class InnerState(TypedDict): - my_key: str - my_other_key: str - - async def inner_1(state: InnerState): - return { - "my_key": state["my_key"] + " here", - "my_other_key": state["my_key"], - } - - async def inner_2(state: InnerState): - return { - "my_key": state["my_key"] + " and there", - "my_other_key": state["my_key"], - } - - inner = StateGraph(InnerState) - inner.add_node("inner_1", inner_1) - inner.add_node("inner_2", inner_2) - inner.add_edge("inner_1", "inner_2") - inner.set_entry_point("inner_1") - inner.set_finish_point("inner_2") - - class State(TypedDict): - my_key: str - - async def outer_1(state: State): - return {"my_key": "hi " + state["my_key"]} - - async def outer_2(state: State): - return {"my_key": state["my_key"] + " and back again"} - - graph = StateGraph(State) - graph.add_node("outer_1", outer_1) - graph.add_node( - "inner", - inner.compile(interrupt_before=["inner_2"]), - ) - graph.add_node("outer_2", outer_2) - graph.set_entry_point("outer_1") - graph.add_edge("outer_1", "inner") - graph.add_edge("inner", "outer_2") - graph.set_finish_point("outer_2") - - app = graph.compile(checkpointer=checkpointer) - - # test invoke w/ nested interrupt - config = {"configurable": {"thread_id": "1"}} - assert await app.ainvoke({"my_key": "my value"}, config, debug=True) == { - "my_key": "hi my value", - } - 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": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - "inner_2", - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert await app.ainvoke(None, config, debug=True) == { - "my_key": "hi my value here and there and back again", - } - 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": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - next=(), - tasks=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # test stream updates w/ nested interrupt - config = {"configurable": {"thread_id": "2"}} - assert [c async for c in app.astream({"my_key": "my value"}, config)] == [ - {"outer_1": {"my_key": "hi my value"}}, - ] - assert [c async for c in app.astream(None, config)] == [ - {"inner": {"my_key": "hi my value here and there"}}, - {"outer_2": {"my_key": "hi my value here and there and back again"}}, - ] - - # test stream values w/ nested interrupt - config = {"configurable": {"thread_id": "3"}} - assert [ - c - async for c in app.astream({"my_key": "my value"}, config, stream_mode="values") - ] == [ - { - "my_key": "my value", - }, - { - "my_key": "hi my value", - }, - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", - }, - { - "my_key": "hi my value here and there and back again", - }, - ] - - # test interrupts BEFORE the node w/ interrupts - app = graph.compile(checkpointer=checkpointer, interrupt_before=["inner"]) - config = {"configurable": {"thread_id": "4"}} - assert [ - c - async for c in app.astream({"my_key": "my value"}, config, stream_mode="values") - ] == [ - { - "my_key": "my value", - }, - { - "my_key": "hi my value", - }, - ] - 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": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # while we're waiting for the node w/ interrupt inside to finish - assert [c async for c in app.astream(None, config, stream_mode="values")] == [] - 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": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - "inner_2", - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", - }, - { - "my_key": "hi my value here and there and back again", - }, - ] - 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": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "4", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # test interrupts AFTER the node w/ interrupts - app = graph.compile(checkpointer=checkpointer, interrupt_after=["inner"]) - config = {"configurable": {"thread_id": "5"}} - assert [ - c - async for c in app.astream({"my_key": "my value"}, config, stream_mode="values") - ] == [ - { - "my_key": "my value", - }, - { - "my_key": "hi my value", - }, - ] - 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": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - name="inner_2", - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there", - }, - ] - # interrupted after "inner" - 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": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "hi my value here and there and back again", - }, - ] - 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": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "5", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - # test restarting from checkpoint_id - config = {"configurable": {"thread_id": "6"}} - app = graph.compile(checkpointer=checkpointer) - await app.ainvoke({"my_key": "my value"}, config, debug=True) - - state_history = [c async for c in app.aget_state_history(config)] - assert state_history == [ - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "inner", - ), - ), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=( - PregelTask( - AnyStr(), - "inner_2", - ), - ), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - child_state_history = [ - c - async for c in app.aget_state_history( - {"configurable": {"thread_id": "6", "checkpoint_ns": "inner"}} - ) - ] - assert child_state_history == [ - StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ), - # there should be a single child checkpoint because we only keep - # one child checkpoint per parent checkpoint (in which child ran) - ] - - # check that child snapshot matches id of parent - child_snapshot = child_state_history[0] - assert ( - child_snapshot.config["configurable"]["checkpoint_id"] - == state_history[0].config["configurable"]["checkpoint_id"] - ) - # check resuming from interrupt w/ checkpoint_id - interrupt_state_snapshot, before_interrupt_state_snapshot = state_history[:2] - before_interrupt_config = before_interrupt_state_snapshot.config - # going to get to interrupt again here, so the output is None - assert await app.ainvoke(None, before_interrupt_config, debug=True) == { - "my_key": "hi my value" - } - # one more "identical" snapshot than before, at top of list - 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": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ), - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ), - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - # going to restart from interrupt - interrupt_config = interrupt_state_snapshot.config - assert await app.ainvoke(None, interrupt_config, debug=True) == { - "my_key": "hi my value here and there and back again", - } - 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": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "outer_2": {"my_key": "hi my value here and there and back again"} - }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - tasks=(PregelTask(AnyStr(), "outer_2"),), - next=("outer_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"inner": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, - tasks=(PregelTask(AnyStr(), "inner_2"),), - next=("inner_2",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_1": { - "my_key": "hi my value here", - "my_other_key": "hi my value", - } - }, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), - next=("inner",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"outer_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - subgraphs={ - "inner": StateSnapshot( - values={ - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - }, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "inner_2": { - "my_key": "hi my value here and there", - "my_other_key": "hi my value here", - } - }, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "inner", - "checkpoint_id": AnyStr(), - } - }, - ) - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - tasks=(PregelTask(AnyStr(), "outer_1"),), - next=("outer_1",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={"parents": {}, "source": "loop", "writes": None, "step": 0}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={}, - tasks=(PregelTask(AnyStr(), "__start__"),), - next=("__start__",), - config={ - "configurable": { - "thread_id": "6", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - ), - ] - - -@pytest.mark.skip("TODO") @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_nested_graph_interrupts_parallel( request: pytest.FixtureRequest, checkpointer_name: str @@ -8619,11 +6872,13 @@ async def test_nested_graph_interrupts_parallel( # - the writes of outer are persisted in 1st call and used in 2nd call, ie outer isn't called again (because we dont see outer_1 output again in 2nd stream) # test stream updates w/ nested interrupt config = {"configurable": {"thread_id": "2"}} - assert [c async for c in app.astream({"my_key": ""}, config)] == [ + assert [c async for c in app.astream({"my_key": ""}, config, subgraphs=True)] == [ # we got to parallel node first - {"outer_1": {"my_key": " and parallel"}}, + ((), {"outer_1": {"my_key": " and parallel"}}), + ((AnyStr("inner:"),), {"inner_1": {"my_key": "got here", "my_other_key": ""}}), ] assert [c async for c in app.astream(None, config)] == [ + {"outer_1": {"my_key": " and parallel"}, "__metadata__": {"cached": True}}, {"inner": {"my_key": "got here and there"}}, {"outer_2": {"my_key": " and back again"}}, ] @@ -8633,17 +6888,11 @@ async def test_nested_graph_interrupts_parallel( assert [ c async for c in app.astream({"my_key": ""}, config, stream_mode="values") ] == [ - { - "my_key": "", - }, + {"my_key": ""}, ] assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "got here and there and parallel", - }, - { - "my_key": "got here and there and parallel and back again", - }, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, ] # # test interrupts BEFORE the parallel node @@ -8651,16 +6900,14 @@ async def test_nested_graph_interrupts_parallel( config = {"configurable": {"thread_id": "4"}} assert [ c async for c in app.astream({"my_key": ""}, config, stream_mode="values") - ] == [{"my_key": ""}] + ] == [ + {"my_key": ""}, + ] # while we're waiting for the node w/ interrupt inside to finish assert [c async for c in app.astream(None, config, stream_mode="values")] == [] assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "got here and there and parallel", - }, - { - "my_key": "got here and there and parallel and back again", - }, + {"my_key": "got here and there and parallel"}, + {"my_key": "got here and there and parallel and back again"}, ] # test interrupts AFTER the parallel node @@ -8668,18 +6915,17 @@ async def test_nested_graph_interrupts_parallel( config = {"configurable": {"thread_id": "5"}} assert [ c async for c in app.astream({"my_key": ""}, config, stream_mode="values") - ] == [{"my_key": ""}] + ] == [ + {"my_key": ""}, + ] assert [c async for c in app.astream(None, config, stream_mode="values")] == [ {"my_key": "got here and there and parallel"}, ] assert [c async for c in app.astream(None, config, stream_mode="values")] == [ - { - "my_key": "got here and there and parallel and back again", - }, + {"my_key": "got here and there and parallel and back again"}, ] -@pytest.mark.skip("TODO") @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_doubly_nested_graph_interrupts( request: pytest.FixtureRequest, checkpointer_name: str @@ -8892,6 +7138,9 @@ async def test_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -9039,6 +7288,9 @@ async def test_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("inner:"): AnyStr()} + ), } }, metadata={ @@ -9070,6 +7322,9 @@ async def test_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("inner:"): AnyStr()} + ), } }, metadata={ @@ -9096,6 +7351,9 @@ async def test_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("inner:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("inner:"): AnyStr()} + ), } }, metadata={ @@ -9413,6 +7671,9 @@ async def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -9446,6 +7707,13 @@ async def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, metadata={ @@ -9495,6 +7763,15 @@ async def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr( + re.compile(r"child:.+|child1:") + ): AnyStr(), + } + ), } }, metadata={ @@ -9527,6 +7804,9 @@ async def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -9614,145 +7894,150 @@ async def test_doubly_nested_graph_state( ) # get outer graph history outer_history = [c async for c in app.aget_state_history(config)] - assert outer_history == [ - StateSnapshot( - values={"my_key": "hi my value here and there and back again"}, - tasks=(), - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": { - "parent_2": {"my_key": "hi my value here and there and back again"} + assert ( + outer_history[0] + == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } }, - "step": 3, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "hi my value here and there"}, - next=("parent_2",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"child": {"my_key": "hi my value here and there"}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="parent_2"),), - ), - StateSnapshot( - values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), - "child", - state={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": AnyStr("child"), + metadata={ + "parents": {}, + "source": "loop", + "writes": { + "parent_2": { + "my_key": "hi my value here and there and back again" } }, - ), + "step": 3, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, ), - next=("child",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": {"parent_1": {"my_key": "hi my value"}}, - "step": 1, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - ), - StateSnapshot( - values={"my_key": "my value"}, - next=("parent_1",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "loop", - "writes": None, - "step": 0, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - tasks=(PregelTask(id=AnyStr(), name="parent_1"),), - ), - StateSnapshot( - values={}, - next=("__start__",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "parents": {}, - "source": "input", - "writes": {"my_key": "my value"}, - "step": -1, - }, - created_at=AnyStr(), - parent_config=None, - tasks=(PregelTask(id=AnyStr(), name="__start__"),), - ), - ] + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("parent_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"child": {"my_key": "hi my value here and there"}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="parent_2"),), + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child", + state={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": AnyStr("child"), + } + }, + ), + ), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("parent_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "loop", + "writes": None, + "step": 0, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + tasks=(PregelTask(id=AnyStr(), name="parent_1"),), + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "parents": {}, + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + tasks=(PregelTask(id=AnyStr(), name="__start__"),), + ), + ][0] + ) # get child graph history child_history = [ c async for c in app.aget_state_history(outer_history[2].tasks[0].state) @@ -9766,6 +8051,9 @@ async def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -9792,6 +8080,9 @@ async def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -9829,6 +8120,9 @@ async def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr("child:"), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + {"": AnyStr(), AnyStr("child:"): AnyStr()} + ), } }, metadata={ @@ -9855,6 +8149,13 @@ async def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, metadata={ @@ -9886,6 +8187,13 @@ async def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, metadata={ @@ -9917,6 +8225,13 @@ async def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, metadata={ @@ -9948,6 +8263,13 @@ async def test_doubly_nested_graph_state( "thread_id": "1", "checkpoint_ns": AnyStr(), "checkpoint_id": AnyStr(), + "checkpoint_map": AnyDict( + { + "": AnyStr(), + AnyStr("child:"): AnyStr(), + AnyStr(re.compile(r"child:.+|child1:")): AnyStr(), + } + ), } }, metadata={ @@ -9967,6 +8289,19 @@ async def test_doubly_nested_graph_state( ), ] + # replay grandchild checkpoint + assert [ + c async for c in app.astream(None, grandchild_history[2].config, subgraphs=True) + ] == [ + ( + (AnyStr("child:"), AnyStr("child_1:")), + { + "grandchild_1": {"my_key": "hi my value here"}, + "__metadata__": {"cached": True}, + }, + ) + ] + @pytest.mark.skip("TODO") @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)