diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index f53bee256..91b49a162 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -255,9 +255,13 @@ class BasePostgresSaver(BaseCheckpointSaver): if config: wheres.append("thread_id = %s ") param_values.append(config["configurable"]["thread_id"]) - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - wheres.append("checkpoint_ns = %s") - param_values.append(checkpoint_ns) + if checkpoint_ns := config["configurable"].get("checkpoint_ns"): + wheres.append("checkpoint_ns = %s") + param_values.append(checkpoint_ns) + + if checkpoint_id := get_checkpoint_id(config): + wheres.append("checkpoint_id = %s ") + param_values.append(checkpoint_id) # construct predicate for metadata filter if filter: diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index e94cf32ae..6f9f7d78b 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -87,29 +87,15 @@ class TestAsyncPostgresSaver: search_results_4 = [c async for c in saver.alist(None, filter=query_4)] assert len(search_results_4) == 0 - # search by config (defaults to root graph checkpoints) + # search by config (defaults to checkpoints across all namespaces) search_results_5 = [ c async for c in saver.alist({"configurable": {"thread_id": "thread-2"}}) ] - assert len(search_results_5) == 1 - assert search_results_5[0].config["configurable"]["checkpoint_ns"] == "" - - # search by config and checkpoint_ns - search_results_6 = [ - c - async for c in saver.alist( - { - "configurable": { - "thread_id": "thread-2", - "checkpoint_ns": "inner", - } - } - ) - ] - assert len(search_results_6) == 1 - assert ( - search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner" - ) + assert len(search_results_5) == 2 + assert { + search_results_5[0].config["configurable"]["checkpoint_ns"], + search_results_5[1].config["configurable"]["checkpoint_ns"], + } == {"", "inner"} # TODO: test before and limit params diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index dfae82907..a2fbcbd88 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/libs/checkpoint-postgres/tests/test_sync.py @@ -88,27 +88,14 @@ class TestPostgresSaver: search_results_4 = list(saver.list(None, filter=query_4)) assert len(search_results_4) == 0 - # search by config (defaults to root graph checkpoints) + # search by config (defaults to checkpoints across all namespaces) search_results_5 = list( saver.list({"configurable": {"thread_id": "thread-2"}}) ) - assert len(search_results_5) == 1 - assert search_results_5[0].config["configurable"]["checkpoint_ns"] == "" - - # search by config and checkpoint_ns - search_results_6 = list( - saver.list( - { - "configurable": { - "thread_id": "thread-2", - "checkpoint_ns": "inner", - } - } - ) - ) - assert len(search_results_6) == 1 - assert ( - search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner" - ) + assert len(search_results_5) == 2 + assert { + search_results_5[0].config["configurable"]["checkpoint_ns"], + search_results_5[1].config["configurable"]["checkpoint_ns"], + } == {"", "inner"} # TODO: test before and limit params diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py index 6e1baf5ae..0e1e06fcc 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py @@ -70,9 +70,13 @@ def search_where( if config is not None: wheres.append("thread_id = ?") param_values.append(config["configurable"]["thread_id"]) - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - wheres.append("checkpoint_ns = ?") - param_values.append(checkpoint_ns) + if checkpoint_ns := config["configurable"].get("checkpoint_ns"): + wheres.append("checkpoint_ns = ?") + param_values.append(checkpoint_ns) + + if checkpoint_id := get_checkpoint_id(config): + wheres.append("checkpoint_id = ?") + param_values.append(checkpoint_id) # construct predicate for metadata filter if filter: diff --git a/libs/checkpoint-sqlite/tests/test_aiosqlite.py b/libs/checkpoint-sqlite/tests/test_aiosqlite.py index 59f830dae..038030172 100644 --- a/libs/checkpoint-sqlite/tests/test_aiosqlite.py +++ b/libs/checkpoint-sqlite/tests/test_aiosqlite.py @@ -84,29 +84,15 @@ class TestAsyncSqliteSaver: search_results_4 = [c async for c in saver.alist(None, filter=query_4)] assert len(search_results_4) == 0 - # search by config (defaults to root graph checkpoints) + # search by config (defaults to checkpoints across all namespaces) search_results_5 = [ c async for c in saver.alist({"configurable": {"thread_id": "thread-2"}}) ] - assert len(search_results_5) == 1 - assert search_results_5[0].config["configurable"]["checkpoint_ns"] == "" - - # search by config and checkpoint_ns - search_results_6 = [ - c - async for c in saver.alist( - { - "configurable": { - "thread_id": "thread-2", - "checkpoint_ns": "inner", - } - } - ) - ] - assert len(search_results_6) == 1 - assert ( - search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner" - ) + assert len(search_results_5) == 2 + assert { + search_results_5[0].config["configurable"]["checkpoint_ns"], + search_results_5[1].config["configurable"]["checkpoint_ns"], + } == {"", "inner"} # TODO: test before and limit params diff --git a/libs/checkpoint-sqlite/tests/test_sqlite.py b/libs/checkpoint-sqlite/tests/test_sqlite.py index 2147cca87..99b7a3728 100644 --- a/libs/checkpoint-sqlite/tests/test_sqlite.py +++ b/libs/checkpoint-sqlite/tests/test_sqlite.py @@ -87,28 +87,15 @@ class TestSqliteSaver: search_results_4 = list(saver.list(None, filter=query_4)) assert len(search_results_4) == 0 - # search by config (defaults to root graph checkpoints) + # search by config (defaults to checkpoints across all namespaces) search_results_5 = list( saver.list({"configurable": {"thread_id": "thread-2"}}) ) - assert len(search_results_5) == 1 - assert search_results_5[0].config["configurable"]["checkpoint_ns"] == "" - - # search by config and checkpoint_ns - search_results_6 = list( - saver.list( - { - "configurable": { - "thread_id": "thread-2", - "checkpoint_ns": "inner", - } - } - ) - ) - assert len(search_results_6) == 1 - assert ( - search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner" - ) + assert len(search_results_5) == 2 + assert { + search_results_5[0].config["configurable"]["checkpoint_ns"], + search_results_5[1].config["configurable"]["checkpoint_ns"], + } == {"", "inner"} # TODO: test before and limit params diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 52102ffef..6918de87a 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -177,62 +177,77 @@ class MemorySaver( Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. """ thread_ids = (config["configurable"]["thread_id"],) if config else self.storage - checkpoint_ns = ( - config["configurable"].get("checkpoint_ns", "") if config else "" + config_checkpoint_ns = ( + config["configurable"].get("checkpoint_ns") if config else None ) + config_checkpoint_id = get_checkpoint_id(config) if config else None for thread_id in thread_ids: - for checkpoint_id, (checkpoint, metadata_b, parent_checkpoint_id) in sorted( - self.storage[thread_id][checkpoint_ns].items(), - key=lambda x: x[0], - reverse=True, - ): - # filter by checkpoint ID - if ( - before - and (before_checkpoint_id := get_checkpoint_id(before)) - and checkpoint_id >= before_checkpoint_id - ): + for checkpoint_ns in self.storage[thread_id].keys(): + if config_checkpoint_ns and checkpoint_ns != config_checkpoint_ns: continue - # filter by metadata - metadata = self.serde.loads_typed(metadata_b) - if filter and not all( - query_value == metadata[query_key] - for query_key, query_value in filter.items() + for checkpoint_id, ( + checkpoint, + metadata_b, + parent_checkpoint_id, + ) in sorted( + self.storage[thread_id][checkpoint_ns].items(), + key=lambda x: x[0], + reverse=True, ): - continue + # filter by checkpoint ID from config + if config_checkpoint_id and checkpoint_id != config_checkpoint_id: + continue - # limit search results - if limit is not None and limit <= 0: - break - elif limit is not None: - limit -= 1 + # filter by checkpoint ID from `before` config + if ( + before + and (before_checkpoint_id := get_checkpoint_id(before)) + and checkpoint_id >= before_checkpoint_id + ): + continue - writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values() + # filter by metadata + metadata = self.serde.loads_typed(metadata_b) + if filter and not all( + query_value == metadata.get(query_key) + for query_key, query_value in filter.items() + ): + continue - yield CheckpointTuple( - config={ - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": checkpoint_id, + # limit search results + if limit is not None and limit <= 0: + break + elif limit is not None: + limit -= 1 + + writes = self.writes[ + (thread_id, checkpoint_ns, checkpoint_id) + ].values() + + yield CheckpointTuple( + config={ + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": checkpoint_id, + } + }, + checkpoint=self.serde.loads_typed(checkpoint), + metadata=metadata, + parent_config={ + "configurable": { + "thread_id": thread_id, + "checkpoint_ns": checkpoint_ns, + "checkpoint_id": parent_checkpoint_id, + } } - }, - checkpoint=self.serde.loads_typed(checkpoint), - metadata=metadata, - parent_config={ - "configurable": { - "thread_id": thread_id, - "checkpoint_ns": checkpoint_ns, - "checkpoint_id": parent_checkpoint_id, - } - } - if parent_checkpoint_id - else None, - pending_writes=[ - (id, c, self.serde.loads_typed(v)) for id, c, v in writes - ], - ) + if parent_checkpoint_id + else None, + pending_writes=[ + (id, c, self.serde.loads_typed(v)) for id, c, v in writes + ], + ) def put( self, @@ -338,7 +353,14 @@ class MemorySaver( """ loop = asyncio.get_running_loop() iter = await loop.run_in_executor( - None, partial(self.list, before=before, limit=limit, filter=filter), config + None, + partial( + self.list, + before=before, + limit=limit, + filter=filter, + ), + config, ) while True: # handling StopIteration exception inside coroutine won't work diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index e39a66e76..e80b9a678 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -109,7 +109,7 @@ class JsonPlusSerializer(SerializerProtocol): return self._encode_constructor_args(obj.__class__, args=[obj.value]) elif isinstance(obj, SendProtocol): return self._encode_constructor_args( - obj.__class__, kwargs={"node": obj.node, "arg": obj.arg} + obj.__class__, kwargs={"node": obj.node, "arg": obj.arg, "id": obj.id} ) elif isinstance(obj, (bytes, bytearray)): return self._encode_constructor_args( diff --git a/libs/checkpoint/langgraph/checkpoint/serde/types.py b/libs/checkpoint/langgraph/checkpoint/serde/types.py index d166a6a0d..71588cfa0 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -57,6 +57,7 @@ class SendProtocol(Protocol): # Mirrors langgraph.constants.Send node: str arg: Any + id: str def __hash__(self) -> int: ... diff --git a/libs/checkpoint/tests/test_memory.py b/libs/checkpoint/tests/test_memory.py index a0bc8d738..34c13b2d0 100644 --- a/libs/checkpoint/tests/test_memory.py +++ b/libs/checkpoint/tests/test_memory.py @@ -82,26 +82,20 @@ class TestMemorySaver: assert search_results_2[0].metadata == self.metadata_2 search_results_3 = list(self.memory_saver.list(None, filter=query_3)) - assert len(search_results_3) == 2 + assert len(search_results_3) == 3 search_results_4 = list(self.memory_saver.list(None, filter=query_4)) assert len(search_results_4) == 0 - # search by config (defaults to root graph checkpoints) + # search by config (defaults to checkpoints across all namespaces) search_results_5 = list( self.memory_saver.list({"configurable": {"thread_id": "thread-2"}}) ) - assert len(search_results_5) == 1 - assert search_results_5[0].config["configurable"]["checkpoint_ns"] == "" - - # search by config and checkpoint_ns - search_results_6 = list( - self.memory_saver.list( - {"configurable": {"thread_id": "thread-2", "checkpoint_ns": "inner"}} - ) - ) - assert len(search_results_6) == 1 - assert search_results_6[0].config["configurable"]["checkpoint_ns"] == "inner" + assert len(search_results_5) == 2 + assert { + search_results_5[0].config["configurable"]["checkpoint_ns"], + search_results_5[1].config["configurable"]["checkpoint_ns"], + } == {"", "inner"} # TODO: test before and limit params @@ -110,6 +104,7 @@ class TestMemorySaver: # save checkpoints self.memory_saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {}) self.memory_saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {}) + self.memory_saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {}) # call method / assertions query_1: CheckpointMetadata = {"source": "input"} # search by 1 key @@ -135,7 +130,7 @@ class TestMemorySaver: search_results_3 = [ c async for c in self.memory_saver.alist(None, filter=query_3) ] - assert len(search_results_3) == 2 + assert len(search_results_3) == 3 search_results_4 = [ c async for c in self.memory_saver.alist(None, filter=query_4) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index e2621f4f6..3b322d391 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -1,5 +1,6 @@ from dataclasses import dataclass -from typing import Any, Literal +from typing import Any, Literal, Optional +from uuid import uuid4 INPUT = "__input__" CONFIG_KEY_SEND = "__pregel_send" @@ -31,6 +32,7 @@ START = "__start__" END = "__end__" CHECKPOINT_NAMESPACE_SEPARATOR = "|" +SEND_CHECKPOINT_NAMESPACE_SEPARATOR = ":" class Send: @@ -49,6 +51,7 @@ class Send: Attributes: node (str): The name of the target node to send the message to. arg (Any): The state or message to send to the target node. + id (str): ID associated with the Send. Examples: >>> from typing import Annotated @@ -76,23 +79,26 @@ class Send: node: str arg: Any + id: Optional[str] - def __init__(self, /, node: str, arg: Any) -> None: + def __init__(self, /, node: str, arg: Any, id: Optional[str] = None) -> None: """ Initialize a new instance of the Send class. Args: node (str): The name of the target node to send the message to. arg (Any): The state or message to send to the target node. + id (str): ID associated with the Send. """ self.node = node self.arg = arg + self.id = id or str(uuid4()) def __hash__(self) -> int: - return hash((self.node, self.arg)) + return hash((self.node, self.arg, self.id)) def __repr__(self) -> str: - return f"Send(node={self.node!r}, arg={self.arg!r})" + return f"Send(node={self.node!r}, arg={self.arg!r}, id={self.id!r})" def __eq__(self, value: object) -> bool: return ( diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index d5eeb88e2..3ecbe726f 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -28,6 +28,7 @@ from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.constants import ( CHECKPOINT_NAMESPACE_SEPARATOR, END, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, START, TAG_HIDDEN, Send, @@ -160,10 +161,15 @@ class Graph: *, metadata: Optional[dict[str, Any]] = None, ) -> None: - if isinstance(node, str) and CHECKPOINT_NAMESPACE_SEPARATOR in node: - raise ValueError( - f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names." - ) + if isinstance(node, str): + for character in ( + CHECKPOINT_NAMESPACE_SEPARATOR, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, + ): + if character in node: + raise ValueError( + f"'{character}' is a reserved character and is not allowed in the node names." + ) if self.compiled: logger.warning( diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 73b7df783..a12796548 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -31,7 +31,11 @@ from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.named_barrier_value import NamedBarrierValue from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR, TAG_HIDDEN +from langgraph.constants import ( + CHECKPOINT_NAMESPACE_SEPARATOR, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, + TAG_HIDDEN, +) from langgraph.errors import InvalidUpdateError from langgraph.graph.graph import ( END, @@ -317,10 +321,14 @@ class StateGraph(Graph): if node == END or node == START: raise ValueError(f"Node `{node}` is reserved.") - if CHECKPOINT_NAMESPACE_SEPARATOR in node: - raise ValueError( - f"'{CHECKPOINT_NAMESPACE_SEPARATOR}' is a reserved character and is not allowed in the node names." - ) + for character in ( + CHECKPOINT_NAMESPACE_SEPARATOR, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, + ): + if character in node: + raise ValueError( + f"'{character}' is a reserved character and is not allowed in the node names." + ) try: if isfunction(action) and ( diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index e871815ea..e5648951a 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -54,17 +54,20 @@ from langgraph.channels.base import ( ) from langgraph.checkpoint.base import ( BaseCheckpointSaver, + CheckpointTuple, copy_checkpoint, create_checkpoint, empty_checkpoint, ) from langgraph.constants import ( + CHECKPOINT_NAMESPACE_SEPARATOR, CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_READ, CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, ERROR, INTERRUPT, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, ) from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ManagedValueSpec @@ -80,6 +83,7 @@ from langgraph.pregel.debug import ( print_step_writes, tasks_w_writes, ) +from langgraph.pregel.get_state import assemble_state_snapshot_hierarchy from langgraph.pregel.io import read_channels from langgraph.pregel.loop import AsyncPregelLoop, SyncPregelLoop from langgraph.pregel.manager import AsyncChannelsManager, ChannelsManager @@ -91,7 +95,9 @@ from langgraph.pregel.types import ( StateSnapshot, StreamMode, ) -from langgraph.pregel.utils import get_new_channel_versions +from langgraph.pregel.utils import ( + get_new_channel_versions, +) from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore @@ -178,6 +184,181 @@ class Channel: ) +def _get_checkpoint_ns_to_graph( + graph: Pregel, + checkpoint_ns_to_graph: Optional[dict[str, Pregel]] = None, + checkpoint_ns: str = "", + max_depth: int = 10, +) -> Pregel: + if checkpoint_ns_to_graph is None: + checkpoint_ns_to_graph = {} + + if max_depth <= 0: + raise RecursionError( + "Reached maximum recursion depth while building checkpoint NS -> graph mapping." + ) + + for node_name, node in graph.nodes.items(): + new_checkpoint_ns = ( + f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}" + if checkpoint_ns + else node_name + ) + if isinstance(node.bound, Pregel): + _get_checkpoint_ns_to_graph( + node.bound, checkpoint_ns_to_graph, new_checkpoint_ns, max_depth - 1 + ) + elif isinstance(node.bound, RunnableSequence): + for runnable in node.bound.steps: + if isinstance(runnable, Pregel): + _get_checkpoint_ns_to_graph( + runnable, + checkpoint_ns_to_graph, + new_checkpoint_ns, + max_depth - 1, + ) + + checkpoint_ns_to_graph[checkpoint_ns] = graph + return checkpoint_ns_to_graph + + +def _prepare_state_snapshot( + config: RunnableConfig, + checkpoint_ns_to_graph: dict[str, Pregel], + checkpoint_tuples: Iterator[CheckpointTuple], +) -> StateSnapshot: + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = config["configurable"].get("checkpoint_id") + checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} + for saved in checkpoint_tuples: + saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"] + saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"] + if checkpoint_id and saved_checkpoint_id != checkpoint_id: + continue + + graph_checkpoint_ns = saved_checkpoint_ns.split( + SEND_CHECKPOINT_NAMESPACE_SEPARATOR + )[0] + graph = checkpoint_ns_to_graph.get(graph_checkpoint_ns) + if graph is None: + continue + + with ChannelsManager( + graph.channels, saved.checkpoint, saved.config, skip_context=True + ) as ( + channels, + managed, + ): + next_tasks = prepare_next_tasks( + saved.checkpoint, + graph.nodes, + channels, + managed, + saved.config, + saved.metadata.get("step", -1) + 1, + for_execution=False, + ) + state_snapshot = StateSnapshot( + read_channels(channels, graph.stream_channels_asis), + tuple(t.name for t in next_tasks), + saved.config, + saved.metadata, + saved.checkpoint["ts"], + saved.parent_config, + tasks_w_writes(next_tasks, saved.pending_writes), + ) + + checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot + + if not checkpoint_ns_to_state_snapshots: + return StateSnapshot( + values={}, + next=(), + config=config, + metadata=None, + created_at=None, + parent_config=None, + tasks=(), + ) + + state_snapshot = assemble_state_snapshot_hierarchy( + checkpoint_ns, checkpoint_ns_to_state_snapshots + ) + return state_snapshot + + +async def _prepare_state_snapshot_async( + config: RunnableConfig, + checkpoint_ns_to_graph: dict[str, Pregel], + checkpoint_tuples: AsyncIterator[CheckpointTuple], +) -> StateSnapshot: + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = config["configurable"].get("checkpoint_id") + checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} + async for saved in checkpoint_tuples: + saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"] + saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"] + if checkpoint_id and saved_checkpoint_id != checkpoint_id: + continue + + graph_checkpoint_ns = saved_checkpoint_ns.split( + SEND_CHECKPOINT_NAMESPACE_SEPARATOR + )[0] + graph = checkpoint_ns_to_graph.get(graph_checkpoint_ns) + if graph is None: + continue + + async with AsyncChannelsManager( + graph.channels, saved.checkpoint, saved.config, skip_context=True + ) as (channels, managed): + next_tasks = prepare_next_tasks( + saved.checkpoint, + graph.nodes, + channels, + managed, + saved.config, + saved.metadata.get("step", -1) + 1, + for_execution=False, + ) + state_snapshot = StateSnapshot( + read_channels(channels, graph.stream_channels_asis), + tuple(t.name for t in next_tasks), + saved.config, + saved.metadata, + saved.checkpoint["ts"], + saved.parent_config, + tasks_w_writes(next_tasks, saved.pending_writes), + ) + + checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot + + if not checkpoint_ns_to_state_snapshots: + return StateSnapshot( + values={}, + next=(), + config=config, + metadata=None, + created_at=None, + parent_config=None, + tasks=(), + ) + + state_snapshot = assemble_state_snapshot_hierarchy( + checkpoint_ns, checkpoint_ns_to_state_snapshots + ) + return state_snapshot + + +def _has_nested_interrupts( + graph: Pregel, +) -> bool: + for child in graph.subgraphs: + if child.interrupt_after_nodes or child.interrupt_before_nodes: + return True + else: + return False + + class Pregel( RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]] ): @@ -347,6 +528,18 @@ class Pregel( k for k in self.channels if isinstance(self.channels[k], BaseChannel) ] + @property + def subgraphs(self) -> Iterator[Pregel]: + for node in self.nodes.values(): + if isinstance(node.bound, Pregel): + yield node.bound + yield from node.bound.subgraphs + elif isinstance(node.bound, RunnableSequence): + for runnable in node.bound.steps: + if isinstance(runnable, Pregel): + yield runnable + yield from runnable.subgraphs + def get_state(self, config: RunnableConfig) -> StateSnapshot: """Get the current state of the graph.""" if not self.checkpointer: @@ -354,31 +547,18 @@ class Pregel( config = merge_configs(self.config, config) if self.config else config saved = self.checkpointer.get_tuple(config) - checkpoint = saved.checkpoint if saved else empty_checkpoint() - config = saved.config if saved else config - with ChannelsManager(self.channels, checkpoint, config, skip_context=True) as ( - channels, - managed, - ): - next_tasks = prepare_next_tasks( - checkpoint, - self.nodes, - channels, - managed, - config, - saved.metadata.get("step", -1) + 1 if saved else -1, - for_execution=False, - ) + checkpoint_config = saved.config if saved else config + checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self) - return StateSnapshot( - read_channels(channels, self.stream_channels_asis), - tuple(t.name for t in next_tasks), - saved.config if saved else config, - saved.metadata if saved else None, - saved.checkpoint["ts"] if saved else None, - saved.parent_config if saved else None, - tasks_w_writes(next_tasks, saved.pending_writes if saved else None), - ) + # we only lookup subgraph checkpoints if we actually have subgraphs + if len(set(checkpoint_ns_to_graph)) == 1: + checkpoint_tuples = (saved,) + else: + checkpoint_tuples = self.checkpointer.list(saved.config) + + return _prepare_state_snapshot( + checkpoint_config, checkpoint_ns_to_graph, checkpoint_tuples + ) async def aget_state(self, config: RunnableConfig) -> StateSnapshot: """Get the current state of the graph.""" @@ -387,30 +567,22 @@ class Pregel( config = merge_configs(self.config, config) if self.config else config saved = await self.checkpointer.aget_tuple(config) - checkpoint = saved.checkpoint if saved else empty_checkpoint() + checkpoint_config = saved.config if saved else config + checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self) - config = saved.config if saved else config - async with AsyncChannelsManager( - self.channels, checkpoint, config, skip_context=True - ) as (channels, managed): - next_tasks = prepare_next_tasks( - checkpoint, - self.nodes, - channels, - managed, - config, - saved.metadata.get("step", -1) + 1 if saved else -1, - for_execution=False, - ) - return StateSnapshot( - read_channels(channels, self.stream_channels_asis), - tuple(t.name for t in next_tasks), - saved.config if saved else config, - saved.metadata if saved else None, - saved.checkpoint["ts"] if saved else None, - saved.parent_config if saved else None, - tasks_w_writes(next_tasks, saved.pending_writes if saved else None), - ) + # we only lookup subgraph checkpoints if we actually have subgraphs + if len(set(checkpoint_ns_to_graph)) == 1: + + async def alist_checkpoints(): + yield saved + + checkpoint_tuples = alist_checkpoints() + else: + checkpoint_tuples = self.checkpointer.alist(saved.config) + + return await _prepare_state_snapshot_async( + checkpoint_config, checkpoint_ns_to_graph, checkpoint_tuples + ) def get_state_history( self, @@ -428,39 +600,30 @@ class Pregel( and signature(self.checkpointer.list).parameters.get("filter") is None ): raise ValueError("Checkpointer does not support filtering") - for ( - config, - checkpoint, - metadata, - parent_config, - pending_writes, - ) in self.checkpointer.list( - merge_configs(self.config, config) if self.config else config, - before=before, - limit=limit, - filter=filter, - ): - with ChannelsManager( - self.channels, checkpoint, config, skip_context=True - ) as (channels, managed): - next_tasks = prepare_next_tasks( - checkpoint, - self.nodes, - channels, - managed, - config, - metadata.get("step", -1) + 1, - for_execution=False, - ) - yield StateSnapshot( - read_channels(channels, self.stream_channels_asis), - tuple(t.name for t in next_tasks), - config, - metadata, - checkpoint["ts"], - parent_config, - tasks_w_writes(next_tasks, pending_writes), - ) + + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_ns_to_graph = _get_checkpoint_ns_to_graph(self) + # find all matching checkpoint tuples for parent and subgraphs + checkpoint_tuples = [ + checkpoint_tuple + for checkpoint_tuple in self.checkpointer.list( + merge_configs(self.config, config) if self.config else config, + before=before, + limit=limit, + filter=filter, + ) + ] + for checkpoint_tuple in checkpoint_tuples: + if ( + checkpoint_tuple.config["configurable"]["checkpoint_ns"] + != checkpoint_ns + ): + continue + + state_snapshot = _prepare_state_snapshot( + checkpoint_tuple.config, checkpoint_ns_to_graph, iter(checkpoint_tuples) + ) + yield state_snapshot async def aget_state_history( self, @@ -478,39 +641,36 @@ class Pregel( and signature(self.checkpointer.list).parameters.get("filter") is None ): raise ValueError("Checkpointer does not support filtering") - async for ( - config, - checkpoint, - metadata, - parent_config, - pending_writes, - ) in self.checkpointer.alist( - merge_configs(self.config, config) if self.config else config, - before=before, - limit=limit, - filter=filter, - ): - async with AsyncChannelsManager( - self.channels, checkpoint, config, skip_context=True - ) as (channels, managed): - next_tasks = prepare_next_tasks( - checkpoint, - self.nodes, - channels, - managed, - config, - metadata.get("step", -1) + 1, - for_execution=False, - ) - yield StateSnapshot( - read_channels(channels, self.stream_channels_asis), - tuple(t.name for t in next_tasks), - config, - metadata, - checkpoint["ts"], - parent_config, - tasks_w_writes(next_tasks, pending_writes), - ) + + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_ns_to_graph = _get_checkpoint_ns_to_graph(self) + # find all matching checkpoint tuples for parent and subgraphs + checkpoint_tuples = [ + checkpoint_tuple + async for checkpoint_tuple in self.checkpointer.alist( + merge_configs(self.config, config) if self.config else config, + before=before, + limit=limit, + filter=filter, + ) + ] + + # turn matching checkpoint tuples into an async iterator + async def alist_checkpoints() -> AsyncIterator[CheckpointTuple]: + for checkpoint_tuple in checkpoint_tuples: + yield checkpoint_tuple + + for checkpoint_tuple in checkpoint_tuples: + if ( + checkpoint_tuple.config["configurable"]["checkpoint_ns"] + != checkpoint_ns + ): + continue + + state_snapshot = await _prepare_state_snapshot_async( + checkpoint_tuple.config, checkpoint_ns_to_graph, alist_checkpoints() + ) + yield state_snapshot def update_state( self, @@ -816,7 +976,7 @@ class Pregel( if ( config is not None and config.get("configurable", {}).get(CONFIG_KEY_CHECKPOINTER) - and (interrupt_after or interrupt_before) + and (interrupt_after or interrupt_before or _has_nested_interrupts(self)) ): checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][ CONFIG_KEY_CHECKPOINTER @@ -1023,6 +1183,7 @@ class Pregel( ) else: loop.put_writes(task.id, [(ERROR, exc)]) + else: # save task writes to checkpointer loop.put_writes(task.id, task.writes) @@ -1263,6 +1424,7 @@ class Pregel( ) if not done: break # timed out + for fut in done: task = futures.pop(fut) if exc := _exception(fut): @@ -1273,6 +1435,7 @@ class Pregel( ) else: loop.put_writes(task.id, [(ERROR, exc)]) + else: # save task writes to checkpointer loop.put_writes(task.id, task.writes) @@ -1480,11 +1643,3 @@ def _panic_or_proceed( inflight.pop().cancel() # raise timeout error raise timeout_exc_cls(f"Timed out at step {step}") - - -def _with_mode(mode: StreamMode, on: bool, iter: Iterator[Any]) -> Iterator[Any]: - if on: - for chunk in iter: - yield (mode, chunk) - else: - yield from iter diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 29adcef5e..5b6064306 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -293,9 +293,9 @@ def prepare_next_tasks( "langgraph_task_idx": len(tasks), } checkpoint_ns = ( - f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}" + f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}:{packet.id}" if parent_ns - else packet.node + else f"{packet.node}:{packet.id}" ) task_id = str( uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata))) @@ -343,8 +343,10 @@ def prepare_next_tasks( PregelTaskWrites(packet.node, writes, triggers), config, ), - # in Send we can't checkpoint nested graphs - # as they could be running in parallel + CONFIG_KEY_CHECKPOINTER: checkpointer, + CONFIG_KEY_RESUMING: is_resuming, + "checkpoint_id": checkpoint["id"], + "checkpoint_ns": checkpoint_ns, }, ), triggers, diff --git a/libs/langgraph/langgraph/pregel/get_state.py b/libs/langgraph/langgraph/pregel/get_state.py new file mode 100644 index 000000000..b8efbe6a2 --- /dev/null +++ b/libs/langgraph/langgraph/pregel/get_state.py @@ -0,0 +1,36 @@ +from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR +from langgraph.pregel.types import StateSnapshot + + +def assemble_state_snapshot_hierarchy( + root_checkpoint_ns: str, + checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot], +) -> StateSnapshot: + checkpoint_ns_list_to_visit = sorted( + checkpoint_ns_to_state_snapshots.keys(), + key=lambda x: len(x.split(CHECKPOINT_NAMESPACE_SEPARATOR)), + ) + while checkpoint_ns_list_to_visit: + checkpoint_ns = checkpoint_ns_list_to_visit.pop() + state_snapshot = checkpoint_ns_to_state_snapshots[checkpoint_ns] + *path, subgraph_node = checkpoint_ns.split(CHECKPOINT_NAMESPACE_SEPARATOR) + parent_checkpoint_ns = CHECKPOINT_NAMESPACE_SEPARATOR.join(path) + if subgraph_node and ( + parent_state_snapshot := checkpoint_ns_to_state_snapshots.get( + parent_checkpoint_ns + ) + ): + parent_subgraph_snapshots = { + **(parent_state_snapshot.subgraph_state_snapshots or {}), + subgraph_node: state_snapshot, + } + checkpoint_ns_to_state_snapshots[ + parent_checkpoint_ns + ] = checkpoint_ns_to_state_snapshots[parent_checkpoint_ns]._replace( + subgraph_state_snapshots=parent_subgraph_snapshots + ) + + state_snapshot = checkpoint_ns_to_state_snapshots.pop(root_checkpoint_ns, None) + if state_snapshot is None: + raise ValueError(f"Missing checkpoint for checkpoint NS '{root_checkpoint_ns}'") + return state_snapshot diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index a8c5a7beb..d4881452b 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -92,6 +92,8 @@ class StateSnapshot(NamedTuple): """Config used to fetch the parent snapshot, if any""" tasks: tuple[PregelTask, ...] """Tasks to execute in this step. If already attempted, may contain an error.""" + subgraph_state_snapshots: Optional[dict[str, "StateSnapshot"]] = None + """State snapshots of subgraphs represented as a mapping from checkpoint namespace (`checkpoint_ns`) to snapshot.""" All = Literal["*"] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index ba91c4eb0..685eb7c87 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -60,7 +60,12 @@ from langgraph.prebuilt.chat_agent_executor import ( create_tool_calling_executor, ) from langgraph.prebuilt.tool_node import ToolNode -from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot +from langgraph.pregel import ( + Channel, + GraphRecursionError, + Pregel, + StateSnapshot, +) from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore @@ -8292,7 +8297,10 @@ def test_nested_graph_interrupts( graph = StateGraph(State) graph.add_node("outer_1", outer_1) - graph.add_node("inner", inner.compile(interrupt_before=["inner_2"])) + 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") @@ -8309,7 +8317,12 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask( + AnyStr(), + "inner", + ), + ), next=("inner",), config={ "configurable": { @@ -8331,6 +8344,47 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -8453,6 +8507,42 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -8539,7 +8629,12 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask( + AnyStr(), + "inner", + ), + ), next=("inner",), config={ "configurable": { @@ -8608,7 +8703,12 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask( + AnyStr(), + "inner", + ), + ), next=("inner",), config={ "configurable": { @@ -8630,6 +8730,47 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -8757,6 +8898,42 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -8811,11 +8988,15 @@ def test_nested_graph_interrupts( "my_key": "hi my value", }, ] - # interrupted after "inner" assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask( + AnyStr(), + "inner", + ), + ), next=("inner",), config={ "configurable": { @@ -8837,6 +9018,47 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -8884,6 +9106,7 @@ def test_nested_graph_interrupts( "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"}, @@ -8934,6 +9157,42 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -9058,6 +9317,42 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -9133,6 +9428,42 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -9183,9 +9514,12 @@ def test_nested_graph_interrupts( ] assert child_state_history == [ StateSnapshot( - values={"my_key": "hi my value here"}, - tasks=(), - next=(), + 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", @@ -9255,6 +9589,41 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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"}, @@ -9280,6 +9649,41 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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"}, @@ -9404,6 +9808,42 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "hi my value"}, @@ -9429,6 +9869,42 @@ def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -9513,7 +9989,10 @@ 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) @@ -9597,7 +10076,6 @@ def test_nested_graph_interrupts_parallel( ] -@pytest.mark.skip @pytest.mark.parametrize( "checkpointer_name", ["memory", "sqlite", "postgres", "postgres_pipe"], @@ -9605,7 +10083,7 @@ def test_nested_graph_interrupts_parallel( def test_doubly_nested_graph_interrupts( request: pytest.FixtureRequest, checkpointer_name: str ) -> None: - checkpointer = request.getfixturevalue(checkpointer_name) + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) class State(TypedDict): my_key: str @@ -9632,7 +10110,10 @@ def test_doubly_nested_graph_interrupts( grandchild.set_finish_point("grandchild_2") child = StateGraph(ChildState) - child.add_node("child_1", grandchild.compile(interrupt_before=["grandchild_2"])) + child.add_node( + "child_1", + grandchild.compile(interrupt_before=["grandchild_2"]), + ) child.set_entry_point("child_1") child.set_finish_point("child_1") @@ -9693,6 +10174,926 @@ def test_doubly_nested_graph_interrupts( ] +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_nested_graph_state( + 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 + other_parent_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) + + config = {"configurable": {"thread_id": "1"}} + app.invoke({"my_key": "my value"}, config, debug=True) + # test state w/ nested subgraph state (right after interrupt) + assert app.get_state(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={ + "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(), + } + }, + subgraph_state_snapshots={ + "inner": StateSnapshot( + values={"my_key": "hi my value here", "my_other_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + error=None, + ), + ), + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ) + assert list(app.get_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={ + "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(), + } + }, + subgraph_state_snapshots={ + "inner": StateSnapshot( + values={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + error=None, + ), + ), + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ), + 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={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + subgraph_state_snapshots=None, + ), + ] + app.invoke(None, config, debug=True) + # test state w/ nested subgraph state (after resuming from interrupt) + assert app.get_state(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={ + "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(), + } + }, + ) + # test full history at the end + actual_history = list(app.get_state_history(config)) + expected_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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ), + 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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots={ + "inner": 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": "inner", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ), + 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={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + subgraph_state_snapshots=None, + ), + ] + assert actual_history == expected_history + # test looking up parent state by checkpoint ID + for actual_snapshot, expected_snapshot in zip(actual_history, expected_history): + assert app.get_state(actual_snapshot.config) == expected_snapshot + + +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_doubly_nested_graph_state( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + class State(TypedDict): + my_key: str + + class ChildState(TypedDict): + my_key: str + + class GrandChildState(TypedDict): + my_key: str + + def grandchild_1(state: ChildState): + return {"my_key": state["my_key"] + " here"} + + def grandchild_2(state: ChildState): + return { + "my_key": state["my_key"] + " and there", + } + + grandchild = StateGraph(GrandChildState) + grandchild.add_node("grandchild_1", grandchild_1) + grandchild.add_node("grandchild_2", grandchild_2) + grandchild.add_edge("grandchild_1", "grandchild_2") + grandchild.set_entry_point("grandchild_1") + grandchild.set_finish_point("grandchild_2") + + child = StateGraph(ChildState) + child.add_node( + "child_1", + grandchild.compile(interrupt_before=["grandchild_2"]), + ) + child.set_entry_point("child_1") + child.set_finish_point("child_1") + + def parent_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + def parent_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("parent_1", parent_1) + graph.add_node("child", child.compile()) + graph.add_node("parent_2", parent_2) + graph.set_entry_point("parent_1") + graph.add_edge("parent_1", "child") + graph.add_edge("child", "parent_2") + graph.set_finish_point("parent_2") + + app = graph.compile(checkpointer=checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + app.invoke({"my_key": "my value"}, config, debug=True) + assert app.get_state(config) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "child"),), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots={ + "child": StateSnapshot( + values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "child_1"),), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots={ + "child_1": StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + ), + ), + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child|child_1", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child|child_1", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ) + }, + ) + app.invoke(None, config, debug=True) + assert app.get_state(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={ + "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(), + } + }, + ) + + # test getting snapshot by ID + config = list(app.get_state_history(config))[2].config + # test getting grandchild snapshot + assert app.get_state(config) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "child"),), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots={ + "child": StateSnapshot( + values={"my_key": "hi my value here and there"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"child_1": {"my_key": "hi my value here and there"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots={ + "child_1": StateSnapshot( + values={"my_key": "hi my value here and there"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child|child_1", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "grandchild_2": {"my_key": "hi my value here and there"} + }, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child|child_1", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ) + }, + ) + + +@pytest.mark.repeat(10) +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite", "postgres", "postgres_pipe"], +) +def test_send_to_nested_graphs( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + class OverallState(TypedDict): + subjects: list[str] + jokes: Annotated[list[str], operator.add] + + def continue_to_jokes(state: OverallState): + return [Send("generate_joke", {"subject": s}) for s in state["subjects"]] + + class JokeState(TypedDict): + subject: str + + def edit(state: JokeState): + subject = state["subject"] + return {"subject": f"{subject} - hohoho"} + + # subgraph + subgraph = StateGraph(input=JokeState, output=OverallState) + subgraph.add_node("edit", edit) + subgraph.add_node( + "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]} + ) + subgraph.set_entry_point("edit") + subgraph.add_edge("edit", "generate") + subgraph.set_finish_point("generate") + + # parent graph + builder = StateGraph(OverallState) + builder.add_node( + "generate_joke", + subgraph.compile(interrupt_before=["generate"]), + ) + builder.add_conditional_edges(START, continue_to_jokes) + builder.add_edge("generate_joke", END) + + graph = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} + + # invoke and pause at nested interrupt + assert graph.invoke({"subjects": ["cats", "dogs"]}, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": [], + } + actual_snapshot = graph.get_state(config) + subgraph_nodes = list(actual_snapshot.subgraph_state_snapshots.keys()) + assert len(subgraph_nodes) == 2 + for subgraph_node in subgraph_nodes: + assert subgraph_node.split(":")[0] == "generate_joke" + + subgraph_state_snapshots = { + subgraph_node: graph.get_state( + {"configurable": {"thread_id": "1", "checkpoint_ns": subgraph_node}} + ) + for subgraph_node in subgraph_nodes + } + + expected_snapshot = StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + tasks=( + PregelTask(AnyStr(), "generate_joke"), + PregelTask(AnyStr(), "generate_joke"), + ), + next=("generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=subgraph_state_snapshots, + ) + assert actual_snapshot == expected_snapshot + + # continue past interrupt + assert graph.invoke(None, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + } + + actual_snapshot = graph.get_state(config) + expected_snapshot = StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about dogs - hohoho"]}, + ] + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + assert actual_snapshot == expected_snapshot + + # test full history + actual_history = list(graph.get_state_history(config)) + + # get subgraph node state for expected history + subgraph_state_snapshots = { + subgraph_node: graph.get_state( + {"configurable": {"thread_id": "1", "checkpoint_ns": subgraph_node}} + ) + for subgraph_node in subgraph_nodes + } + expected_history = [ + StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about dogs - hohoho"]}, + ] + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + tasks=( + PregelTask(AnyStr(), "generate_joke"), + PregelTask(AnyStr(), "generate_joke"), + ), + next=("generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=subgraph_state_snapshots, + ), + StateSnapshot( + values={"jokes": []}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"subjects": ["cats", "dogs"]}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + subgraph_state_snapshots=None, + ), + ] + assert actual_history == expected_history + + def test_repeat_condition(snapshot: SnapshotAssertion) -> None: class AgentState(TypedDict): hello: str diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index f59433a07..50aafae86 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -60,7 +60,12 @@ from langgraph.prebuilt.chat_agent_executor import ( ) from langgraph.prebuilt.tool_executor import ToolExecutor from langgraph.prebuilt.tool_node import ToolNode -from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot +from langgraph.pregel import ( + Channel, + GraphRecursionError, + Pregel, + StateSnapshot, +) from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.types import PregelTask from langgraph.store.memory import MemoryStore @@ -6768,7 +6773,10 @@ async def test_nested_graph_interrupts( graph = StateGraph(State) graph.add_node("outer_1", outer_1) - graph.add_node("inner", inner.compile(interrupt_before=["inner_2"])) + 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") @@ -6785,7 +6793,12 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask( + AnyStr(), + "inner", + ), + ), next=("inner",), config={ "configurable": { @@ -6807,6 +6820,47 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -6929,6 +6983,42 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -7020,7 +7110,12 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask( + AnyStr(), + "inner", + ), + ), next=("inner",), config={ "configurable": { @@ -7089,7 +7184,12 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask( + AnyStr(), + "inner", + ), + ), next=("inner",), config={ "configurable": { @@ -7111,6 +7211,47 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -7238,6 +7379,42 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -7298,7 +7475,12 @@ async def test_nested_graph_interrupts( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask( + AnyStr(), + "inner", + ), + ), next=("inner",), config={ "configurable": { @@ -7320,6 +7502,47 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -7367,6 +7590,7 @@ async def test_nested_graph_interrupts( "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"}, @@ -7417,6 +7641,42 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -7541,6 +7801,42 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -7593,7 +7889,12 @@ async def test_nested_graph_interrupts( assert state_history == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask( + AnyStr(), + "inner", + ), + ), next=("inner",), config={ "configurable": { @@ -7615,6 +7916,47 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -7666,9 +8008,12 @@ async def test_nested_graph_interrupts( ] assert child_state_history == [ StateSnapshot( - values={"my_key": "hi my value here"}, - tasks=(), - next=(), + 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", @@ -7708,7 +8053,7 @@ async def test_nested_graph_interrupts( # 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 + # 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" } @@ -7738,6 +8083,41 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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"}, @@ -7763,6 +8143,41 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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"}, @@ -7805,9 +8220,9 @@ async def test_nested_graph_interrupts( parent_config=None, ), ] - # going to resume from interrupt + # going to restart from interrupt interrupt_config = interrupt_state_snapshot.config - assert (await app.ainvoke(None, interrupt_config, debug=True)) == { + 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)] == [ @@ -7887,6 +8302,42 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "hi my value"}, @@ -7912,6 +8363,42 @@ async def test_nested_graph_interrupts( "checkpoint_id": AnyStr(), } }, + subgraph_state_snapshots={ + "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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, ), StateSnapshot( values={"my_key": "my value"}, @@ -7996,7 +8483,10 @@ async 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) @@ -8082,7 +8572,6 @@ async def test_nested_graph_interrupts_parallel( ] -@pytest.mark.skip @pytest.mark.parametrize( "checkpointer_name", ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], @@ -8117,7 +8606,10 @@ async def test_doubly_nested_graph_interrupts( grandchild.set_finish_point("grandchild_2") child = StateGraph(ChildState) - child.add_node("child_1", grandchild.compile(interrupt_before=["grandchild_2"])) + child.add_node( + "child_1", + grandchild.compile(interrupt_before=["grandchild_2"]), + ) child.set_entry_point("child_1") child.set_finish_point("child_1") @@ -8181,6 +8673,933 @@ async def test_doubly_nested_graph_interrupts( ] +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_nested_graph_state( + 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 + other_parent_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) + + config = {"configurable": {"thread_id": "1"}} + await app.ainvoke({"my_key": "my value"}, config, debug=True) + # test state w/ nested subgraph state (right after interrupt) + assert await app.aget_state(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={ + "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(), + } + }, + subgraph_state_snapshots={ + "inner": StateSnapshot( + values={"my_key": "hi my value here", "my_other_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + error=None, + ), + ), + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ) + 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={ + "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(), + } + }, + subgraph_state_snapshots={ + "inner": StateSnapshot( + values={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + error=None, + ), + ), + next=("inner_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "inner", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ), + 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={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + subgraph_state_snapshots=None, + ), + ] + await app.ainvoke(None, config, debug=True) + # test state w/ nested subgraph state (after resuming from interrupt) + assert await app.aget_state(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={ + "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(), + } + }, + ) + # test full history at the end + actual_history = [s async for s in app.aget_state_history(config)] + expected_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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ), + 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={ + "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(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots={ + "inner": 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": "inner", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ), + 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={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + subgraph_state_snapshots=None, + ), + ] + assert actual_history == expected_history + # test looking up parent state by checkpoint ID + for actual_snapshot, expected_snapshot in zip(actual_history, expected_history): + assert await app.aget_state(actual_snapshot.config) == expected_snapshot + + +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_doubly_nested_graph_state( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + class State(TypedDict): + my_key: str + + class ChildState(TypedDict): + my_key: str + + class GrandChildState(TypedDict): + my_key: str + + async def grandchild_1(state: ChildState): + return {"my_key": state["my_key"] + " here"} + + async def grandchild_2(state: ChildState): + return { + "my_key": state["my_key"] + " and there", + } + + grandchild = StateGraph(GrandChildState) + grandchild.add_node("grandchild_1", grandchild_1) + grandchild.add_node("grandchild_2", grandchild_2) + grandchild.add_edge("grandchild_1", "grandchild_2") + grandchild.set_entry_point("grandchild_1") + grandchild.set_finish_point("grandchild_2") + + child = StateGraph(ChildState) + child.add_node( + "child_1", + grandchild.compile(interrupt_before=["grandchild_2"]), + ) + child.set_entry_point("child_1") + child.set_finish_point("child_1") + + async def parent_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + async def parent_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("parent_1", parent_1) + graph.add_node("child", child.compile()) + graph.add_node("parent_2", parent_2) + graph.set_entry_point("parent_1") + graph.add_edge("parent_1", "child") + graph.add_edge("child", "parent_2") + graph.set_finish_point("parent_2") + + app = graph.compile(checkpointer=checkpointer) + + # test invoke w/ nested interrupt + config = {"configurable": {"thread_id": "1"}} + await app.ainvoke({"my_key": "my value"}, config, debug=True) + assert await app.aget_state(config) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "child"),), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots={ + "child": StateSnapshot( + values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "child_1", + ), + ), + next=("child_1",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots={ + "child_1": StateSnapshot( + values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + ), + ), + next=("grandchild_2",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child|child_1", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"grandchild_1": {"my_key": "hi my value here"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child|child_1", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ) + }, + ) + await app.ainvoke(None, config, debug=True) + assert await app.aget_state(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={ + "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(), + } + }, + ) + # test getting snapshot by ID + config = [s async for s in app.aget_state_history(config)][2].config + # test getting grandchild snapshot + assert await app.aget_state(config) == StateSnapshot( + values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "child"),), + next=("child",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "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(), + } + }, + subgraph_state_snapshots={ + "child": StateSnapshot( + values={"my_key": "hi my value here and there"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"child_1": {"my_key": "hi my value here and there"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots={ + "child_1": StateSnapshot( + values={"my_key": "hi my value here and there"}, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child|child_1", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "grandchild_2": {"my_key": "hi my value here and there"} + }, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "child|child_1", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ) + }, + ) + + +@pytest.mark.repeat(10) +@pytest.mark.parametrize( + "checkpointer_name", + ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], +) +async def test_send_to_nested_graphs( + request: pytest.FixtureRequest, checkpointer_name: str +) -> None: + checkpointer = request.getfixturevalue("checkpointer_" + checkpointer_name) + + class OverallState(TypedDict): + subjects: list[str] + jokes: Annotated[list[str], operator.add] + + async def continue_to_jokes(state: OverallState): + return [Send("generate_joke", {"subject": s}) for s in state["subjects"]] + + class JokeState(TypedDict): + subject: str + + async def edit(state: JokeState): + subject = state["subject"] + return {"subject": f"{subject} - hohoho"} + + # subgraph + subgraph = StateGraph(input=JokeState, output=OverallState) + subgraph.add_node("edit", edit) + subgraph.add_node( + "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]} + ) + subgraph.set_entry_point("edit") + subgraph.add_edge("edit", "generate") + subgraph.set_finish_point("generate") + + # parent graph + builder = StateGraph(OverallState) + builder.add_node( + "generate_joke", + subgraph.compile(interrupt_before=["generate"]), + ) + builder.add_conditional_edges(START, continue_to_jokes) + builder.add_edge("generate_joke", END) + + graph = builder.compile(checkpointer=checkpointer) + config = {"configurable": {"thread_id": "1"}} + + # invoke and pause at nested interrupt + assert await graph.ainvoke({"subjects": ["cats", "dogs"]}, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": [], + } + actual_snapshot = await graph.aget_state(config) + subgraph_nodes = list(actual_snapshot.subgraph_state_snapshots.keys()) + assert len(subgraph_nodes) == 2 + for subgraph_node in subgraph_nodes: + assert subgraph_node.split(":")[0] == "generate_joke" + + subgraph_state_snapshots = { + subgraph_node: await graph.aget_state( + {"configurable": {"thread_id": "1", "checkpoint_ns": subgraph_node}} + ) + for subgraph_node in subgraph_nodes + } + expected_snapshot = StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + tasks=( + PregelTask(AnyStr(), "generate_joke"), + PregelTask(AnyStr(), "generate_joke"), + ), + next=("generate_joke", "generate_joke"), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=subgraph_state_snapshots, + ) + assert actual_snapshot == expected_snapshot + + # continue past interrupt + assert await graph.ainvoke(None, config=config) == { + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + } + + actual_snapshot = await graph.aget_state(config) + expected_snapshot = StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about dogs - hohoho"]}, + ] + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + ) + assert actual_snapshot == expected_snapshot + + # test full history + actual_history = [c async for c in graph.aget_state_history(config)] + # get subgraph node state for expected history + subgraph_state_snapshots = { + subgraph_node: await graph.aget_state( + {"configurable": {"thread_id": "1", "checkpoint_ns": subgraph_node}} + ) + for subgraph_node in subgraph_nodes + } + expected_history = [ + StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + tasks=(), + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "generate_joke": [ + {"jokes": ["Joke about cats - hohoho"]}, + {"jokes": ["Joke about dogs - hohoho"]}, + ] + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + next=("generate_joke", "generate_joke"), + tasks=( + PregelTask(AnyStr(), "generate_joke"), + PregelTask(AnyStr(), "generate_joke"), + ), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=subgraph_state_snapshots, + ), + StateSnapshot( + values={"jokes": []}, + tasks=(PregelTask(AnyStr(), "__start__"),), + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"subjects": ["cats", "dogs"]}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + subgraph_state_snapshots=None, + ), + ] + assert actual_history == expected_history + + async def test_checkpoint_metadata() -> None: """This test verifies that a run's configurable fields are merged with the previous checkpoint config for each step in the run.