From ae696d4f3039cb8286bf01325e6970da5206fc63 Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 23 Jul 2024 21:40:04 -0400 Subject: [PATCH] add sync history --- libs/langgraph/langgraph/pregel/__init__.py | 59 ++++-- libs/langgraph/tests/test_pregel.py | 215 ++++++++++++++++++++ 2 files changed, 253 insertions(+), 21 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index f1ab7d5b6..6287c713b 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -502,6 +502,7 @@ 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: @@ -514,28 +515,44 @@ class Pregel( for config, checkpoint, metadata, parent_config, _ in self.checkpointer.list( config, before=before, limit=limit, filter=filter ): - with ChannelsManager( - self.channels, checkpoint, config - ) as channels, ManagedValuesManager( - self.managed_values_dict, ensure_config(config), self - ) as managed: - next_tasks = prepare_next_tasks( - checkpoint, - self.nodes, - channels, - managed, - config, - -1, - for_execution=False, - ) - yield StateSnapshot( - read_channels(channels, self.stream_channels_asis), - tuple(name for name, _ in next_tasks), - config, - metadata, - checkpoint["ts"], - parent_config, + # is there a way to do this more efficiently? + if include_subgraph_state: + checkpoint_tuples = self.checkpointer.list_subgraph_checkpoints(config) + + thread_id_to_state_snapshots: dict[str, StateSnapshot] = { + checkpoint.config["configurable"][ + "thread_id" + ]: self._prepare_state_snapshot(checkpoint, config) + for checkpoint in checkpoint_tuples + } + thread_id = config["configurable"]["thread_id"] + state_snapshot = self._assemble_state_snapshot_hierarchy( + thread_id, thread_id_to_state_snapshots ) + yield state_snapshot + else: + with ChannelsManager( + self.channels, checkpoint, config + ) as channels, ManagedValuesManager( + self.managed_values_dict, ensure_config(config), self + ) as managed: + next_tasks = prepare_next_tasks( + checkpoint, + self.nodes, + channels, + managed, + config, + -1, + for_execution=False, + ) + yield StateSnapshot( + read_channels(channels, self.stream_channels_asis), + tuple(name for name, _ in next_tasks), + config, + metadata, + checkpoint["ts"], + parent_config, + ) async def aget_state_history( self, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index fea1e29b8..0a86aee21 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8786,6 +8786,76 @@ def test_nested_graph_state( ) }, ) + 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", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + }, + subgraph_state_snapshots={ + "inner": StateSnapshot( + values={"my_key": "hi my value here"}, + next=(), + config={ + "configurable": { + "thread_id": "1__inner", + "thread_ts": 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__inner", + "thread_ts": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1", "thread_ts": AnyStr()} + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={"configurable": {"thread_id": "1", "thread_ts": 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( @@ -8826,6 +8896,151 @@ def test_nested_graph_state( ) }, ) + assert list(app.get_state_history(config, include_subgraph_state=False)) == [ + StateSnapshot( + values={"my_key": "hi my value here and there and back again"}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": 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", + "thread_ts": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=("outer_2",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": 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", + "thread_ts": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={"my_key": "hi my value"}, + next=("inner",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"outer_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "thread_ts": 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 thread_ts + # 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__inner", + "thread_ts": 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__inner", + "thread_ts": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ), + StateSnapshot( + values={"my_key": "my value"}, + next=("outer_1",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + StateSnapshot( + values={}, + next=("__start__",), + config={ + "configurable": { + "thread_id": "1", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "input", + "writes": {"my_key": "my value"}, + "step": -1, + }, + created_at=AnyStr(), + parent_config=None, + subgraph_state_snapshots=None, + ), + ] finally: if hasattr(checkpointer, "__exit__"): checkpointer.__exit__(None, None, None)