diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 5c89ed27f..bc4fd48a0 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -506,26 +506,19 @@ class Pregel( yield runnable yield from runnable.subgraphs - def get_state( - self, config: RunnableConfig, *, include_subgraph_state: bool = False - ) -> StateSnapshot: + def get_state(self, config: RunnableConfig) -> StateSnapshot: """Get the current state of the graph.""" if not self.checkpointer: raise ValueError("No checkpointer set") checkpoint_tuple = self.checkpointer.get_tuple(config) - if include_subgraph_state: - checkpoint_tuples = self.checkpointer.list(config) - else: - checkpoint_tuples = iter([checkpoint_tuple] if checkpoint_tuple else []) - checkpoint_config = checkpoint_tuple.config if checkpoint_tuple else config checkpoint_ns = checkpoint_config["configurable"].get("checkpoint_ns", "") checkpoint_id = checkpoint_config["configurable"].get("checkpoint_id") checkpoint_ns_to_checkpoint_id: dict[str, str] = {} checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} checkpoint_ns_to_graph: dict[str, Pregel] = {} - for checkpoint_tuple in checkpoint_tuples: + for checkpoint_tuple in self.checkpointer.list(config): saved_checkpoint_ns = checkpoint_tuple.config["configurable"][ "checkpoint_ns" ] @@ -567,31 +560,19 @@ class Pregel( ) return state_snapshot - async def aget_state( - self, config: RunnableConfig, *, include_subgraph_state: bool = False - ) -> StateSnapshot: + async def aget_state(self, config: RunnableConfig) -> StateSnapshot: """Get the current state of the graph.""" if not self.checkpointer: raise ValueError("No checkpointer set") checkpoint_tuple = await self.checkpointer.aget_tuple(config) - if include_subgraph_state: - checkpoint_tuples = self.checkpointer.alist(config) - else: - - async def alist_checkpoints(): - if checkpoint_tuple: - yield checkpoint_tuple - - checkpoint_tuples = alist_checkpoints() - checkpoint_config = checkpoint_tuple.config if checkpoint_tuple else config checkpoint_ns = checkpoint_config["configurable"].get("checkpoint_ns", "") checkpoint_id = checkpoint_config["configurable"].get("checkpoint_id") checkpoint_ns_to_checkpoint_id: dict[str, str] = {} checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} checkpoint_ns_to_graph: dict[str, Pregel] = {} - async for checkpoint_tuple in checkpoint_tuples: + async for checkpoint_tuple in self.checkpointer.alist(config): saved_checkpoint_ns = checkpoint_tuple.config["configurable"][ "checkpoint_ns" ] @@ -640,7 +621,6 @@ class Pregel( filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, - include_subgraph_state: bool = False, ) -> Iterator[StateSnapshot]: """Get the history of the state of the graph.""" if not self.checkpointer: @@ -662,20 +642,8 @@ class Pregel( # only list root checkpoints here continue - if include_subgraph_state: - state_snapshot = self.get_state( - checkpoint_tuple.config, include_subgraph_state=True - ) - yield state_snapshot - else: - graph = _get_subgraph( - self, - checkpoint_tuple.config["configurable"]["checkpoint_ns"], - ) - yield _prepare_state_snapshot( - checkpoint_tuple, - graph, - ) + state_snapshot = self.get_state(checkpoint_tuple.config) + yield state_snapshot async def aget_state_history( self, @@ -684,7 +652,6 @@ class Pregel( filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, - include_subgraph_state: bool = False, ) -> AsyncIterator[StateSnapshot]: """Get the history of the state of the graph.""" if not self.checkpointer: @@ -706,20 +673,8 @@ class Pregel( # only list root checkpoints here continue - if include_subgraph_state: - state_snapshot = await self.aget_state( - checkpoint_tuple.config, include_subgraph_state=True - ) - yield state_snapshot - else: - graph = _get_subgraph( - self, - checkpoint_tuple.config["configurable"]["checkpoint_ns"], - ) - yield await _prepare_state_snapshot_async( - checkpoint_tuple, - graph, - ) + state_snapshot = await self.aget_state(checkpoint_tuple.config) + yield state_snapshot def update_state( self, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 7bdc357d7..24bfca34d 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7783,6 +7783,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", + }, + 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"}, @@ -7900,6 +7935,41 @@ 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=(), + 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"}, @@ -8071,6 +8141,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", + }, + 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"}, @@ -8193,6 +8298,41 @@ 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=(), + 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"}, @@ -8245,7 +8385,6 @@ 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"}, @@ -8270,6 +8409,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", + }, + 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"}, @@ -8315,6 +8489,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"}, @@ -8363,6 +8538,41 @@ 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=(), + 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"}, @@ -8482,6 +8692,41 @@ 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=(), + 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"}, @@ -8554,6 +8799,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", + }, + 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"}, @@ -8672,6 +8952,40 @@ 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", + }, + 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"}, @@ -8696,6 +9010,40 @@ 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", + }, + 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"}, @@ -8815,6 +9163,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", + }, + 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"}, @@ -8839,6 +9222,41 @@ 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=(), + 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"}, @@ -9165,32 +9583,7 @@ def test_nested_graph_state( 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( + assert app.get_state(config) == StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), config={ @@ -9246,7 +9639,7 @@ def test_nested_graph_state( ) }, ) - assert list(app.get_state_history(config, include_subgraph_state=True)) == [ + assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), @@ -9349,7 +9742,7 @@ def test_nested_graph_state( ] 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( + assert app.get_state(config) == StateSnapshot( values={"my_key": "hi my value here and there and back again"}, next=(), config={ @@ -9421,7 +9814,6 @@ def test_nested_graph_state( "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], } }, - include_subgraph_state=True, ) == StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), @@ -9448,7 +9840,7 @@ def test_nested_graph_state( subgraph_state_snapshots={"inner": child_snapshot}, ) # test full history at the end - assert list(app.get_state_history(config, include_subgraph_state=True)) == [ + assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, next=(), @@ -9665,31 +10057,6 @@ def test_doubly_nested_graph_state( 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"}, - 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=None, - ) - assert app.get_state(config, include_subgraph_state=True) == StateSnapshot( values={"my_key": "hi my value"}, next=("child",), config={ @@ -9763,7 +10130,7 @@ def test_doubly_nested_graph_state( }, ) app.invoke(None, config, debug=True) - assert app.get_state(config, include_subgraph_state=True) == StateSnapshot( + assert app.get_state(config) == StateSnapshot( values={"my_key": "hi my value here and there and back again"}, next=(), config={ @@ -9821,7 +10188,6 @@ def test_doubly_nested_graph_state( # test getting child snapshot child_snapshot = app.get_state( {"configurable": {"thread_id": "1", "checkpoint_ns": "child"}}, - include_subgraph_state=True, ) assert child_snapshot == StateSnapshot( values={"my_key": "hi my value here and there"}, @@ -9856,7 +10222,6 @@ def test_doubly_nested_graph_state( "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], } }, - include_subgraph_state=True, ) == StateSnapshot( values={"my_key": "hi my value"}, next=("child",), @@ -9935,7 +10300,7 @@ def test_send_to_nested_graphs( "subjects": ["cats", "dogs"], "jokes": [], } - actual_snapshot = graph.get_state(config, include_subgraph_state=True) + 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: @@ -9977,7 +10342,7 @@ def test_send_to_nested_graphs( "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], } - actual_snapshot = graph.get_state(config, include_subgraph_state=True) + actual_snapshot = graph.get_state(config) expected_snapshot = StateSnapshot( values={ "subjects": ["cats", "dogs"], @@ -10013,7 +10378,7 @@ def test_send_to_nested_graphs( assert actual_snapshot == expected_snapshot # test full history - actual_history = list(graph.get_state_history(config, include_subgraph_state=True)) + actual_history = list(graph.get_state_history(config)) # get subgraph node state for expected history subgraph_state_snapshots = { diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 87d9240e0..7aae5308b 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6278,6 +6278,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", + }, + 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"}, @@ -6395,6 +6430,41 @@ 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=(), + 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"}, @@ -6571,6 +6641,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", + }, + 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"}, @@ -6693,6 +6798,41 @@ 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=(), + 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"}, @@ -6772,6 +6912,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", + }, + 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"}, @@ -6817,6 +6992,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"}, @@ -6865,6 +7041,41 @@ 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=(), + 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"}, @@ -6984,6 +7195,41 @@ 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=(), + 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"}, @@ -7055,6 +7301,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", + }, + 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"}, @@ -7174,6 +7455,40 @@ 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", + }, + 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"}, @@ -7198,6 +7513,40 @@ 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", + }, + 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"}, @@ -7317,6 +7666,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", + }, + 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"}, @@ -7341,6 +7725,41 @@ 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=(), + 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"}, @@ -7672,32 +8091,7 @@ async def test_nested_graph_state( 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 await app.aget_state(config, include_subgraph_state=True) == StateSnapshot( + assert await app.aget_state(config) == StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), config={ @@ -7756,9 +8150,7 @@ async def test_nested_graph_state( ) }, ) - assert [ - s async for s in app.aget_state_history(config, include_subgraph_state=True) - ] == [ + assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), @@ -7861,7 +8253,7 @@ async def test_nested_graph_state( ] 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( + assert await app.aget_state(config) == StateSnapshot( values={"my_key": "hi my value here and there and back again"}, next=(), config={ @@ -7933,7 +8325,6 @@ async def test_nested_graph_state( "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], } }, - include_subgraph_state=True, ) == StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), @@ -7960,9 +8351,7 @@ async def test_nested_graph_state( subgraph_state_snapshots={"inner": child_snapshot}, ) # test full history at the end - assert [ - s async for s in app.aget_state_history(config, include_subgraph_state=True) - ] == [ + assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, next=(), @@ -8179,31 +8568,6 @@ async def test_doubly_nested_graph_state( 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"}, - 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=None, - ) - assert await app.aget_state(config, include_subgraph_state=True) == StateSnapshot( values={"my_key": "hi my value"}, next=("child",), config={ @@ -8277,7 +8641,7 @@ async def test_doubly_nested_graph_state( }, ) await app.ainvoke(None, config, debug=True) - assert await app.aget_state(config, include_subgraph_state=True) == StateSnapshot( + assert await app.aget_state(config) == StateSnapshot( values={"my_key": "hi my value here and there and back again"}, next=(), config={ @@ -8335,7 +8699,6 @@ async def test_doubly_nested_graph_state( # test getting child snapshot child_snapshot = await app.aget_state( {"configurable": {"thread_id": "1", "checkpoint_ns": "child"}}, - include_subgraph_state=True, ) assert child_snapshot == StateSnapshot( values={"my_key": "hi my value here and there"}, @@ -8370,7 +8733,6 @@ async def test_doubly_nested_graph_state( "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], } }, - include_subgraph_state=True, ) == StateSnapshot( values={"my_key": "hi my value"}, next=("child",), @@ -8449,7 +8811,7 @@ async def test_send_to_nested_graphs( "subjects": ["cats", "dogs"], "jokes": [], } - actual_snapshot = await graph.aget_state(config, include_subgraph_state=True) + 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: @@ -8490,7 +8852,7 @@ async def test_send_to_nested_graphs( "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], } - actual_snapshot = await graph.aget_state(config, include_subgraph_state=True) + actual_snapshot = await graph.aget_state(config) expected_snapshot = StateSnapshot( values={ "subjects": ["cats", "dogs"], @@ -8526,9 +8888,7 @@ async def test_send_to_nested_graphs( assert actual_snapshot == expected_snapshot # test full history - actual_history = [ - c async for c in graph.aget_state_history(config, include_subgraph_state=True) - ] + 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(