diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 86c8b0eec..4f0f08a84 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -259,6 +259,7 @@ class BaseCheckpointSaver(ABC): filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, + include_nested_checkpoints: bool = False, ) -> Iterator[CheckpointTuple]: """List checkpoints that match the given criteria. @@ -350,6 +351,7 @@ class BaseCheckpointSaver(ABC): filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, + include_nested_checkpoints: bool = False, ) -> AsyncIterator[CheckpointTuple]: """Asynchronously list checkpoints that match the given criteria. diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index d989b1fbc..5e60ebaf0 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -157,6 +157,7 @@ class MemorySaver( filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, + include_nested_checkpoints: bool = False, ) -> Iterator[CheckpointTuple]: """List checkpoints from the in-memory storage. @@ -177,53 +178,67 @@ class MemorySaver( config["configurable"].get("checkpoint_ns", "") if config else "" ) 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 - ): - 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() - ): - continue - - # limit search results - if limit is not None and limit <= 0: - break - elif limit is not None: - limit -= 1 - - 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, - } - } - if parent_checkpoint_id - else None, + checkpoint_ns_iter = ( + ( + key + for key in self.storage[thread_id].keys() + if key.startswith(checkpoint_ns) ) + if include_nested_checkpoints + else [checkpoint_ns] + ) + for checkpoint_ns in checkpoint_ns_iter: + 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 + ): + 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() + ): + continue + + # limit search results + if limit is not None and limit <= 0: + break + elif limit is not None: + limit -= 1 + + 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, + } + } + if parent_checkpoint_id + else None, + ) def put( self, @@ -315,6 +330,7 @@ class MemorySaver( filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, + include_nested_checkpoints: bool = False, ) -> AsyncIterator[CheckpointTuple]: """Asynchronous version of list. @@ -335,6 +351,7 @@ class MemorySaver( before=before, limit=limit, filter=filter, + include_nested_checkpoints=include_nested_checkpoints, ), config, ) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 2a31ec281..34c12604c 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -65,12 +65,12 @@ from langgraph.checkpoint.base import ( empty_checkpoint, ) from langgraph.constants import ( + CHECKPOINT_NAMESPACE_SEPARATOR, CONFIG_KEY_CHECKPOINTER, CONFIG_KEY_READ, CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, INTERRUPT, - THREAD_ID_SEPARATOR, ) from langgraph.errors import GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ( @@ -422,35 +422,36 @@ class Pregel( @staticmethod def _assemble_state_snapshot_hierarchy( - root_thread_id: str, thread_id_to_state_snapshots: dict[str, StateSnapshot] + root_checkpoint_ns: str, + checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot], ) -> StateSnapshot: - thread_ids_to_visit = sorted( - thread_id_to_state_snapshots.keys(), - key=lambda x: len(x.split(THREAD_ID_SEPARATOR)), + checkpoint_ns_list_to_visit = sorted( + checkpoint_ns_to_state_snapshots.keys(), + key=lambda x: len(x.split(CHECKPOINT_NAMESPACE_SEPARATOR)), ) - while thread_ids_to_visit: - thread_id = thread_ids_to_visit.pop() - state_snapshot = thread_id_to_state_snapshots[thread_id] - *path, subgraph_node = thread_id.split(THREAD_ID_SEPARATOR) - parent_thread_id = THREAD_ID_SEPARATOR.join(path) - if parent_thread_id and ( - parent_state_snapshot := thread_id_to_state_snapshots.get( - parent_thread_id + 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, } - thread_id_to_state_snapshots[ - parent_thread_id - ] = thread_id_to_state_snapshots[parent_thread_id]._replace( + 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 = thread_id_to_state_snapshots.pop(root_thread_id, None) + state_snapshot = checkpoint_ns_to_state_snapshots.pop(root_checkpoint_ns, None) if state_snapshot is None: - raise ValueError(f"Missing checkpoint for thread ID '{root_thread_id}'") + raise ValueError(f"Missing checkpoint for thread ID '{root_checkpoint_ns}'") return state_snapshot def get_state( @@ -461,36 +462,51 @@ class Pregel( raise ValueError("No checkpointer set") if include_subgraph_state: - checkpoint_tuples = self.checkpointer.list(config, as_prefix=True) + checkpoint_tuples = self.checkpointer.list( + config, include_nested_checkpoints=True + ) else: checkpoint_tuples = iter([self.checkpointer.get_tuple(config)]) - thread_id = config["configurable"]["thread_id"] - thread_ts = config["configurable"].get("thread_ts") - thread_id_to_thread_ts: dict[str, str] = {} - thread_id_to_state_snapshots: dict[str, StateSnapshot] = {} + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = config["configurable"].get("checkpoint_id") + checkpoint_ns_to_checkpoint_id: dict[str, str] = {} + checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} for checkpoint_tuple in checkpoint_tuples: - checkpoint_thread_id = checkpoint_tuple.config["configurable"]["thread_id"] - checkpoint_thread_ts = checkpoint_tuple.config["configurable"]["thread_ts"] - if thread_ts and thread_ts != checkpoint_thread_ts: + saved_checkpoint_ns = checkpoint_tuple.config["configurable"][ + "checkpoint_ns" + ] + saved_checkpoint_id = checkpoint_tuple.config["configurable"][ + "checkpoint_id" + ] + if checkpoint_id and checkpoint_id != saved_checkpoint_id: continue - existing_thread_ts = thread_id_to_thread_ts.get(checkpoint_thread_id) - # keep only most recent thread_ts - if existing_thread_ts is None or checkpoint_thread_ts > existing_thread_ts: + existing_checkpoint_id = checkpoint_ns_to_checkpoint_id.get( + saved_checkpoint_ns + ) + # keep only most recent checkpoint_id + if ( + existing_checkpoint_id is None + or saved_checkpoint_id > existing_checkpoint_id + ): state_snapshot = self._prepare_state_snapshot(checkpoint_tuple, config) - thread_id_to_state_snapshots[checkpoint_thread_id] = state_snapshot - thread_id_to_thread_ts[checkpoint_thread_id] = checkpoint_thread_ts + checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot + checkpoint_ns_to_checkpoint_id[ + saved_checkpoint_ns + ] = saved_checkpoint_id - if not thread_id_to_state_snapshots: - error_msg = f"Could not find checkpoints for thread ID '{thread_id}'" - if thread_ts: - error_msg += f" and thread TS '{thread_ts}'" + if not checkpoint_ns_to_state_snapshots: + error_msg = ( + f"Could not find checkpoints for checkpoint NS '{checkpoint_ns}'" + ) + if checkpoint_id: + error_msg += f" and checkpoint ID '{checkpoint_id}'" raise ValueError(error_msg) state_snapshot = self._assemble_state_snapshot_hierarchy( - thread_id, thread_id_to_state_snapshots + checkpoint_ns, checkpoint_ns_to_state_snapshots ) return state_snapshot @@ -502,7 +518,9 @@ class Pregel( raise ValueError("No checkpointer set") if include_subgraph_state: - checkpoint_tuples = self.checkpointer.alist(config, as_prefix=True) + checkpoint_tuples = self.checkpointer.alist( + config, include_nested_checkpoints=True + ) else: async def alist_checkpoints(): @@ -510,34 +528,45 @@ class Pregel( checkpoint_tuples = alist_checkpoints() - thread_id = config["configurable"]["thread_id"] - thread_ts = config["configurable"].get("thread_ts") - thread_id_to_thread_ts: dict[str, str] = {} - thread_id_to_state_snapshots: dict[str, StateSnapshot] = {} + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + checkpoint_id = config["configurable"].get("checkpoint_id") + checkpoint_ns_to_checkpoint_id: dict[str, str] = {} + checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} async for checkpoint_tuple in checkpoint_tuples: - checkpoint_thread_id = checkpoint_tuple.config["configurable"]["thread_id"] - checkpoint_thread_ts = checkpoint_tuple.config["configurable"]["thread_ts"] - if thread_ts and thread_ts != checkpoint_thread_ts: + saved_checkpoint_ns = checkpoint_tuple.config["configurable"][ + "checkpoint_ns" + ] + saved_checkpoint_id = checkpoint_tuple.config["configurable"][ + "checkpoint_id" + ] + if checkpoint_id and checkpoint_id != saved_checkpoint_id: continue - existing_thread_ts = thread_id_to_thread_ts.get(checkpoint_thread_id) - # keep only most recent thread_ts - if existing_thread_ts is None or checkpoint_thread_ts > existing_thread_ts: - state_snapshot = await self._prepare_state_snapshot_async( - checkpoint_tuple, config - ) - thread_id_to_state_snapshots[checkpoint_thread_id] = state_snapshot - thread_id_to_thread_ts[checkpoint_thread_id] = checkpoint_thread_ts + existing_checkpoint_id = checkpoint_ns_to_checkpoint_id.get( + saved_checkpoint_ns + ) + # keep only most recent checkpoint_id + if ( + existing_checkpoint_id is None + or saved_checkpoint_id > existing_checkpoint_id + ): + state_snapshot = self._prepare_state_snapshot(checkpoint_tuple, config) + checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot + checkpoint_ns_to_checkpoint_id[ + saved_checkpoint_ns + ] = saved_checkpoint_id - if not thread_id_to_state_snapshots: - error_msg = f"Could not find checkpoints for thread ID '{thread_id}'" - if thread_ts: - error_msg += f" and thread TS '{thread_ts}'" + if not checkpoint_ns_to_state_snapshots: + error_msg = ( + f"Could not find checkpoints for checkpoint NS '{checkpoint_ns}'" + ) + if checkpoint_id: + error_msg += f" and checkpoint ID '{checkpoint_id}'" raise ValueError(error_msg) state_snapshot = self._assemble_state_snapshot_hierarchy( - thread_id, thread_id_to_state_snapshots + checkpoint_ns, checkpoint_ns_to_state_snapshots ) return state_snapshot @@ -584,7 +613,7 @@ class Pregel( -1, for_execution=False, ) - + yield StateSnapshot( read_channels(channels, self.stream_channels_asis), tuple(t.name for t in next_tasks), @@ -620,10 +649,11 @@ class Pregel( ) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter): if include_subgraph_state: state_snapshot = await self.aget_state( - config, include_subgraph_state=True) + config, include_subgraph_state=True + ) yield state_snapshot else: - async with AsyncChannelsManager( + async with AsyncChannelsManager( { k: LastValue(None) if isinstance(c, Context) else c for k, c in self.channels.items() diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index af7f0920a..1530fcda4 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9093,6 +9093,456 @@ 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 + + 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, include_subgraph_state=False) == StateSnapshot( + values={"my_key": "hi my value"}, + 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=None, + ) + assert app.get_state(config, include_subgraph_state=True) == StateSnapshot( + values={"my_key": "hi my value"}, + 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"}, + next=(), + 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, include_subgraph_state=True)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + 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"}, + next=(), + 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"}, + 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={}, + 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, include_subgraph_state=True) == StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + 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={ + "inner": StateSnapshot( + values={"my_key": "hi my value here and there"}, + 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, + ) + }, + ) + assert list(app.get_state_history(config, include_subgraph_state=True)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + 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"}, + 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"}, + 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(), + } + }, + # TODO: this is likely very confusing for an end user, and we'll probably need to update this. + # right now this is happening due to us overwriting the + # subgraph snapshot after we finish the graph with while the checkpoint_id + # is the same as when we interrupted + subgraph_state_snapshots={ + "inner": StateSnapshot( + values={"my_key": "hi my value here and there"}, + 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"}, + 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={}, + 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, + ), + ] + + 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 67b62941b..064c48448 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7600,6 +7600,458 @@ 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 + + def inner_1(state: InnerState): + return { + "my_key": state["my_key"] + " here", + "my_other_key": state["my_key"], + } + + def inner_2(state: InnerState): + return { + "my_key": state["my_key"] + " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + + def outer_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + def outer_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("outer_1", outer_1) + graph.add_node("inner", inner.compile(interrupt_before=["inner_2"])) + graph.add_node("outer_2", outer_2) + graph.set_entry_point("outer_1") + graph.add_edge("outer_1", "inner") + graph.add_edge("inner", "outer_2") + graph.set_finish_point("outer_2") + + app = graph.compile(checkpointer=checkpointer) + + 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, include_subgraph_state=False) == StateSnapshot( + values={"my_key": "hi my value"}, + 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=None, + ) + assert app.get_state(config, include_subgraph_state=True) == StateSnapshot( + values={"my_key": "hi my value"}, + 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"}, + next=(), + 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, include_subgraph_state=True)) == [ + StateSnapshot( + values={"my_key": "hi my value"}, + 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"}, + next=(), + 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"}, + 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={}, + 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, include_subgraph_state=True) == StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + 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={ + "inner": StateSnapshot( + values={"my_key": "hi my value here and there"}, + 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, + ) + }, + ) + assert [ + s async for s in app.aget_state_history(config, include_subgraph_state=True) + ] == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + 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"}, + 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"}, + 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(), + } + }, + # TODO: this is likely very confusing for an end user, and we'll probably need to update this. + # right now this is happening due to us overwriting the + # subgraph snapshot after we finish the graph with while the checkpoint_id + # is the same as when we interrupted + subgraph_state_snapshots={ + "inner": StateSnapshot( + values={"my_key": "hi my value here and there"}, + 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"}, + 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={}, + 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, + ), + ] + + 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.