From 261cdf88a5aea13ea3c958af2e7c51d753b5a0b0 Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 23 Jul 2024 16:43:26 -0400 Subject: [PATCH 01/41] langgraph: update get_state to handle nested subgraph state --- libs/langgraph/langgraph/checkpoint/base.py | 5 ++ libs/langgraph/langgraph/checkpoint/memory.py | 33 +++++++++ libs/langgraph/langgraph/checkpoint/sqlite.py | 43 ++++++++++++ libs/langgraph/langgraph/constants.py | 2 + libs/langgraph/langgraph/pregel/__init__.py | 67 +++++++++++++++---- libs/langgraph/langgraph/pregel/algo.py | 5 +- libs/langgraph/langgraph/pregel/types.py | 2 + libs/langgraph/tests/test_pregel.py | 44 ++++++------ 8 files changed, 166 insertions(+), 35 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/base.py b/libs/langgraph/langgraph/checkpoint/base.py index b6694db89..fc53833c1 100644 --- a/libs/langgraph/langgraph/checkpoint/base.py +++ b/libs/langgraph/langgraph/checkpoint/base.py @@ -231,6 +231,11 @@ class BaseCheckpointSaver(ABC): """ raise NotImplementedError + def list_subgraph_checkpoints( + self, config: RunnableConfig + ) -> Iterator[CheckpointTuple]: + raise NotImplementedError + def put( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index 72b8c93db..55246be4b 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -171,6 +171,39 @@ class MemorySaver(BaseCheckpointSaver): else None, ) + def list_subgraph_checkpoints( + self, config: RunnableConfig + ) -> Iterator[CheckpointTuple]: + thread_id_prefix = config["configurable"]["thread_id"] + matching_thread_ids = [ + key for key in self.storage.keys() if key.startswith(thread_id_prefix) + ] + for thread_id in matching_thread_ids: + ts = config["configurable"].get("thread_ts") + if not ts: + if checkpoints := self.storage[thread_id]: + ts = max(checkpoints.keys()) + + if saved := self.storage[thread_id].get(ts): + checkpoint, metadata, parent_ts = saved + writes = self.writes[(thread_id, ts)] + yield CheckpointTuple( + config={"configurable": {"thread_id": thread_id, "thread_ts": ts}}, + checkpoint=self.serde.loads(checkpoint), + metadata=self.serde.loads(metadata), + pending_writes=[ + (id, c, self.serde.loads(v)) for id, c, v in writes + ], + parent_config={ + "configurable": { + "thread_id": thread_id, + "thread_ts": parent_ts, + } + } + if parent_ts + else None, + ) + def put( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/checkpoint/sqlite.py b/libs/langgraph/langgraph/checkpoint/sqlite.py index eee6e05c7..f6bf4b857 100644 --- a/libs/langgraph/langgraph/checkpoint/sqlite.py +++ b/libs/langgraph/langgraph/checkpoint/sqlite.py @@ -357,6 +357,49 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): ), ) + def list_subgraph_checkpoints( + self, config: RunnableConfig + ) -> Iterator[CheckpointTuple]: + with self.cursor(transaction=False) as cur: + if config["configurable"].get("thread_ts"): + cur.execute( + "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id LIKE ? || '%' AND thread_ts = ?", + ( + str(config["configurable"]["thread_id"]), + str(config["configurable"]["thread_ts"]), + ), + ) + else: + cur.execute( + """SELECT checkpoints.thread_id, checkpoints.thread_ts, checkpoints.parent_ts, checkpoints.checkpoint, checkpoints.metadata + FROM checkpoints + INNER JOIN ( + SELECT thread_id, MAX(thread_ts) as thread_ts + FROM checkpoints + WHERE thread_id LIKE ? || '%' + GROUP BY thread_id + ) latest_checkpoints + ON checkpoints.thread_id = latest_checkpoints.thread_id AND checkpoints.thread_ts = latest_checkpoints.thread_ts + ORDER BY checkpoints.thread_id, checkpoints.thread_ts DESC""", + (str(config["configurable"]["thread_id"]),), + ) + for thread_id, thread_ts, parent_ts, value, metadata in cur: + yield CheckpointTuple( + {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}}, + self.serde.loads(value), + self.serde.loads(metadata) if metadata is not None else {}, + ( + { + "configurable": { + "thread_id": thread_id, + "thread_ts": parent_ts, + } + } + if parent_ts + else None + ), + ) + def put( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index f3aeb6a2e..edab6d6b7 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -18,6 +18,8 @@ RESERVED = { } TAG_HIDDEN = "langsmith:hidden" +THREAD_ID_SEPARATOR = "__" + START = "__start__" END = "__end__" diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index c0edcacc3..c04f3e9e8 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -59,6 +59,7 @@ from langgraph.channels.manager import ( ) from langgraph.checkpoint.base import ( BaseCheckpointSaver, + CheckpointTuple, copy_checkpoint, empty_checkpoint, ) @@ -68,6 +69,7 @@ from langgraph.constants import ( CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, INTERRUPT, + THREAD_ID_SEPARATOR, ) from langgraph.errors import GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ( @@ -350,12 +352,9 @@ class Pregel( if is_managed_value(v) } - def get_state(self, config: RunnableConfig) -> StateSnapshot: - """Get the current state of the graph.""" - if not self.checkpointer: - raise ValueError("No checkpointer set") - - saved = self.checkpointer.get_tuple(config) + def _prepare_state_snapshot( + self, saved: CheckpointTuple, config: RunnableConfig + ) -> StateSnapshot: checkpoint = saved.checkpoint if saved else empty_checkpoint() config = saved.config if saved else config with ChannelsManager( @@ -373,14 +372,58 @@ class Pregel( for_execution=False, ) return StateSnapshot( - read_channels(channels, self.stream_channels_asis), - tuple(name for name, _ 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, + values=read_channels(channels, self.stream_channels_asis), + next=tuple(name for name, _ in next_tasks), + config=saved.config if saved else config, + metadata=saved.metadata if saved else None, + created_at=saved.checkpoint["ts"] if saved else None, + parent_config=saved.parent_config if saved else None, ) + @staticmethod + def _assemble_state_snapshot_hierarchy( + root_thread_id: str, subgraph_state_snapshots: dict[str, StateSnapshot] + ) -> StateSnapshot: + thread_ids_to_visit = sorted( + subgraph_state_snapshots.keys(), + key=lambda x: len(x.split(THREAD_ID_SEPARATOR)), + ) + while thread_ids_to_visit: + thread_id = thread_ids_to_visit.pop() + state_snapshot = subgraph_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 THREAD_ID_SEPARATOR in parent_thread_id: + parent_subgraph_snapshots = ( + subgraph_state_snapshots[parent_thread_id].subgraph_state_snapshots + or {} + ) + parent_subgraph_snapshots[subgraph_node] = state_snapshot + subgraph_state_snapshots[parent_thread_id] = subgraph_state_snapshots[ + parent_thread_id + ]._replace(subgraph_state_snapshots=parent_subgraph_snapshots) + + state_snapshot = subgraph_state_snapshots.pop(root_thread_id) + return state_snapshot + + def get_state(self, config: RunnableConfig) -> StateSnapshot: + """Get the current state of the graph.""" + if not self.checkpointer: + raise ValueError("No checkpointer set") + + subgraph_state_snapshots: dict[str, StateSnapshot] = { + checkpoint.config["configurable"][ + "thread_id" + ]: self._prepare_state_snapshot(checkpoint, config) + for checkpoint in self.checkpointer.list_subgraph_checkpoints(config) + } + + thread_id = config["configurable"]["thread_id"] + state_snapshot = self._assemble_state_snapshot_hierarchy( + thread_id, subgraph_state_snapshots + ) + return state_snapshot + async def aget_state(self, config: RunnableConfig) -> StateSnapshot: """Get the current state of the graph.""" if not self.checkpointer: diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 13ce7e183..00d9e7809 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -36,6 +36,7 @@ from langgraph.constants import ( RESERVED, TAG_HIDDEN, TASKS, + THREAD_ID_SEPARATOR, Send, ) from langgraph.errors import EmptyChannelError, InvalidUpdateError @@ -345,7 +346,9 @@ def prepare_next_tasks( if parent_thread_id := config.get("configurable", {}).get( "thread_id" ): - thread_id: Optional[str] = f"{parent_thread_id}-{name}" + thread_id: Optional[ + str + ] = f"{parent_thread_id}{THREAD_ID_SEPARATOR}{name}" else: thread_id = None writes = deque() diff --git a/libs/langgraph/langgraph/pregel/types.py b/libs/langgraph/langgraph/pregel/types.py index 19d3c3301..d0648cf3b 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -85,6 +85,8 @@ class StateSnapshot(NamedTuple): """Timestamp of snapshot creation""" parent_config: Optional[RunnableConfig] = None """Config used to fetch the parent snapshot, if any""" + subgraph_state_snapshots: Optional[dict[str, "StateSnapshot"]] = None + """State snapshots of subgraphs represented as a mapping from thread ID suffix to snapshot.""" All = Literal["*"] diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index db902eca0..b218fabda 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -568,53 +568,53 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: ) # start execution, stop at inbox - assert app.invoke(2, {"configurable": {"thread_id": 1}}) is None + assert app.invoke(2, {"configurable": {"thread_id": "1"}}) is None # inbox == 3 - checkpoint = memory.get({"configurable": {"thread_id": 1}}) + checkpoint = memory.get({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"]["inbox"] == 3 # resume execution, finish - assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 4 + assert app.invoke(None, {"configurable": {"thread_id": "1"}}) == 4 # start execution again, stop at inbox - assert app.invoke(20, {"configurable": {"thread_id": 1}}) is None + assert app.invoke(20, {"configurable": {"thread_id": "1"}}) is None # inbox == 21 - checkpoint = memory.get({"configurable": {"thread_id": 1}}) + checkpoint = memory.get({"configurable": {"thread_id": "1"}}) assert checkpoint is not None assert checkpoint["channel_values"]["inbox"] == 21 # send a new value in, interrupting the previous execution - assert app.invoke(3, {"configurable": {"thread_id": 1}}) is None - assert app.invoke(None, {"configurable": {"thread_id": 1}}) == 5 + assert app.invoke(3, {"configurable": {"thread_id": "1"}}) is None + assert app.invoke(None, {"configurable": {"thread_id": "1"}}) == 5 # start execution again, stopping at inbox - assert app.invoke(20, {"configurable": {"thread_id": 2}}) is None + assert app.invoke(20, {"configurable": {"thread_id": "2"}}) is None # inbox == 21 - snapshot = app.get_state({"configurable": {"thread_id": 2}}) + snapshot = app.get_state({"configurable": {"thread_id": "2"}}) assert snapshot.values["inbox"] == 21 assert snapshot.next == ("two",) # update the state, resume - app.update_state({"configurable": {"thread_id": 2}}, 25, as_node="one") - assert app.invoke(None, {"configurable": {"thread_id": 2}}) == 26 + app.update_state({"configurable": {"thread_id": "2"}}, 25, as_node="one") + assert app.invoke(None, {"configurable": {"thread_id": "2"}}) == 26 # no pending tasks - snapshot = app.get_state({"configurable": {"thread_id": 2}}) + snapshot = app.get_state({"configurable": {"thread_id": "2"}}) assert snapshot.next == () # list history - thread1 = {"configurable": {"thread_id": 1}} + thread1 = {"configurable": {"thread_id": "1"}} assert [c for c in app.get_state_history(thread1)] == [ StateSnapshot( values={"inbox": 4, "output": 5, "input": 3}, next=(), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -627,7 +627,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("two",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -640,7 +640,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("one",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -653,7 +653,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("two",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -666,7 +666,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("one",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -679,7 +679,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=(), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -692,7 +692,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("two",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -705,7 +705,7 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: next=("one",), config={ "configurable": { - "thread_id": 1, + "thread_id": "1", "thread_ts": AnyStr(), } }, @@ -1080,7 +1080,7 @@ def test_pending_writes_resume(checkpointer: BaseCheckpointSaver) -> None: builder.add_edge(START, "two") graph = builder.compile(checkpointer=checkpointer) - thread1: RunnableConfig = {"configurable": {"thread_id": 1}} + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} with pytest.raises(ConnectionError, match="I'm not good"): graph.invoke({"value": 1}, thread1) From 92ae8f48177ddc75e877881bbe82f62ad21868d7 Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 23 Jul 2024 19:51:44 -0400 Subject: [PATCH 02/41] async methods --- .../langgraph/checkpoint/aiosqlite.py | 4 +- libs/langgraph/langgraph/checkpoint/memory.py | 13 ++++ libs/langgraph/langgraph/pregel/__init__.py | 68 ++++++++++++------- libs/langgraph/tests/test_pregel_async.py | 4 +- 4 files changed, 59 insertions(+), 30 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/aiosqlite.py b/libs/langgraph/langgraph/checkpoint/aiosqlite.py index 21db080d1..401982fdf 100644 --- a/libs/langgraph/langgraph/checkpoint/aiosqlite.py +++ b/libs/langgraph/langgraph/checkpoint/aiosqlite.py @@ -352,7 +352,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): ) -> AsyncIterator[CheckpointTuple]: async with self.conn.cursor() as cur: if config["configurable"].get("thread_ts"): - cur.execute( + await cur.execute( "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id LIKE ? || '%' AND thread_ts = ?", ( str(config["configurable"]["thread_id"]), @@ -360,7 +360,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): ), ) else: - cur.execute( + await cur.execute( """SELECT checkpoints.thread_id, checkpoints.thread_ts, checkpoints.parent_ts, checkpoints.checkpoint, checkpoints.metadata FROM checkpoints INNER JOIN ( diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index 55246be4b..348aacdf0 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -311,6 +311,19 @@ class MemorySaver(BaseCheckpointSaver): else: break + async def alist_subgraph_checkpoints( + self, config: RunnableConfig + ) -> AsyncIterator[CheckpointTuple]: + loop = asyncio.get_running_loop() + iter = await loop.run_in_executor(None, self.list_subgraph_checkpoints, config) + while True: + # handling StopIteration exception inside coroutine won't work + # as expected, so using next() with default value to break the loop + if item := await loop.run_in_executor(None, next, iter, None): + yield item + else: + break + async def aput( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index c04f3e9e8..e3241eae4 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -380,6 +380,34 @@ class Pregel( parent_config=saved.parent_config if saved else None, ) + async def _prepare_state_snapshot_async( + self, saved: CheckpointTuple, config: RunnableConfig + ) -> StateSnapshot: + checkpoint = saved.checkpoint if saved else empty_checkpoint() + config = saved.config if saved else config + async with AsyncChannelsManager( + self.channels, checkpoint, config + ) as channels, AsyncManagedValuesManager( + 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, + ) + return StateSnapshot( + values=read_channels(channels, self.stream_channels_asis), + next=tuple(name for name, _ in next_tasks), + config=saved.config if saved else config, + metadata=saved.metadata if saved else None, + created_at=saved.checkpoint["ts"] if saved else None, + parent_config=saved.parent_config if saved else None, + ) + @staticmethod def _assemble_state_snapshot_hierarchy( root_thread_id: str, subgraph_state_snapshots: dict[str, StateSnapshot] @@ -403,7 +431,9 @@ class Pregel( parent_thread_id ]._replace(subgraph_state_snapshots=parent_subgraph_snapshots) - state_snapshot = subgraph_state_snapshots.pop(root_thread_id) + state_snapshot = subgraph_state_snapshots.pop(root_thread_id, None) + if state_snapshot is None: + raise ValueError(f"Missing snapshot for thread ID '{root_thread_id}'") return state_snapshot def get_state(self, config: RunnableConfig) -> StateSnapshot: @@ -429,32 +459,18 @@ class Pregel( if not self.checkpointer: raise ValueError("No checkpointer set") - saved = await self.checkpointer.aget_tuple(config) - checkpoint = saved.checkpoint if saved else empty_checkpoint() + subgraph_state_snapshots: dict[str, StateSnapshot] = { + checkpoint.config["configurable"][ + "thread_id" + ]: await self._prepare_state_snapshot_async(checkpoint, config) + async for checkpoint in self.checkpointer.alist_subgraph_checkpoints(config) + } - config = saved.config if saved else config - async with AsyncChannelsManager( - self.channels, checkpoint, config - ) as channels, AsyncManagedValuesManager( - 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, - ) - return StateSnapshot( - read_channels(channels, self.stream_channels_asis), - tuple(name for name, _ 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, - ) + thread_id = config["configurable"]["thread_id"] + state_snapshot = self._assemble_state_snapshot_hierarchy( + thread_id, subgraph_state_snapshots + ) + return state_snapshot def get_state_history( self, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 76a182d69..7d88f286e 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -1204,7 +1204,7 @@ async def test_pending_writes_resume(checkpointer: BaseCheckpointSaver) -> None: builder.add_edge(START, "two") graph = builder.compile(checkpointer=checkpointer) - thread1: RunnableConfig = {"configurable": {"thread_id": 1}} + thread1: RunnableConfig = {"configurable": {"thread_id": "1"}} with pytest.raises(ValueError, match="I'm not good"): await graph.ainvoke({"value": 1}, thread1) @@ -6079,7 +6079,7 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert times_called == 1 -@pytest.mark.repeat(10) +# @pytest.mark.repeat(10) @pytest.mark.parametrize( "checkpointer_fct", [ From e615aabf14385152c021ce1686cdff4318bc0aff Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 23 Jul 2024 20:14:53 -0400 Subject: [PATCH 03/41] cleanup names + make subgraph state optional --- libs/langgraph/langgraph/pregel/__init__.py | 49 +++++++++++++-------- 1 file changed, 30 insertions(+), 19 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index e3241eae4..511a5fd9c 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -410,65 +410,76 @@ class Pregel( @staticmethod def _assemble_state_snapshot_hierarchy( - root_thread_id: str, subgraph_state_snapshots: dict[str, StateSnapshot] + root_thread_id: str, thread_id_to_state_snapshots: dict[str, StateSnapshot] ) -> StateSnapshot: thread_ids_to_visit = sorted( - subgraph_state_snapshots.keys(), + thread_id_to_state_snapshots.keys(), key=lambda x: len(x.split(THREAD_ID_SEPARATOR)), ) while thread_ids_to_visit: thread_id = thread_ids_to_visit.pop() - state_snapshot = subgraph_state_snapshots[thread_id] + 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 THREAD_ID_SEPARATOR in parent_thread_id: - parent_subgraph_snapshots = ( - subgraph_state_snapshots[parent_thread_id].subgraph_state_snapshots - or {} - ) - parent_subgraph_snapshots[subgraph_node] = state_snapshot - subgraph_state_snapshots[parent_thread_id] = subgraph_state_snapshots[ + if parent_thread_id and (parent_state_snapshot := thread_id_to_state_snapshots.get(parent_thread_id)): + 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(subgraph_state_snapshots=parent_subgraph_snapshots) - state_snapshot = subgraph_state_snapshots.pop(root_thread_id, None) + state_snapshot = thread_id_to_state_snapshots.pop(root_thread_id, None) if state_snapshot is None: raise ValueError(f"Missing snapshot for thread ID '{root_thread_id}'") return state_snapshot - def get_state(self, config: RunnableConfig) -> StateSnapshot: + def get_state(self, config: RunnableConfig, *, include_subgraph_state: bool = False) -> StateSnapshot: """Get the current state of the graph.""" if not self.checkpointer: raise ValueError("No checkpointer set") - subgraph_state_snapshots: dict[str, StateSnapshot] = { + if include_subgraph_state: + checkpoint_tuples = self.checkpointer.list_subgraph_checkpoints(config) + else: + checkpoint_tuples = iter([self.checkpointer.get_tuple(config)]) + + thread_id_to_state_snapshots: dict[str, StateSnapshot] = { checkpoint.config["configurable"][ "thread_id" ]: self._prepare_state_snapshot(checkpoint, config) - for checkpoint in self.checkpointer.list_subgraph_checkpoints(config) + for checkpoint in checkpoint_tuples } thread_id = config["configurable"]["thread_id"] state_snapshot = self._assemble_state_snapshot_hierarchy( - thread_id, subgraph_state_snapshots + thread_id, thread_id_to_state_snapshots ) return state_snapshot - async def aget_state(self, config: RunnableConfig) -> StateSnapshot: + async def aget_state(self, config: RunnableConfig, *, include_subgraph_state: bool = False) -> StateSnapshot: """Get the current state of the graph.""" if not self.checkpointer: raise ValueError("No checkpointer set") - subgraph_state_snapshots: dict[str, StateSnapshot] = { + if include_subgraph_state: + checkpoint_tuples = self.checkpointer.alist_subgraph_checkpoints(config) + else: + async def alist_checkpoints(): + yield await self.checkpointer.aget_tuple(config) + checkpoint_tuples = alist_checkpoints() + + thread_id_to_state_snapshots: dict[str, StateSnapshot] = { checkpoint.config["configurable"][ "thread_id" ]: await self._prepare_state_snapshot_async(checkpoint, config) - async for checkpoint in self.checkpointer.alist_subgraph_checkpoints(config) + async for checkpoint in checkpoint_tuples } thread_id = config["configurable"]["thread_id"] state_snapshot = self._assemble_state_snapshot_hierarchy( - thread_id, subgraph_state_snapshots + thread_id, thread_id_to_state_snapshots ) return state_snapshot From b43ef6440f2cc6a2f5801131d2aa783dc3c6244a Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 23 Jul 2024 20:38:20 -0400 Subject: [PATCH 04/41] tests --- .../langgraph/checkpoint/aiosqlite.py | 1 + libs/langgraph/langgraph/checkpoint/base.py | 2 + libs/langgraph/langgraph/checkpoint/memory.py | 2 + libs/langgraph/langgraph/checkpoint/sqlite.py | 1 + libs/langgraph/langgraph/pregel/__init__.py | 24 +- libs/langgraph/tests/test_pregel.py | 354 +++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 365 +++++++++++++++++- 7 files changed, 742 insertions(+), 7 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/aiosqlite.py b/libs/langgraph/langgraph/checkpoint/aiosqlite.py index 401982fdf..e0e982537 100644 --- a/libs/langgraph/langgraph/checkpoint/aiosqlite.py +++ b/libs/langgraph/langgraph/checkpoint/aiosqlite.py @@ -350,6 +350,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): async def alist_subgraph_checkpoints( self, config: RunnableConfig ) -> AsyncIterator[CheckpointTuple]: + # TODO: docstring async with self.conn.cursor() as cur: if config["configurable"].get("thread_ts"): await cur.execute( diff --git a/libs/langgraph/langgraph/checkpoint/base.py b/libs/langgraph/langgraph/checkpoint/base.py index f2b0b0f63..d2b0df3da 100644 --- a/libs/langgraph/langgraph/checkpoint/base.py +++ b/libs/langgraph/langgraph/checkpoint/base.py @@ -234,6 +234,7 @@ class BaseCheckpointSaver(ABC): def list_subgraph_checkpoints( self, config: RunnableConfig ) -> Iterator[CheckpointTuple]: + # TODO: docstring raise NotImplementedError def put( @@ -331,6 +332,7 @@ class BaseCheckpointSaver(ABC): async def alist_subgraph_checkpoints( self, config: RunnableConfig ) -> AsyncIterator[CheckpointTuple]: + # TODO: docstring raise NotImplementedError async def aput( diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index 348aacdf0..8160de238 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -174,6 +174,7 @@ class MemorySaver(BaseCheckpointSaver): def list_subgraph_checkpoints( self, config: RunnableConfig ) -> Iterator[CheckpointTuple]: + # TODO: docstring thread_id_prefix = config["configurable"]["thread_id"] matching_thread_ids = [ key for key in self.storage.keys() if key.startswith(thread_id_prefix) @@ -314,6 +315,7 @@ class MemorySaver(BaseCheckpointSaver): async def alist_subgraph_checkpoints( self, config: RunnableConfig ) -> AsyncIterator[CheckpointTuple]: + # TODO: docstring loop = asyncio.get_running_loop() iter = await loop.run_in_executor(None, self.list_subgraph_checkpoints, config) while True: diff --git a/libs/langgraph/langgraph/checkpoint/sqlite.py b/libs/langgraph/langgraph/checkpoint/sqlite.py index f6bf4b857..6578764ca 100644 --- a/libs/langgraph/langgraph/checkpoint/sqlite.py +++ b/libs/langgraph/langgraph/checkpoint/sqlite.py @@ -360,6 +360,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): def list_subgraph_checkpoints( self, config: RunnableConfig ) -> Iterator[CheckpointTuple]: + # TODO: docstring with self.cursor(transaction=False) as cur: if config["configurable"].get("thread_ts"): cur.execute( diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 511a5fd9c..f1ab7d5b6 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -421,21 +421,29 @@ class Pregel( 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)): + if parent_thread_id and ( + parent_state_snapshot := thread_id_to_state_snapshots.get( + parent_thread_id + ) + ): parent_subgraph_snapshots = { **(parent_state_snapshot.subgraph_state_snapshots or {}), - subgraph_node: state_snapshot + subgraph_node: state_snapshot, } - thread_id_to_state_snapshots[parent_thread_id] = thread_id_to_state_snapshots[ + thread_id_to_state_snapshots[ parent_thread_id - ]._replace(subgraph_state_snapshots=parent_subgraph_snapshots) + ] = thread_id_to_state_snapshots[parent_thread_id]._replace( + subgraph_state_snapshots=parent_subgraph_snapshots + ) state_snapshot = thread_id_to_state_snapshots.pop(root_thread_id, None) if state_snapshot is None: raise ValueError(f"Missing snapshot for thread ID '{root_thread_id}'") return state_snapshot - def get_state(self, config: RunnableConfig, *, include_subgraph_state: bool = False) -> StateSnapshot: + def get_state( + self, config: RunnableConfig, *, include_subgraph_state: bool = False + ) -> StateSnapshot: """Get the current state of the graph.""" if not self.checkpointer: raise ValueError("No checkpointer set") @@ -458,7 +466,9 @@ 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, *, include_subgraph_state: bool = False + ) -> StateSnapshot: """Get the current state of the graph.""" if not self.checkpointer: raise ValueError("No checkpointer set") @@ -466,8 +476,10 @@ class Pregel( if include_subgraph_state: checkpoint_tuples = self.checkpointer.alist_subgraph_checkpoints(config) else: + async def alist_checkpoints(): yield await self.checkpointer.aget_tuple(config) + checkpoint_tuples = alist_checkpoints() thread_id_to_state_snapshots: dict[str, StateSnapshot] = { diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 907fdb0f9..fea1e29b8 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8674,6 +8674,360 @@ def test_doubly_nested_graph_interrupts(checkpointer: BaseCheckpointSaver) -> No checkpointer.__exit__(None, None, None) +@pytest.mark.parametrize( + "checkpointer_fct", + [ + lambda: MemorySaverAssertImmutable(put_sleep=0.2), + lambda: SqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +def test_nested_graph_state( + checkpointer_fct: Callable[[], BaseCheckpointSaver], +) -> None: + try: + checkpointer = checkpointer_fct() + + 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", "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=None, + ) + assert app.get_state(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, + ) + }, + ) + 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", "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={ + "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, + ) + }, + ) + finally: + if hasattr(checkpointer, "__exit__"): + checkpointer.__exit__(None, None, None) + + +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + SqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +def test_doubly_nested_graph_state(checkpointer: BaseCheckpointSaver) -> None: + try: + + 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"}, + next=("child",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + subgraph_state_snapshots=None, + ) + assert app.get_state(config, include_subgraph_state=True) == StateSnapshot( + values={"my_key": "hi my value"}, + next=("child",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + subgraph_state_snapshots={ + "child": StateSnapshot( + values={"my_key": "hi my value"}, + next=(), + config={ + "configurable": {"thread_id": "1__child", "thread_ts": AnyStr()} + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1__child", "thread_ts": AnyStr()} + }, + subgraph_state_snapshots={ + "child_1": StateSnapshot( + values={"my_key": "hi my value here"}, + next=(), + config={ + "configurable": { + "thread_id": "1__child__child_1", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "grandchild_1": {"my_key": "hi my value here"} + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1__child__child_1", + "thread_ts": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ) + }, + ) + app.invoke(None, config, debug=True) + 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", "thread_ts": 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", "thread_ts": AnyStr()}}, + subgraph_state_snapshots={ + "child": StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + config={ + "configurable": {"thread_id": "1__child", "thread_ts": 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__child", "thread_ts": AnyStr()} + }, + subgraph_state_snapshots={ + "child_1": StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + config={ + "configurable": { + "thread_id": "1__child__child_1", + "thread_ts": 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__child__child_1", + "thread_ts": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ) + }, + ) + + finally: + if hasattr(checkpointer, "__exit__"): + checkpointer.__exit__(None, None, 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 7d88f286e..2e40dc0b1 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6079,7 +6079,7 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: assert times_called == 1 -# @pytest.mark.repeat(10) +@pytest.mark.repeat(10) @pytest.mark.parametrize( "checkpointer_fct", [ @@ -7209,6 +7209,369 @@ async def test_doubly_nested_graph_interrupts( await checkpointer.__aexit__(None, None, None) +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + AsyncSqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +async def test_nested_graph_state(checkpointer: BaseCheckpointSaver) -> None: + try: + + class InnerState(TypedDict): + my_key: str + my_other_key: str + + async def inner_1(state: InnerState): + return { + "my_key": state["my_key"] + " here", + "my_other_key": state["my_key"], + } + + async def inner_2(state: InnerState): + return { + "my_key": state["my_key"] + " and there", + "my_other_key": state["my_key"], + } + + inner = StateGraph(InnerState) + inner.add_node("inner_1", inner_1) + inner.add_node("inner_2", inner_2) + inner.add_edge("inner_1", "inner_2") + inner.set_entry_point("inner_1") + inner.set_finish_point("inner_2") + + class State(TypedDict): + my_key: str + + async def outer_1(state: State): + return {"my_key": "hi " + state["my_key"]} + + async def outer_2(state: State): + return {"my_key": state["my_key"] + " and back again"} + + graph = StateGraph(State) + graph.add_node("outer_1", outer_1) + graph.add_node("inner", inner.compile(interrupt_before=["inner_2"])) + graph.add_node("outer_2", outer_2) + graph.set_entry_point("outer_1") + graph.add_edge("outer_1", "inner") + graph.add_edge("inner", "outer_2") + graph.set_finish_point("outer_2") + + app = graph.compile(checkpointer=checkpointer) + + 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", "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=None, + ) + assert await app.aget_state( + 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, + ) + }, + ) + 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", "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={ + "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, + ) + }, + ) + finally: + if hasattr(checkpointer, "__aexit__"): + await checkpointer.__aexit__(None, None, None) + + +@pytest.mark.parametrize( + "checkpointer", + [ + MemorySaverAssertImmutable(), + AsyncSqliteSaver.from_conn_string(":memory:"), + ], + ids=[ + "memory", + "sqlite", + ], +) +async def test_doubly_nested_graph_state( + checkpointer: BaseCheckpointSaver, +) -> None: + try: + + 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"}, + next=("child",), + config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={"configurable": {"thread_id": "1", "thread_ts": 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={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + metadata={ + "source": "loop", + "writes": {"parent_1": {"my_key": "hi my value"}}, + "step": 1, + }, + created_at=AnyStr(), + parent_config={"configurable": {"thread_id": "1", "thread_ts": AnyStr()}}, + subgraph_state_snapshots={ + "child": StateSnapshot( + values={"my_key": "hi my value"}, + next=(), + config={ + "configurable": {"thread_id": "1__child", "thread_ts": AnyStr()} + }, + metadata={"source": "loop", "writes": None, "step": 0}, + created_at=AnyStr(), + parent_config={ + "configurable": {"thread_id": "1__child", "thread_ts": AnyStr()} + }, + subgraph_state_snapshots={ + "child_1": StateSnapshot( + values={"my_key": "hi my value here"}, + next=(), + config={ + "configurable": { + "thread_id": "1__child__child_1", + "thread_ts": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": { + "grandchild_1": {"my_key": "hi my value here"} + }, + "step": 1, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1__child__child_1", + "thread_ts": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ) + }, + ) + await app.ainvoke(None, config, debug=True) + 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", "thread_ts": 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", "thread_ts": AnyStr()}}, + subgraph_state_snapshots={ + "child": StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + config={ + "configurable": {"thread_id": "1__child", "thread_ts": 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__child", "thread_ts": AnyStr()} + }, + subgraph_state_snapshots={ + "child_1": StateSnapshot( + values={"my_key": "hi my value here and there"}, + next=(), + config={ + "configurable": { + "thread_id": "1__child__child_1", + "thread_ts": 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__child__child_1", + "thread_ts": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ) + }, + ) + }, + ) + + finally: + if hasattr(checkpointer, "__aexit__"): + await checkpointer.__aexit__(None, None, 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. From ae696d4f3039cb8286bf01325e6970da5206fc63 Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 23 Jul 2024 21:40:04 -0400 Subject: [PATCH 05/41] 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) From a7d48465dafdc328e79f591edb956ad5e1a11bfd Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 24 Jul 2024 12:44:35 -0400 Subject: [PATCH 06/41] use .list for looking up prefix-matched checkpoints --- .../langgraph/checkpoint/aiosqlite.py | 47 +--- libs/langgraph/langgraph/checkpoint/base.py | 13 +- libs/langgraph/langgraph/checkpoint/memory.py | 74 ++---- libs/langgraph/langgraph/checkpoint/sqlite.py | 52 +---- libs/langgraph/langgraph/pregel/__init__.py | 131 +++++++---- libs/langgraph/tests/test_pregel.py | 2 +- libs/langgraph/tests/test_pregel_async.py | 219 ++++++++++++++++++ 7 files changed, 333 insertions(+), 205 deletions(-) diff --git a/libs/langgraph/langgraph/checkpoint/aiosqlite.py b/libs/langgraph/langgraph/checkpoint/aiosqlite.py index e0e982537..5ac9a9197 100644 --- a/libs/langgraph/langgraph/checkpoint/aiosqlite.py +++ b/libs/langgraph/langgraph/checkpoint/aiosqlite.py @@ -306,6 +306,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, + as_prefix: bool = False, ) -> AsyncIterator[CheckpointTuple]: """List checkpoints from the database asynchronously. @@ -322,7 +323,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): AsyncIterator[CheckpointTuple]: An asynchronous iterator of matching checkpoint tuples. """ await self.setup() - where, param_values = search_where(config, filter, before) + where, param_values = search_where(config, filter, before, as_prefix=as_prefix) query = f"""SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints {where} @@ -347,50 +348,6 @@ class AsyncSqliteSaver(BaseCheckpointSaver, AbstractAsyncContextManager): ), ) - async def alist_subgraph_checkpoints( - self, config: RunnableConfig - ) -> AsyncIterator[CheckpointTuple]: - # TODO: docstring - async with self.conn.cursor() as cur: - if config["configurable"].get("thread_ts"): - await cur.execute( - "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id LIKE ? || '%' AND thread_ts = ?", - ( - str(config["configurable"]["thread_id"]), - str(config["configurable"]["thread_ts"]), - ), - ) - else: - await cur.execute( - """SELECT checkpoints.thread_id, checkpoints.thread_ts, checkpoints.parent_ts, checkpoints.checkpoint, checkpoints.metadata - FROM checkpoints - INNER JOIN ( - SELECT thread_id, MAX(thread_ts) as thread_ts - FROM checkpoints - WHERE thread_id LIKE ? || '%' - GROUP BY thread_id - ) latest_checkpoints - ON checkpoints.thread_id = latest_checkpoints.thread_id AND checkpoints.thread_ts = latest_checkpoints.thread_ts - ORDER BY checkpoints.thread_id, checkpoints.thread_ts DESC""", - (str(config["configurable"]["thread_id"]),), - ) - async for thread_id, thread_ts, parent_ts, value, metadata in cur: - yield CheckpointTuple( - {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}}, - self.serde.loads(value), - self.serde.loads(metadata) if metadata is not None else {}, - ( - { - "configurable": { - "thread_id": thread_id, - "thread_ts": parent_ts, - } - } - if parent_ts - else None - ), - ) - async def aput( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/checkpoint/base.py b/libs/langgraph/langgraph/checkpoint/base.py index d2b0df3da..cca902050 100644 --- a/libs/langgraph/langgraph/checkpoint/base.py +++ b/libs/langgraph/langgraph/checkpoint/base.py @@ -214,6 +214,7 @@ class BaseCheckpointSaver(ABC): filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, + as_prefix: bool = False, ) -> Iterator[CheckpointTuple]: """List checkpoints that match the given criteria. @@ -231,12 +232,6 @@ class BaseCheckpointSaver(ABC): """ raise NotImplementedError - def list_subgraph_checkpoints( - self, config: RunnableConfig - ) -> Iterator[CheckpointTuple]: - # TODO: docstring - raise NotImplementedError - def put( self, config: RunnableConfig, @@ -329,12 +324,6 @@ class BaseCheckpointSaver(ABC): raise NotImplementedError yield - async def alist_subgraph_checkpoints( - self, config: RunnableConfig - ) -> AsyncIterator[CheckpointTuple]: - # TODO: docstring - raise NotImplementedError - async def aput( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/checkpoint/memory.py b/libs/langgraph/langgraph/checkpoint/memory.py index 8160de238..1aa89f881 100644 --- a/libs/langgraph/langgraph/checkpoint/memory.py +++ b/libs/langgraph/langgraph/checkpoint/memory.py @@ -119,6 +119,7 @@ class MemorySaver(BaseCheckpointSaver): filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, + as_prefix: bool = False, ) -> Iterator[CheckpointTuple]: """List checkpoints from the in-memory storage. @@ -134,7 +135,19 @@ class MemorySaver(BaseCheckpointSaver): Yields: Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. """ - thread_ids = (config["configurable"]["thread_id"],) if config else self.storage + if config: + config_thread_id = config["configurable"]["thread_id"] + if as_prefix: + thread_ids = ( + thread_id + for thread_id in self.storage + if thread_id.startswith(config_thread_id) + ) + else: + thread_ids = (config_thread_id,) + else: + thread_ids = self.storage.keys() + for thread_id in thread_ids: for ts, (checkpoint, metadata_b, parent_ts) in sorted( self.storage[thread_id].items(), key=lambda x: x[0], reverse=True @@ -171,40 +184,6 @@ class MemorySaver(BaseCheckpointSaver): else None, ) - def list_subgraph_checkpoints( - self, config: RunnableConfig - ) -> Iterator[CheckpointTuple]: - # TODO: docstring - thread_id_prefix = config["configurable"]["thread_id"] - matching_thread_ids = [ - key for key in self.storage.keys() if key.startswith(thread_id_prefix) - ] - for thread_id in matching_thread_ids: - ts = config["configurable"].get("thread_ts") - if not ts: - if checkpoints := self.storage[thread_id]: - ts = max(checkpoints.keys()) - - if saved := self.storage[thread_id].get(ts): - checkpoint, metadata, parent_ts = saved - writes = self.writes[(thread_id, ts)] - yield CheckpointTuple( - config={"configurable": {"thread_id": thread_id, "thread_ts": ts}}, - checkpoint=self.serde.loads(checkpoint), - metadata=self.serde.loads(metadata), - pending_writes=[ - (id, c, self.serde.loads(v)) for id, c, v in writes - ], - parent_config={ - "configurable": { - "thread_id": thread_id, - "thread_ts": parent_ts, - } - } - if parent_ts - else None, - ) - def put( self, config: RunnableConfig, @@ -288,6 +267,7 @@ class MemorySaver(BaseCheckpointSaver): filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, + as_prefix: bool = False, ) -> AsyncIterator[CheckpointTuple]: """Asynchronous version of list. @@ -302,7 +282,15 @@ class MemorySaver(BaseCheckpointSaver): """ 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, + as_prefix=as_prefix, + ), + config, ) while True: # handling StopIteration exception inside coroutine won't work @@ -312,20 +300,6 @@ class MemorySaver(BaseCheckpointSaver): else: break - async def alist_subgraph_checkpoints( - self, config: RunnableConfig - ) -> AsyncIterator[CheckpointTuple]: - # TODO: docstring - loop = asyncio.get_running_loop() - iter = await loop.run_in_executor(None, self.list_subgraph_checkpoints, config) - while True: - # handling StopIteration exception inside coroutine won't work - # as expected, so using next() with default value to break the loop - if item := await loop.run_in_executor(None, next, iter, None): - yield item - else: - break - async def aput( self, config: RunnableConfig, diff --git a/libs/langgraph/langgraph/checkpoint/sqlite.py b/libs/langgraph/langgraph/checkpoint/sqlite.py index 6578764ca..c1c0b2217 100644 --- a/libs/langgraph/langgraph/checkpoint/sqlite.py +++ b/libs/langgraph/langgraph/checkpoint/sqlite.py @@ -301,6 +301,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): filter: Optional[Dict[str, Any]] = None, before: Optional[RunnableConfig] = None, limit: Optional[int] = None, + as_prefix: bool = False, ) -> Iterator[CheckpointTuple]: """List checkpoints from the database. @@ -331,7 +332,7 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): >>> print(checkpoints) [CheckpointTuple(...), ...] """ - where, param_values = search_where(config, filter, before) + where, param_values = search_where(config, filter, before, as_prefix=as_prefix) query = f"""SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints {where} @@ -357,50 +358,6 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): ), ) - def list_subgraph_checkpoints( - self, config: RunnableConfig - ) -> Iterator[CheckpointTuple]: - # TODO: docstring - with self.cursor(transaction=False) as cur: - if config["configurable"].get("thread_ts"): - cur.execute( - "SELECT thread_id, thread_ts, parent_ts, checkpoint, metadata FROM checkpoints WHERE thread_id LIKE ? || '%' AND thread_ts = ?", - ( - str(config["configurable"]["thread_id"]), - str(config["configurable"]["thread_ts"]), - ), - ) - else: - cur.execute( - """SELECT checkpoints.thread_id, checkpoints.thread_ts, checkpoints.parent_ts, checkpoints.checkpoint, checkpoints.metadata - FROM checkpoints - INNER JOIN ( - SELECT thread_id, MAX(thread_ts) as thread_ts - FROM checkpoints - WHERE thread_id LIKE ? || '%' - GROUP BY thread_id - ) latest_checkpoints - ON checkpoints.thread_id = latest_checkpoints.thread_id AND checkpoints.thread_ts = latest_checkpoints.thread_ts - ORDER BY checkpoints.thread_id, checkpoints.thread_ts DESC""", - (str(config["configurable"]["thread_id"]),), - ) - for thread_id, thread_ts, parent_ts, value, metadata in cur: - yield CheckpointTuple( - {"configurable": {"thread_id": thread_id, "thread_ts": thread_ts}}, - self.serde.loads(value), - self.serde.loads(metadata) if metadata is not None else {}, - ( - { - "configurable": { - "thread_id": thread_id, - "thread_ts": parent_ts, - } - } - if parent_ts - else None - ), - ) - def put( self, config: RunnableConfig, @@ -592,6 +549,7 @@ def search_where( config: Optional[RunnableConfig], filter: Optional[Dict[str, Any]], before: Optional[RunnableConfig] = None, + as_prefix: bool = False, ) -> Tuple[str, Sequence[Any]]: """Return WHERE clause predicates for (a)search() given metadata filter and `before` config. @@ -606,7 +564,9 @@ def search_where( # construct predicate for config filter if config is not None: - wheres.append("thread_id = ?") + thread_filter = "thread_id LIKE ? || '%'" if as_prefix else "thread_id = ?" + wheres.append(thread_filter) + param_values.append(config["configurable"]["thread_id"]) # construct predicate for metadata filter diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 6287c713b..c0b3a511d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -438,7 +438,7 @@ class Pregel( state_snapshot = thread_id_to_state_snapshots.pop(root_thread_id, None) if state_snapshot is None: - raise ValueError(f"Missing snapshot for thread ID '{root_thread_id}'") + raise ValueError(f"Missing checkpoint for thread ID '{root_thread_id}'") return state_snapshot def get_state( @@ -449,18 +449,34 @@ class Pregel( raise ValueError("No checkpointer set") if include_subgraph_state: - checkpoint_tuples = self.checkpointer.list_subgraph_checkpoints(config) + checkpoint_tuples = self.checkpointer.list(config, as_prefix=True) else: checkpoint_tuples = iter([self.checkpointer.get_tuple(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"] + thread_ts = config["configurable"].get("thread_ts") + thread_id_to_thread_ts: dict[str, str] = {} + thread_id_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: + 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 = 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 + + 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}'" + + raise ValueError(error_msg) + state_snapshot = self._assemble_state_snapshot_hierarchy( thread_id, thread_id_to_state_snapshots ) @@ -474,7 +490,7 @@ class Pregel( raise ValueError("No checkpointer set") if include_subgraph_state: - checkpoint_tuples = self.checkpointer.alist_subgraph_checkpoints(config) + checkpoint_tuples = self.checkpointer.alist(config, as_prefix=True) else: async def alist_checkpoints(): @@ -482,14 +498,32 @@ class Pregel( checkpoint_tuples = alist_checkpoints() - thread_id_to_state_snapshots: dict[str, StateSnapshot] = { - checkpoint.config["configurable"][ - "thread_id" - ]: await self._prepare_state_snapshot_async(checkpoint, config) - async for checkpoint in checkpoint_tuples - } - 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] = {} + 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: + 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 + + 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}'" + + raise ValueError(error_msg) + state_snapshot = self._assemble_state_snapshot_hierarchy( thread_id, thread_id_to_state_snapshots ) @@ -515,20 +549,8 @@ class Pregel( for config, checkpoint, metadata, parent_config, _ in self.checkpointer.list( config, before=before, limit=limit, filter=filter ): - # 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 - ) + state_snapshot = self.get_state(config, include_subgraph_state=True) yield state_snapshot else: with ChannelsManager( @@ -561,6 +583,7 @@ 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: @@ -577,28 +600,34 @@ class Pregel( parent_config, _, ) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter): - async with AsyncChannelsManager( - self.channels, checkpoint, config - ) as channels, AsyncManagedValuesManager( - 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, + if include_subgraph_state: + state_snapshot = await self.aget_state( + config, include_subgraph_state=True ) + yield state_snapshot + else: + async with AsyncChannelsManager( + self.channels, checkpoint, config + ) as channels, AsyncManagedValuesManager( + 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, + ) def update_state( self, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0a86aee21..34e4028ff 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8896,7 +8896,7 @@ def test_nested_graph_state( ) }, ) - assert list(app.get_state_history(config, include_subgraph_state=False)) == [ + 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=(), diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 2e40dc0b1..a9653b546 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7322,6 +7322,78 @@ async def test_nested_graph_state(checkpointer: BaseCheckpointSaver) -> None: ) }, ) + assert [ + c async for c in app.aget_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, + ), + ] await app.ainvoke(None, config, debug=True) # test state w/ nested subgraph state (after resuming from interrupt) assert await app.aget_state( @@ -7364,6 +7436,153 @@ async def test_nested_graph_state(checkpointer: BaseCheckpointSaver) -> None: ) }, ) + assert [ + c async for c 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", + "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, "__aexit__"): await checkpointer.__aexit__(None, None, None) From 322cfc46d3acbbd8cc3efda0e40bb70ef524af93 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 12 Aug 2024 10:36:25 -0400 Subject: [PATCH 07/41] cleanup --- .../langgraph/checkpoint/base/__init__.py | 2 + .../langgraph/checkpoint/memory/__init__.py | 109 +++-- libs/langgraph/langgraph/pregel/__init__.py | 150 +++--- libs/langgraph/tests/test_pregel.py | 450 +++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 452 ++++++++++++++++++ 5 files changed, 1057 insertions(+), 106 deletions(-) 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. From 9948125745de3acafa1cbdef9e8f90e1fe484126 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 12 Aug 2024 14:00:38 -0400 Subject: [PATCH 08/41] add checkpointer=INHERIT_CHECKPOINTER --- libs/langgraph/langgraph/graph/graph.py | 5 +- libs/langgraph/langgraph/graph/state.py | 4 +- libs/langgraph/langgraph/pregel/__init__.py | 23 +- libs/langgraph/tests/test_pregel.py | 280 +++++++++++++++++++- libs/langgraph/tests/test_pregel_async.py | 278 ++++++++++++++++++- 5 files changed, 568 insertions(+), 22 deletions(-) diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index dc927446d..c17373820 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -24,7 +24,6 @@ from langchain_core.runnables.graph import Graph as DrawableGraph from langchain_core.runnables.graph import Node as DrawableNode from langgraph.channels.ephemeral_value import EphemeralValue -from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.constants import ( CHECKPOINT_NAMESPACE_SEPARATOR, END, @@ -33,7 +32,7 @@ from langgraph.constants import ( Send, ) from langgraph.errors import InvalidUpdateError -from langgraph.pregel import Channel, Pregel +from langgraph.pregel import Channel, CheckpointerType, Pregel from langgraph.pregel.read import PregelNode from langgraph.pregel.types import All from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry @@ -369,7 +368,7 @@ class Graph: def compile( self, - checkpointer: Optional[BaseCheckpointSaver] = None, + checkpointer: Optional[CheckpointerType] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, debug: bool = False, diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 5810aaf29..e9971a5f1 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -29,7 +29,6 @@ from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitFo 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.errors import InvalidUpdateError from langgraph.graph.graph import ( @@ -41,6 +40,7 @@ from langgraph.graph.graph import ( Send, ) from langgraph.managed.base import ManagedValue, is_managed_value +from langgraph.pregel import CheckpointerType from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.types import All, RetryPolicy from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry @@ -373,7 +373,7 @@ class StateGraph(Graph): def compile( self, - checkpointer: Optional[BaseCheckpointSaver] = None, + checkpointer: Optional[CheckpointerType] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, debug: bool = False, diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 34c12604c..dd105e0ba 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -13,6 +13,7 @@ from typing import ( Callable, Dict, Iterator, + Literal, Mapping, Optional, Sequence, @@ -115,6 +116,10 @@ WriteValue = Union[ ] +INHERIT_CHECKPOINTER = "inherit_checkpointer" +CheckpointerType = Union[BaseCheckpointSaver, Literal["inherit_checkpointer"]] + + class Channel: @overload @classmethod @@ -218,7 +223,7 @@ class Pregel( debug: bool = Field(default_factory=get_debug) """Whether to print debug information during execution. Defaults to False.""" - checkpointer: Optional[BaseCheckpointSaver] = None + checkpointer: Optional[CheckpointerType] = None """Checkpointer used to save and load graph state. Defaults to None.""" retry_policy: Optional[RetryPolicy] = None @@ -275,6 +280,7 @@ class Pregel( + ( self.checkpointer.config_specs if self.checkpointer is not None + and self.checkpointer != INHERIT_CHECKPOINTER else [] ) + ( @@ -950,10 +956,23 @@ class Pregel( if config and config.get("configurable", {}).get(CONFIG_KEY_READ) is not None: # if being called as a node in another graph, always use values mode stream_mode = ["values"] + + if self.checkpointer is None: + raise ValueError( + "Missing checkpointer for subgraph. " + "Please compile the subgraph graph with checkpointer=INHERIT_CHECKPOINTER (from langgraph.pregel import INHERIT_CHECKPOINTER)." + ) + + if self.checkpointer != INHERIT_CHECKPOINTER: + raise ValueError( + "Custom checkpointers for subgraphs are not allowed. " + "Please compile the subgraph graph with checkpointer=INHERIT_CHECKPOINTER (from langgraph.pregel import INHERIT_CHECKPOINTER)." + ) + if ( config is not None and config.get("configurable", {}).get(CONFIG_KEY_CHECKPOINTER) - and (interrupt_after or interrupt_before) + and self.checkpointer == INHERIT_CHECKPOINTER ): checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][ CONFIG_KEY_CHECKPOINTER diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 1530fcda4..0f4fbb5b2 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -60,7 +60,13 @@ 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 ( + INHERIT_CHECKPOINTER, + Channel, + GraphRecursionError, + Pregel, + StateSnapshot, +) from langgraph.pregel.retry import RetryPolicy from tests.any_str import AnyStr from tests.memory_assert import ( @@ -7737,7 +7743,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"], checkpointer=INHERIT_CHECKPOINTER), + ) graph.add_node("outer_2", outer_2) graph.set_entry_point("outer_1") graph.add_edge("outer_1", "inner") @@ -8913,7 +8922,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"], checkpointer=INHERIT_CHECKPOINTER), + ) graph.add_node("outer_1", outer_1) graph.add_node("outer_2", outer_2) @@ -8997,7 +9009,6 @@ def test_nested_graph_interrupts_parallel( ] -@pytest.mark.skip @pytest.mark.parametrize( "checkpointer_name", ["memory", "sqlite", "postgres", "postgres_pipe"], @@ -9005,7 +9016,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 @@ -9032,7 +9043,12 @@ 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"], checkpointer=INHERIT_CHECKPOINTER + ), + ) child.set_entry_point("child_1") child.set_finish_point("child_1") @@ -9044,7 +9060,7 @@ def test_doubly_nested_graph_interrupts( graph = StateGraph(State) graph.add_node("parent_1", parent_1) - graph.add_node("child", child.compile()) + graph.add_node("child", child.compile(checkpointer=INHERIT_CHECKPOINTER)) graph.add_node("parent_2", parent_2) graph.set_entry_point("parent_1") graph.add_edge("parent_1", "child") @@ -9136,7 +9152,10 @@ def test_nested_graph_state( 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"], checkpointer=INHERIT_CHECKPOINTER), + ) graph.add_node("outer_2", outer_2) graph.set_entry_point("outer_1") graph.add_edge("outer_1", "inner") @@ -9543,6 +9562,251 @@ def test_nested_graph_state( ] +@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"], checkpointer=INHERIT_CHECKPOINTER + ), + ) + 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(checkpointer=INHERIT_CHECKPOINTER)) + 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"}, + 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={ + "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"}, + next=(), + 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"}, + next=(), + 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, 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": { + "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(), + } + }, + subgraph_state_snapshots={ + "child": StateSnapshot( + values={"my_key": "hi my value here and there"}, + 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"}, + 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, + ) + }, + ) + }, + ) + + 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 064c48448..ee5efaa37 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -56,7 +56,13 @@ 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 ( + INHERIT_CHECKPOINTER, + Channel, + GraphRecursionError, + Pregel, + StateSnapshot, +) from langgraph.pregel.retry import RetryPolicy from tests.any_str import AnyStr from tests.memory_assert import ( @@ -6232,7 +6238,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"], checkpointer=INHERIT_CHECKPOINTER), + ) graph.add_node("outer_2", outer_2) graph.set_entry_point("outer_1") graph.add_edge("outer_1", "inner") @@ -7415,7 +7424,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"], checkpointer=INHERIT_CHECKPOINTER), + ) graph.add_node("outer_1", outer_1) graph.add_node("outer_2", outer_2) @@ -7501,7 +7513,6 @@ async def test_nested_graph_interrupts_parallel( ] -@pytest.mark.skip @pytest.mark.parametrize( "checkpointer_name", ["memory", "sqlite_aio", "postgres_aio", "postgres_aio_pipe"], @@ -7536,7 +7547,12 @@ 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"], checkpointer=INHERIT_CHECKPOINTER + ), + ) child.set_entry_point("child_1") child.set_finish_point("child_1") @@ -7548,7 +7564,7 @@ async def test_doubly_nested_graph_interrupts( graph = StateGraph(State) graph.add_node("parent_1", parent_1) - graph.add_node("child", child.compile()) + graph.add_node("child", child.compile(checkpointer=INHERIT_CHECKPOINTER)) graph.add_node("parent_2", parent_2) graph.set_entry_point("parent_1") graph.add_edge("parent_1", "child") @@ -7643,7 +7659,10 @@ async def test_nested_graph_state( 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"], checkpointer=INHERIT_CHECKPOINTER), + ) graph.add_node("outer_2", outer_2) graph.set_entry_point("outer_1") graph.add_edge("outer_1", "inner") @@ -8052,6 +8071,251 @@ async def test_nested_graph_state( ] +@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 + + 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"], checkpointer=INHERIT_CHECKPOINTER + ), + ) + 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(checkpointer=INHERIT_CHECKPOINTER)) + 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"}, + 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={ + "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"}, + next=(), + 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"}, + next=(), + 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, 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": { + "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(), + } + }, + subgraph_state_snapshots={ + "child": StateSnapshot( + values={"my_key": "hi my value here and there"}, + 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"}, + 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, + ) + }, + ) + }, + ) + + 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. From fb05bdc2bfda23759cc894b0332d933c775423a9 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 12 Aug 2024 14:45:36 -0400 Subject: [PATCH 09/41] return all checkpoints from .list --- .../langgraph/checkpoint/postgres/base.py | 3 - .../langgraph/checkpoint/sqlite/utils.py | 3 - .../langgraph/checkpoint/base/__init__.py | 2 - .../langgraph/checkpoint/memory/__init__.py | 14 +- libs/langgraph/langgraph/pregel/__init__.py | 201 +++++++----------- libs/langgraph/tests/test_pregel_async.py | 22 +- 6 files changed, 90 insertions(+), 155 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index e9ca12305..038bb6677 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -253,9 +253,6 @@ 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) # construct predicate for metadata filter if filter: diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py index 6e1baf5ae..56034ea34 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py @@ -70,9 +70,6 @@ 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) # construct predicate for metadata filter if filter: diff --git a/libs/checkpoint/langgraph/checkpoint/base/__init__.py b/libs/checkpoint/langgraph/checkpoint/base/__init__.py index 4f0f08a84..86c8b0eec 100644 --- a/libs/checkpoint/langgraph/checkpoint/base/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/base/__init__.py @@ -259,7 +259,6 @@ 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. @@ -351,7 +350,6 @@ 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 5e60ebaf0..3bb7f8803 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -157,7 +157,6 @@ 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. @@ -178,16 +177,7 @@ class MemorySaver( config["configurable"].get("checkpoint_ns", "") if config else "" ) for thread_id in thread_ids: - 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_ns in self.storage[thread_id].keys(): for checkpoint_id, ( checkpoint, metadata_b, @@ -330,7 +320,6 @@ 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. @@ -351,7 +340,6 @@ 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 dd105e0ba..9c591979f 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -360,70 +360,64 @@ class Pregel( if is_managed_value(v) } - def _prepare_state_snapshot( - self, saved: CheckpointTuple, config: RunnableConfig - ) -> StateSnapshot: - checkpoint = saved.checkpoint if saved else empty_checkpoint() - config = saved.config if saved else config + def _prepare_state_snapshot(self, saved: CheckpointTuple) -> StateSnapshot: with ChannelsManager( { k: LastValue(None) if isinstance(c, Context) else c for k, c in self.channels.items() }, - checkpoint, - config, + saved.checkpoint, + saved.config, ) as channels, ManagedValuesManager( - self.managed_values_dict, ensure_config(config) + self.managed_values_dict, ensure_config(saved.config) ) as managed: next_tasks = prepare_next_tasks( - checkpoint, + saved.checkpoint, self.nodes, channels, managed, - config, + saved.config, -1, for_execution=False, ) return StateSnapshot( values=read_channels(channels, self.stream_channels_asis), next=tuple(t.name for t in next_tasks), - config=saved.config if saved else config, - metadata=saved.metadata if saved else None, - created_at=saved.checkpoint["ts"] if saved else None, - parent_config=saved.parent_config if saved else None, + config=saved.config, + metadata=saved.metadata, + created_at=saved.checkpoint["ts"], + parent_config=saved.parent_config, ) async def _prepare_state_snapshot_async( - self, saved: CheckpointTuple, config: RunnableConfig + self, saved: CheckpointTuple ) -> StateSnapshot: - checkpoint = saved.checkpoint if saved else empty_checkpoint() - config = saved.config if saved else config async with AsyncChannelsManager( { k: LastValue(None) if isinstance(c, Context) else c for k, c in self.channels.items() }, - checkpoint, - config, + saved.checkpoint, + saved.config, ) as channels, AsyncManagedValuesManager( - self.managed_values_dict, ensure_config(config) + self.managed_values_dict, ensure_config(saved.config) ) as managed: next_tasks = prepare_next_tasks( - checkpoint, + saved.checkpoint, self.nodes, channels, managed, - config, + saved.config, -1, for_execution=False, ) return StateSnapshot( values=read_channels(channels, self.stream_channels_asis), next=tuple(t.name for t in next_tasks), - config=saved.config if saved else config, - metadata=saved.metadata if saved else None, - created_at=saved.checkpoint["ts"] if saved else None, - parent_config=saved.parent_config if saved else None, + config=saved.config, + metadata=saved.metadata, + created_at=saved.checkpoint["ts"], + parent_config=saved.parent_config, ) @staticmethod @@ -457,7 +451,9 @@ class Pregel( 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_checkpoint_ns}'") + raise ValueError( + f"Missing checkpoint for checkpoint NS '{root_checkpoint_ns}'" + ) return state_snapshot def get_state( @@ -468,9 +464,7 @@ class Pregel( raise ValueError("No checkpointer set") if include_subgraph_state: - checkpoint_tuples = self.checkpointer.list( - config, include_nested_checkpoints=True - ) + checkpoint_tuples = self.checkpointer.list(config) else: checkpoint_tuples = iter([self.checkpointer.get_tuple(config)]) @@ -496,20 +490,16 @@ class Pregel( existing_checkpoint_id is None or saved_checkpoint_id > existing_checkpoint_id ): - state_snapshot = self._prepare_state_snapshot(checkpoint_tuple, config) + state_snapshot = self._prepare_state_snapshot(checkpoint_tuple) checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot checkpoint_ns_to_checkpoint_id[ saved_checkpoint_ns ] = saved_checkpoint_id if not checkpoint_ns_to_state_snapshots: - error_msg = ( - f"Could not find checkpoints for checkpoint NS '{checkpoint_ns}'" + return StateSnapshot( + values={}, next=(), config=config, checkpoint=empty_checkpoint() ) - if checkpoint_id: - error_msg += f" and checkpoint ID '{checkpoint_id}'" - - raise ValueError(error_msg) state_snapshot = self._assemble_state_snapshot_hierarchy( checkpoint_ns, checkpoint_ns_to_state_snapshots @@ -524,9 +514,7 @@ class Pregel( raise ValueError("No checkpointer set") if include_subgraph_state: - checkpoint_tuples = self.checkpointer.alist( - config, include_nested_checkpoints=True - ) + checkpoint_tuples = self.checkpointer.alist(config) else: async def alist_checkpoints(): @@ -556,20 +544,18 @@ class Pregel( existing_checkpoint_id is None or saved_checkpoint_id > existing_checkpoint_id ): - state_snapshot = self._prepare_state_snapshot(checkpoint_tuple, config) + state_snapshot = await self._prepare_state_snapshot_async( + checkpoint_tuple + ) checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot checkpoint_ns_to_checkpoint_id[ saved_checkpoint_ns ] = saved_checkpoint_id if not checkpoint_ns_to_state_snapshots: - error_msg = ( - f"Could not find checkpoints for checkpoint NS '{checkpoint_ns}'" + return StateSnapshot( + values={}, next=(), config=config, checkpoint=empty_checkpoint() ) - if checkpoint_id: - error_msg += f" and checkpoint ID '{checkpoint_id}'" - - raise ValueError(error_msg) state_snapshot = self._assemble_state_snapshot_hierarchy( checkpoint_ns, checkpoint_ns_to_state_snapshots @@ -593,41 +579,25 @@ 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, _ in self.checkpointer.list( + + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + for checkpoint_tuple in self.checkpointer.list( config, before=before, limit=limit, filter=filter ): + if ( + checkpoint_tuple.config["configurable"]["checkpoint_ns"] + != checkpoint_ns + ): + # only list root checkpoints here + continue + if include_subgraph_state: - state_snapshot = self.get_state(config, include_subgraph_state=True) + state_snapshot = self.get_state( + checkpoint_tuple.config, include_subgraph_state=True + ) yield state_snapshot else: - with ChannelsManager( - { - k: LastValue(None) if isinstance(c, Context) else c - for k, c in self.channels.items() - }, - checkpoint, - config, - ) as channels, ManagedValuesManager( - self.managed_values_dict, ensure_config(config) - ) 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(t.name for t in next_tasks), - config, - metadata, - checkpoint["ts"], - parent_config, - ) + yield self._prepare_state_snapshot(checkpoint_tuple) async def aget_state_history( self, @@ -646,46 +616,25 @@ 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, - _, - ) in self.checkpointer.alist(config, before=before, limit=limit, filter=filter): + + checkpoint_ns = config["configurable"].get("checkpoint_ns", "") + async for checkpoint_tuple in self.checkpointer.alist( + config, before=before, limit=limit, filter=filter + ): + if ( + checkpoint_tuple.config["configurable"]["checkpoint_ns"] + != checkpoint_ns + ): + # only list root checkpoints here + continue + if include_subgraph_state: state_snapshot = await self.aget_state( - config, include_subgraph_state=True + checkpoint_tuple.config, include_subgraph_state=True ) yield state_snapshot else: - async with AsyncChannelsManager( - { - k: LastValue(None) if isinstance(c, Context) else c - for k, c in self.channels.items() - }, - checkpoint, - config, - ) as channels, AsyncManagedValuesManager( - self.managed_values_dict, ensure_config(config) - ) 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(t.name for t in next_tasks), - config, - metadata, - checkpoint["ts"], - parent_config, - ) + yield await self._prepare_state_snapshot_async(checkpoint_tuple) def update_state( self, @@ -953,32 +902,36 @@ class Pregel( stream_mode = stream_mode if stream_mode is not None else self.stream_mode if not isinstance(stream_mode, list): stream_mode = [stream_mode] + + if config and config.get("configurable", {}).get(CONFIG_KEY_CHECKPOINTER): + parent_checkpointer = config["configurable"][CONFIG_KEY_CHECKPOINTER] + else: + parent_checkpointer = None + if config and config.get("configurable", {}).get(CONFIG_KEY_READ) is not None: # if being called as a node in another graph, always use values mode stream_mode = ["values"] - if self.checkpointer is None: + if parent_checkpointer is not None and self.checkpointer is None: raise ValueError( "Missing checkpointer for subgraph. " "Please compile the subgraph graph with checkpointer=INHERIT_CHECKPOINTER (from langgraph.pregel import INHERIT_CHECKPOINTER)." ) - if self.checkpointer != INHERIT_CHECKPOINTER: + if ( + parent_checkpointer is not None + and self.checkpointer != INHERIT_CHECKPOINTER + ): raise ValueError( "Custom checkpointers for subgraphs are not allowed. " "Please compile the subgraph graph with checkpointer=INHERIT_CHECKPOINTER (from langgraph.pregel import INHERIT_CHECKPOINTER)." ) - if ( - config is not None - and config.get("configurable", {}).get(CONFIG_KEY_CHECKPOINTER) - and self.checkpointer == INHERIT_CHECKPOINTER - ): - checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][ - CONFIG_KEY_CHECKPOINTER - ] - else: - checkpointer = self.checkpointer + checkpointer = ( + parent_checkpointer + if parent_checkpointer is not None + else self.checkpointer + ) return ( debug, stream_mode, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index ee5efaa37..b94a73494 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7629,13 +7629,13 @@ async def test_nested_graph_state( my_key: str my_other_key: str - def inner_1(state: InnerState): + async def inner_1(state: InnerState): return { "my_key": state["my_key"] + " here", "my_other_key": state["my_key"], } - def inner_2(state: InnerState): + async def inner_2(state: InnerState): return { "my_key": state["my_key"] + " and there", "my_other_key": state["my_key"], @@ -7651,10 +7651,10 @@ async def test_nested_graph_state( class State(TypedDict): my_key: str - def outer_1(state: State): + async def outer_1(state: State): return {"my_key": "hi " + state["my_key"]} - def outer_2(state: State): + async def outer_2(state: State): return {"my_key": state["my_key"] + " and back again"} graph = StateGraph(State) @@ -7699,7 +7699,7 @@ async def test_nested_graph_state( }, subgraph_state_snapshots=None, ) - assert app.get_state(config, include_subgraph_state=True) == StateSnapshot( + assert await app.aget_state(config, include_subgraph_state=True) == StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), config={ @@ -7755,7 +7755,9 @@ async def test_nested_graph_state( ) }, ) - assert list(app.get_state_history(config, include_subgraph_state=True)) == [ + assert [ + s async for s in app.aget_state_history(config, include_subgraph_state=True) + ] == [ StateSnapshot( values={"my_key": "hi my value"}, next=("inner",), @@ -8089,10 +8091,10 @@ async def test_doubly_nested_graph_state( class GrandChildState(TypedDict): my_key: str - def grandchild_1(state: ChildState): + async def grandchild_1(state: ChildState): return {"my_key": state["my_key"] + " here"} - def grandchild_2(state: ChildState): + async def grandchild_2(state: ChildState): return { "my_key": state["my_key"] + " and there", } @@ -8114,10 +8116,10 @@ async def test_doubly_nested_graph_state( child.set_entry_point("child_1") child.set_finish_point("child_1") - def parent_1(state: State): + async def parent_1(state: State): return {"my_key": "hi " + state["my_key"]} - def parent_2(state: State): + async def parent_2(state: State): return {"my_key": state["my_key"] + " and back again"} graph = StateGraph(State) From 503304458721134eb927a3dc057d0061e76c8e37 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 12 Aug 2024 16:12:47 -0400 Subject: [PATCH 10/41] update checkpointer tests --- libs/checkpoint-postgres/tests/test_async.py | 26 +++++-------------- libs/checkpoint-postgres/tests/test_sync.py | 25 +++++------------- .../checkpoint-sqlite/tests/test_aiosqlite.py | 26 +++++-------------- libs/checkpoint-sqlite/tests/test_sqlite.py | 25 +++++------------- .../langgraph/checkpoint/memory/__init__.py | 5 +--- libs/checkpoint/tests/test_memory.py | 23 +++++++--------- 6 files changed, 34 insertions(+), 96 deletions(-) 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/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 3bb7f8803..76901e2e0 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -173,9 +173,6 @@ 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 "" - ) for thread_id in thread_ids: for checkpoint_ns in self.storage[thread_id].keys(): for checkpoint_id, ( @@ -198,7 +195,7 @@ class MemorySaver( # filter by metadata metadata = self.serde.loads_typed(metadata_b) if filter and not all( - query_value == metadata[query_key] + query_value == metadata.get(query_key) for query_key, query_value in filter.items() ): continue 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) From 6d4cdc94560b66dc424628115e49da1d45faed79 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 12 Aug 2024 16:29:15 -0400 Subject: [PATCH 11/41] opt-in --- libs/langgraph/langgraph/pregel/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 9c591979f..bf491c6e3 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -912,9 +912,9 @@ class Pregel( # if being called as a node in another graph, always use values mode stream_mode = ["values"] - if parent_checkpointer is not None and self.checkpointer is None: + if (interrupt_before or interrupt_after) and parent_checkpointer is not None and self.checkpointer is None: raise ValueError( - "Missing checkpointer for subgraph. " + "Missing checkpointer for a subgraph with interrupts. " "Please compile the subgraph graph with checkpointer=INHERIT_CHECKPOINTER (from langgraph.pregel import INHERIT_CHECKPOINTER)." ) From abe9b7c08eb887932a0f553f87040043ac2c4b13 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 12 Aug 2024 16:32:22 -0400 Subject: [PATCH 12/41] lint --- libs/langgraph/langgraph/pregel/__init__.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index bf491c6e3..d2b996d08 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -912,7 +912,11 @@ class Pregel( # if being called as a node in another graph, always use values mode stream_mode = ["values"] - if (interrupt_before or interrupt_after) and parent_checkpointer is not None and self.checkpointer is None: + if ( + (interrupt_before or interrupt_after) + and parent_checkpointer is not None + and self.checkpointer is None + ): raise ValueError( "Missing checkpointer for a subgraph with interrupts. " "Please compile the subgraph graph with checkpointer=INHERIT_CHECKPOINTER (from langgraph.pregel import INHERIT_CHECKPOINTER)." From f65d9b2b7d05b4cc0d09a41c346b177ac81042f8 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 12 Aug 2024 19:50:07 -0400 Subject: [PATCH 13/41] pass subgraph nodes/channels --- libs/langgraph/langgraph/pregel/__init__.py | 85 ++++++++++++++++++--- libs/langgraph/tests/test_pregel.py | 10 +-- libs/langgraph/tests/test_pregel_async.py | 10 +-- 3 files changed, 85 insertions(+), 20 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index d2b996d08..13aacf0e5 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -360,11 +360,41 @@ class Pregel( if is_managed_value(v) } - def _prepare_state_snapshot(self, saved: CheckpointTuple) -> StateSnapshot: + def _get_nodes_and_channels( + self, checkpoint_ns: str + ) -> tuple[Mapping[str, PregelNode], Mapping[str, BaseChannel]]: + if checkpoint_ns == "": + return self.nodes, self.channels + + path = checkpoint_ns.split(CHECKPOINT_NAMESPACE_SEPARATOR) + nodes = self.nodes + channels = self.channels + for subgraph_node_name in path: + if subgraph_node_name not in nodes: + raise ValueError(f"Couldn't find node '{subgraph_node_name}'.") + + subgraph_node = nodes[subgraph_node_name].get_node() + + if not isinstance(subgraph_node, RunnableSequence): + break + + first_step = subgraph_node.steps[0] + if isinstance(first_step, Pregel): + nodes = first_step.nodes + channels = first_step.channels + + return nodes, channels + + def _prepare_state_snapshot( + self, + saved: CheckpointTuple, + nodes: Mapping[str, PregelNode], + channels: Mapping[str, BaseChannel], + ) -> StateSnapshot: with ChannelsManager( { k: LastValue(None) if isinstance(c, Context) else c - for k, c in self.channels.items() + for k, c in channels.items() }, saved.checkpoint, saved.config, @@ -373,7 +403,7 @@ class Pregel( ) as managed: next_tasks = prepare_next_tasks( saved.checkpoint, - self.nodes, + nodes, channels, managed, saved.config, @@ -390,12 +420,15 @@ class Pregel( ) async def _prepare_state_snapshot_async( - self, saved: CheckpointTuple + self, + saved: CheckpointTuple, + nodes: Mapping[str, PregelNode], + channels: Mapping[str, BaseChannel], ) -> StateSnapshot: async with AsyncChannelsManager( { k: LastValue(None) if isinstance(c, Context) else c - for k, c in self.channels.items() + for k, c in channels.items() }, saved.checkpoint, saved.config, @@ -404,7 +437,7 @@ class Pregel( ) as managed: next_tasks = prepare_next_tasks( saved.checkpoint, - self.nodes, + nodes, channels, managed, saved.config, @@ -472,6 +505,9 @@ class Pregel( checkpoint_id = config["configurable"].get("checkpoint_id") checkpoint_ns_to_checkpoint_id: dict[str, str] = {} checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} + checkpoint_ns_to_nodes_and_channels: dict[ + str, tuple[Mapping[str, PregelNode], Mapping[str, BaseChannel]] + ] = {} for checkpoint_tuple in checkpoint_tuples: saved_checkpoint_ns = checkpoint_tuple.config["configurable"][ "checkpoint_ns" @@ -490,7 +526,17 @@ class Pregel( existing_checkpoint_id is None or saved_checkpoint_id > existing_checkpoint_id ): - state_snapshot = self._prepare_state_snapshot(checkpoint_tuple) + if saved_checkpoint_ns not in checkpoint_ns_to_nodes_and_channels: + checkpoint_ns_to_nodes_and_channels[ + saved_checkpoint_ns + ] = self._get_nodes_and_channels(saved_checkpoint_ns) + + nodes, channels = checkpoint_ns_to_nodes_and_channels[ + saved_checkpoint_ns + ] + state_snapshot = self._prepare_state_snapshot( + checkpoint_tuple, nodes, channels + ) checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot checkpoint_ns_to_checkpoint_id[ saved_checkpoint_ns @@ -526,6 +572,9 @@ class Pregel( checkpoint_id = config["configurable"].get("checkpoint_id") checkpoint_ns_to_checkpoint_id: dict[str, str] = {} checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} + checkpoint_ns_to_nodes_and_channels: dict[ + str, tuple[Mapping[str, PregelNode], Mapping[str, BaseChannel]] + ] = {} async for checkpoint_tuple in checkpoint_tuples: saved_checkpoint_ns = checkpoint_tuple.config["configurable"][ "checkpoint_ns" @@ -544,8 +593,16 @@ class Pregel( existing_checkpoint_id is None or saved_checkpoint_id > existing_checkpoint_id ): + if saved_checkpoint_ns not in checkpoint_ns_to_nodes_and_channels: + checkpoint_ns_to_nodes_and_channels[ + saved_checkpoint_ns + ] = self._get_nodes_and_channels(saved_checkpoint_ns) + + nodes, channels = checkpoint_ns_to_nodes_and_channels[ + saved_checkpoint_ns + ] state_snapshot = await self._prepare_state_snapshot_async( - checkpoint_tuple + checkpoint_tuple, nodes, channels ) checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot checkpoint_ns_to_checkpoint_id[ @@ -597,7 +654,10 @@ class Pregel( ) yield state_snapshot else: - yield self._prepare_state_snapshot(checkpoint_tuple) + nodes, channels = self._get_nodes_and_channels( + checkpoint_tuple.config["configurable"]["checkpoint_ns"] + ) + yield self._prepare_state_snapshot(checkpoint_tuple, nodes, channels) async def aget_state_history( self, @@ -634,7 +694,12 @@ class Pregel( ) yield state_snapshot else: - yield await self._prepare_state_snapshot_async(checkpoint_tuple) + nodes, channels = self._get_nodes_and_channels( + checkpoint_tuple.config["configurable"]["checkpoint_ns"] + ) + yield await self._prepare_state_snapshot_async( + checkpoint_tuple, nodes, channels + ) def update_state( self, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 0f4fbb5b2..8780bc042 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8604,7 +8604,7 @@ def test_nested_graph_interrupts( assert child_state_history == [ StateSnapshot( values={"my_key": "hi my value here"}, - next=(), + next=("inner_2",), config={ "configurable": { "thread_id": "6", @@ -9218,7 +9218,7 @@ def test_nested_graph_state( subgraph_state_snapshots={ "inner": StateSnapshot( values={"my_key": "hi my value here"}, - next=(), + next=("inner_2",), config={ "configurable": { "thread_id": "1", @@ -9275,7 +9275,7 @@ def test_nested_graph_state( subgraph_state_snapshots={ "inner": StateSnapshot( values={"my_key": "hi my value here"}, - next=(), + next=("inner_2",), config={ "configurable": { "thread_id": "1", @@ -9676,7 +9676,7 @@ def test_doubly_nested_graph_state( subgraph_state_snapshots={ "child": StateSnapshot( values={"my_key": "hi my value"}, - next=(), + next=("child_1",), config={ "configurable": { "thread_id": "1", @@ -9696,7 +9696,7 @@ def test_doubly_nested_graph_state( subgraph_state_snapshots={ "child_1": StateSnapshot( values={"my_key": "hi my value here"}, - next=(), + next=("grandchild_2",), config={ "configurable": { "thread_id": "1", diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b94a73494..af624f639 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7106,7 +7106,7 @@ async def test_nested_graph_interrupts( assert child_state_history == [ StateSnapshot( values={"my_key": "hi my value here"}, - next=(), + next=("inner_2",), config={ "configurable": { "thread_id": "6", @@ -7725,7 +7725,7 @@ async def test_nested_graph_state( subgraph_state_snapshots={ "inner": StateSnapshot( values={"my_key": "hi my value here"}, - next=(), + next=("inner_2",), config={ "configurable": { "thread_id": "1", @@ -7784,7 +7784,7 @@ async def test_nested_graph_state( subgraph_state_snapshots={ "inner": StateSnapshot( values={"my_key": "hi my value here"}, - next=(), + next=("inner_2",), config={ "configurable": { "thread_id": "1", @@ -8187,7 +8187,7 @@ async def test_doubly_nested_graph_state( subgraph_state_snapshots={ "child": StateSnapshot( values={"my_key": "hi my value"}, - next=(), + next=("child_1",), config={ "configurable": { "thread_id": "1", @@ -8207,7 +8207,7 @@ async def test_doubly_nested_graph_state( subgraph_state_snapshots={ "child_1": StateSnapshot( values={"my_key": "hi my value here"}, - next=(), + next=("grandchild_2",), config={ "configurable": { "thread_id": "1", From 0135c6f743742af1887dca841939c3f5c69c4d4e Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 13 Aug 2024 10:07:52 -0400 Subject: [PATCH 14/41] correct check for using parent checkpointer --- libs/langgraph/langgraph/pregel/__init__.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 13aacf0e5..0b7dd76c9 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -989,6 +989,7 @@ class Pregel( if ( parent_checkpointer is not None + and self.checkpointer is not None and self.checkpointer != INHERIT_CHECKPOINTER ): raise ValueError( @@ -996,11 +997,14 @@ class Pregel( "Please compile the subgraph graph with checkpointer=INHERIT_CHECKPOINTER (from langgraph.pregel import INHERIT_CHECKPOINTER)." ) - checkpointer = ( - parent_checkpointer - if parent_checkpointer is not None - else self.checkpointer - ) + if ( + parent_checkpointer is not None + and self.checkpointer == INHERIT_CHECKPOINTER + ): + checkpointer = parent_checkpointer + else: + checkpointer = self.checkpointer + return ( debug, stream_mode, From 32952747116d769ecbdb51dc2d39fc7ceef91c55 Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 13 Aug 2024 10:23:53 -0400 Subject: [PATCH 15/41] fix empty snapshot --- libs/langgraph/langgraph/pregel/__init__.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 0b7dd76c9..58e34dd59 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -499,7 +499,8 @@ class Pregel( if include_subgraph_state: checkpoint_tuples = self.checkpointer.list(config) else: - checkpoint_tuples = iter([self.checkpointer.get_tuple(config)]) + checkpoint_tuple = self.checkpointer.get_tuple(config) + checkpoint_tuples = iter([checkpoint_tuple] if checkpoint_tuple else []) checkpoint_ns = config["configurable"].get("checkpoint_ns", "") checkpoint_id = config["configurable"].get("checkpoint_id") @@ -544,7 +545,7 @@ class Pregel( if not checkpoint_ns_to_state_snapshots: return StateSnapshot( - values={}, next=(), config=config, checkpoint=empty_checkpoint() + values={}, next=(), config=config, metadata=None, created_at=None ) state_snapshot = self._assemble_state_snapshot_hierarchy( @@ -564,7 +565,9 @@ class Pregel( else: async def alist_checkpoints(): - yield await self.checkpointer.aget_tuple(config) + checkpoint_tuple = await self.checkpointer.aget_tuple(config) + if checkpoint_tuple: + yield checkpoint_tuple checkpoint_tuples = alist_checkpoints() @@ -611,7 +614,7 @@ class Pregel( if not checkpoint_ns_to_state_snapshots: return StateSnapshot( - values={}, next=(), config=config, checkpoint=empty_checkpoint() + values={}, next=(), config=config, metadata=None, created_at=None ) state_snapshot = self._assemble_state_snapshot_hierarchy( From 58887a5a3b93cc7c1b8da8122d7f54bb227b97f3 Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 13 Aug 2024 16:23:30 -0400 Subject: [PATCH 16/41] checkpoints/interrupts for subgraphs triggered by sends --- .../langgraph/checkpoint/serde/jsonplus.py | 2 +- .../langgraph/checkpoint/serde/types.py | 1 + libs/langgraph/langgraph/constants.py | 12 +- libs/langgraph/langgraph/graph/graph.py | 14 +- libs/langgraph/langgraph/graph/state.py | 18 +- libs/langgraph/langgraph/pregel/__init__.py | 16 +- libs/langgraph/langgraph/pregel/algo.py | 8 +- libs/langgraph/tests/test_pregel.py | 360 +++++++++++++++++ libs/langgraph/tests/test_pregel_async.py | 362 ++++++++++++++++++ 9 files changed, 774 insertions(+), 19 deletions(-) diff --git a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py index 2b7bce5d4..08feeabfa 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/jsonplus.py @@ -66,7 +66,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 de61ef78b..cc5c1fa8b 100644 --- a/libs/checkpoint/langgraph/checkpoint/serde/types.py +++ b/libs/checkpoint/langgraph/checkpoint/serde/types.py @@ -55,6 +55,7 @@ class SendProtocol(Protocol): # Mirrors langgraph.constants.Send node: str arg: Any + id: str def __hash__(self) -> int: ... diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index ae447d62e..89a2effa7 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -1,4 +1,5 @@ -from typing import Any +from typing import Any, Optional +from uuid import uuid4 INPUT = "__input__" CONFIG_KEY_SEND = "__pregel_send" @@ -22,6 +23,7 @@ START = "__start__" END = "__end__" CHECKPOINT_NAMESPACE_SEPARATOR = "|" +SEND_CHECKPOINT_NAMESPACE_SEPARATOR = ":" class Send: @@ -40,6 +42,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 @@ -67,23 +70,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)) 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 c17373820..6fc4eb8cd 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -27,6 +27,7 @@ from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.constants import ( CHECKPOINT_NAMESPACE_SEPARATOR, END, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, START, TAG_HIDDEN, Send, @@ -159,10 +160,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 e9971a5f1..8e7298616 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -29,7 +29,11 @@ from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitFo from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.channels.named_barrier_value import NamedBarrierValue -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, @@ -312,10 +316,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 58e34dd59..96f720be5 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -72,6 +72,7 @@ from langgraph.constants import ( CONFIG_KEY_RESUMING, CONFIG_KEY_SEND, INTERRUPT, + SEND_CHECKPOINT_NAMESPACE_SEPARATOR, ) from langgraph.errors import GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ( @@ -370,6 +371,15 @@ class Pregel( nodes = self.nodes channels = self.channels for subgraph_node_name in path: + # if we have this separator it means we have a node that was triggered by Send + if SEND_CHECKPOINT_NAMESPACE_SEPARATOR in subgraph_node_name: + name_parts = subgraph_node_name.split( + SEND_CHECKPOINT_NAMESPACE_SEPARATOR + ) + if len(name_parts) != 2: + raise ValueError(f"Malformed node name '{subgraph_node_name}'") + + subgraph_node_name = name_parts[0] if subgraph_node_name not in nodes: raise ValueError(f"Couldn't find node '{subgraph_node_name}'.") @@ -1190,8 +1200,7 @@ class Pregel( ) if not done: break # timed out - for fut in done: - task = futures.pop(fut) + for fut, task in zip(done, [futures.pop(fut) for fut in done]): if fut.exception() is not None: # we got an exception, break out of while loop # exception will be handled in panic_or_proceed @@ -1435,8 +1444,7 @@ class Pregel( ) if not done: break # timed out - for fut in done: - task = futures.pop(fut) + for fut, task in zip(done, [futures.pop(fut) for fut in done]): if fut.exception() is not None: # we got an exception, break out of while loop # exception will be handled in panic_or_proceed diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 2ddbcd873..ec8c18650 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -280,9 +280,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))) @@ -318,6 +318,10 @@ def prepare_next_tasks( PregelTaskWrites(packet.node, writes, triggers), config, ), + CONFIG_KEY_CHECKPOINTER: checkpointer, + CONFIG_KEY_RESUMING: is_resuming, + "checkpoint_id": checkpoint["id"], + "checkpoint_ns": checkpoint_ns, # in Send we can't checkpoint nested graphs # as they could be running in parallel }, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 8780bc042..881487c93 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9807,6 +9807,366 @@ def test_doubly_nested_graph_state( ) +@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( + checkpointer=INHERIT_CHECKPOINTER, 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, include_subgraph_state=True) + 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" + + expected_snapshot = StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + 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_nodes[0]: StateSnapshot( + values={"jokes": []}, + next=("generate",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": []}, + next=("generate",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ) + 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, include_subgraph_state=True) + subgraph_nodes, _ = zip( + *( + sorted( + actual_snapshot.subgraph_state_snapshots.items(), + key=lambda x: x[1].values["jokes"][0], + ) + ) + ) + assert len(subgraph_nodes) == 2 + for subgraph_node in subgraph_nodes: + assert subgraph_node.split(":")[0] == "generate_joke" + expected_snapshot = StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + 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={ + subgraph_nodes[0]: StateSnapshot( + values={"jokes": ["Joke about cats - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": ["Joke about dogs - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ) + assert actual_snapshot == expected_snapshot + + # test full history + actual_history = list(graph.get_state_history(config, include_subgraph_state=True)) + expected_history = [ + StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + 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"), + 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_nodes[0]: StateSnapshot( + values={"jokes": ["Joke about cats - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": ["Joke about dogs - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ), + StateSnapshot( + values={"jokes": []}, + 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 af624f639..43eab4ea6 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -8318,6 +8318,368 @@ async def test_doubly_nested_graph_state( ) +@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( + checkpointer=INHERIT_CHECKPOINTER, 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, include_subgraph_state=True) + 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" + + expected_snapshot = StateSnapshot( + values={"subjects": ["cats", "dogs"], "jokes": []}, + 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_nodes[0]: StateSnapshot( + values={"jokes": []}, + next=("generate",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": []}, + next=("generate",), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ) + 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, include_subgraph_state=True) + subgraph_nodes, _ = zip( + *( + sorted( + actual_snapshot.subgraph_state_snapshots.items(), + key=lambda x: x[1].values["jokes"][0], + ) + ) + ) + assert len(subgraph_nodes) == 2 + for subgraph_node in subgraph_nodes: + assert subgraph_node.split(":")[0] == "generate_joke" + expected_snapshot = StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + 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={ + subgraph_nodes[0]: StateSnapshot( + values={"jokes": ["Joke about cats - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": ["Joke about dogs - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ) + assert actual_snapshot == expected_snapshot + + # test full history + actual_history = [ + c async for c in graph.aget_state_history(config, include_subgraph_state=True) + ] + expected_history = [ + StateSnapshot( + values={ + "subjects": ["cats", "dogs"], + "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], + }, + 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"), + 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_nodes[0]: StateSnapshot( + values={"jokes": ["Joke about cats - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[0], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + subgraph_nodes[1]: StateSnapshot( + values={"jokes": ["Joke about dogs - hohoho"]}, + next=(), + config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + metadata={ + "source": "loop", + "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, + "step": 2, + }, + created_at=AnyStr(), + parent_config={ + "configurable": { + "thread_id": "1", + "checkpoint_ns": subgraph_nodes[1], + "checkpoint_id": AnyStr(), + } + }, + subgraph_state_snapshots=None, + ), + }, + ), + StateSnapshot( + values={"jokes": []}, + 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. From d9618880a3525a79635f8b2f5e98f76af26ed28c Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 13 Aug 2024 19:46:13 -0400 Subject: [PATCH 17/41] update logic for latest snapshot's subgraph snapshots --- libs/langgraph/langgraph/pregel/__init__.py | 18 +- libs/langgraph/tests/test_pregel.py | 376 +++++++++----------- libs/langgraph/tests/test_pregel_async.py | 375 +++++++++---------- 3 files changed, 351 insertions(+), 418 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 96f720be5..81f08d98f 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -506,14 +506,15 @@ class Pregel( 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_tuple = self.checkpointer.get_tuple(config) checkpoint_tuples = iter([checkpoint_tuple] if checkpoint_tuple else []) - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - checkpoint_id = config["configurable"].get("checkpoint_id") + 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_nodes_and_channels: dict[ @@ -526,7 +527,7 @@ class Pregel( saved_checkpoint_id = checkpoint_tuple.config["configurable"][ "checkpoint_id" ] - if checkpoint_id and checkpoint_id != saved_checkpoint_id: + if checkpoint_id != saved_checkpoint_id: continue existing_checkpoint_id = checkpoint_ns_to_checkpoint_id.get( @@ -570,19 +571,20 @@ class Pregel( 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(): - checkpoint_tuple = await self.checkpointer.aget_tuple(config) if checkpoint_tuple: yield checkpoint_tuple checkpoint_tuples = alist_checkpoints() - checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - checkpoint_id = config["configurable"].get("checkpoint_id") + 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_nodes_and_channels: dict[ @@ -595,7 +597,7 @@ class Pregel( saved_checkpoint_id = checkpoint_tuple.config["configurable"][ "checkpoint_id" ] - if checkpoint_id and checkpoint_id != saved_checkpoint_id: + if checkpoint_id != saved_checkpoint_id: continue existing_checkpoint_id = checkpoint_ns_to_checkpoint_id.get( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 881487c93..b64d97b57 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -9373,39 +9373,77 @@ def test_nested_graph_state( "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, - ) - }, ) + # test loading inner snapshot + child_snapshot = app.get_state( + {"configurable": {"thread_id": "1", "checkpoint_ns": "inner"}} + ) + assert child_snapshot == 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, + ) + # test looking up parent state by checkpoint ID + assert app.get_state( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], + } + }, + 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": child_snapshot}, + ) + # test full history at the end assert list(app.get_state_history(config, include_subgraph_state=True)) == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, @@ -9482,10 +9520,6 @@ def test_nested_graph_state( "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"}, @@ -9749,61 +9783,99 @@ def test_doubly_nested_graph_state( "checkpoint_id": AnyStr(), } }, - subgraph_state_snapshots={ - "child": StateSnapshot( - values={"my_key": "hi my value here and there"}, - 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"}, - 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, - ) - }, - ) + ) + # test getting grandchild snapshot + grandchild_snapshot = app.get_state( + {"configurable": {"thread_id": "1", "checkpoint_ns": "child|child_1"}} + ) + assert grandchild_snapshot == StateSnapshot( + values={"my_key": "hi my value here and there"}, + 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, + ) + # 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"}, + 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": grandchild_snapshot}, + ) + # test getting parent snapshot for a checkpoint ID + assert app.get_state( + { + "configurable": { + "thread_id": "1", + "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], + } + }, + include_subgraph_state=True, + ) == 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={"child": child_snapshot}, ) @@ -9939,17 +10011,6 @@ def test_send_to_nested_graphs( } actual_snapshot = graph.get_state(config, include_subgraph_state=True) - subgraph_nodes, _ = zip( - *( - sorted( - actual_snapshot.subgraph_state_snapshots.items(), - key=lambda x: x[1].values["jokes"][0], - ) - ) - ) - assert len(subgraph_nodes) == 2 - for subgraph_node in subgraph_nodes: - assert subgraph_node.split(":")[0] == "generate_joke" expected_snapshot = StateSnapshot( values={ "subjects": ["cats", "dogs"], @@ -9981,63 +10042,19 @@ def test_send_to_nested_graphs( "checkpoint_id": AnyStr(), } }, - subgraph_state_snapshots={ - subgraph_nodes[0]: StateSnapshot( - values={"jokes": ["Joke about cats - hohoho"]}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - subgraph_nodes[1]: StateSnapshot( - values={"jokes": ["Joke about dogs - hohoho"]}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - }, ) assert actual_snapshot == expected_snapshot # test full history actual_history = list(graph.get_state_history(config, include_subgraph_state=True)) + + # 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={ @@ -10091,58 +10108,7 @@ def test_send_to_nested_graphs( "checkpoint_id": AnyStr(), } }, - subgraph_state_snapshots={ - subgraph_nodes[0]: StateSnapshot( - values={"jokes": ["Joke about cats - hohoho"]}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - subgraph_nodes[1]: StateSnapshot( - values={"jokes": ["Joke about dogs - hohoho"]}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - }, + subgraph_state_snapshots=subgraph_state_snapshots, ), StateSnapshot( values={"jokes": []}, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 43eab4ea6..7b2e3222b 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7882,39 +7882,77 @@ async def test_nested_graph_state( "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, - ) - }, ) + # test loading inner snapshot + child_snapshot = await app.aget_state( + {"configurable": {"thread_id": "1", "checkpoint_ns": "inner"}} + ) + assert child_snapshot == 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, + ) + # test looking up parent state by checkpoint ID + assert await app.aget_state( + { + "configurable": { + "thread_id": "1", + "checkpoint_ns": "", + "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], + } + }, + 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": child_snapshot}, + ) + # test full history at the end assert [ s async for s in app.aget_state_history(config, include_subgraph_state=True) ] == [ @@ -7993,10 +8031,6 @@ async def test_nested_graph_state( "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"}, @@ -8260,61 +8294,99 @@ async def test_doubly_nested_graph_state( "checkpoint_id": AnyStr(), } }, - subgraph_state_snapshots={ - "child": StateSnapshot( - values={"my_key": "hi my value here and there"}, - 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"}, - 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, - ) - }, - ) + ) + # test getting grandchild snapshot + grandchild_snapshot = await app.aget_state( + {"configurable": {"thread_id": "1", "checkpoint_ns": "child|child_1"}} + ) + assert grandchild_snapshot == StateSnapshot( + values={"my_key": "hi my value here and there"}, + 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, + ) + # 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"}, + 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": grandchild_snapshot}, + ) + # test getting parent snapshot for a checkpoint ID + assert await app.aget_state( + { + "configurable": { + "thread_id": "1", + "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], + } + }, + include_subgraph_state=True, + ) == 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={"child": child_snapshot}, ) @@ -8450,17 +8522,6 @@ async def test_send_to_nested_graphs( } actual_snapshot = await graph.aget_state(config, include_subgraph_state=True) - subgraph_nodes, _ = zip( - *( - sorted( - actual_snapshot.subgraph_state_snapshots.items(), - key=lambda x: x[1].values["jokes"][0], - ) - ) - ) - assert len(subgraph_nodes) == 2 - for subgraph_node in subgraph_nodes: - assert subgraph_node.split(":")[0] == "generate_joke" expected_snapshot = StateSnapshot( values={ "subjects": ["cats", "dogs"], @@ -8492,58 +8553,6 @@ async def test_send_to_nested_graphs( "checkpoint_id": AnyStr(), } }, - subgraph_state_snapshots={ - subgraph_nodes[0]: StateSnapshot( - values={"jokes": ["Joke about cats - hohoho"]}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - subgraph_nodes[1]: StateSnapshot( - values={"jokes": ["Joke about dogs - hohoho"]}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - }, ) assert actual_snapshot == expected_snapshot @@ -8551,6 +8560,13 @@ async def test_send_to_nested_graphs( actual_history = [ c async for c in graph.aget_state_history(config, include_subgraph_state=True) ] + # 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={ @@ -8604,58 +8620,7 @@ async def test_send_to_nested_graphs( "checkpoint_id": AnyStr(), } }, - subgraph_state_snapshots={ - subgraph_nodes[0]: StateSnapshot( - values={"jokes": ["Joke about cats - hohoho"]}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"generate": {"jokes": ["Joke about cats - hohoho"]}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - subgraph_nodes[1]: StateSnapshot( - values={"jokes": ["Joke about dogs - hohoho"]}, - next=(), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - metadata={ - "source": "loop", - "writes": {"generate": {"jokes": ["Joke about dogs - hohoho"]}}, - "step": 2, - }, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - }, + subgraph_state_snapshots=subgraph_state_snapshots, ), StateSnapshot( values={"jokes": []}, From 409b915a3f38196fc95edfb1d12495f36a0da12e Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 14 Aug 2024 11:54:23 -0400 Subject: [PATCH 18/41] code review --- libs/langgraph/langgraph/constants.py | 2 +- libs/langgraph/langgraph/pregel/__init__.py | 333 +++++++++++--------- 2 files changed, 181 insertions(+), 154 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 89a2effa7..bc74619ff 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -86,7 +86,7 @@ class Send: 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}, id={self.id!r})" diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 81f08d98f..ffa53da85 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -195,6 +195,145 @@ class Channel: ) +def _get_nodes_and_channels( + nodes: Mapping[str, PregelNode], + channels: Mapping[str, BaseChannel], + checkpoint_ns: str, +) -> tuple[Mapping[str, PregelNode], Mapping[str, BaseChannel]]: + if checkpoint_ns == "": + return nodes, channels + + path = checkpoint_ns.split(CHECKPOINT_NAMESPACE_SEPARATOR) + for subgraph_node_name in path: + # if we have this separator it means we have a node that was triggered by Send + if SEND_CHECKPOINT_NAMESPACE_SEPARATOR in subgraph_node_name: + name_parts = subgraph_node_name.split(SEND_CHECKPOINT_NAMESPACE_SEPARATOR) + if len(name_parts) != 2: + raise ValueError(f"Malformed node name '{subgraph_node_name}'") + + subgraph_node_name = name_parts[0] + if subgraph_node_name not in nodes: + raise ValueError(f"Couldn't find node '{subgraph_node_name}'.") + + subgraph_node = nodes[subgraph_node_name].get_node() + + if not isinstance(subgraph_node, RunnableSequence): + break + + first_step = subgraph_node.steps[0] + if isinstance(first_step, Pregel): + nodes = first_step.nodes + channels = first_step.channels + + return nodes, channels + + +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 + + +def _prepare_state_snapshot( + saved: CheckpointTuple, + nodes: Mapping[str, PregelNode], + channels: Mapping[str, BaseChannel], + managed_values_dict: dict[str, ManagedValueSpec], + select_channels: str | list[str], +) -> StateSnapshot: + with ChannelsManager( + { + k: LastValue(None) if isinstance(c, Context) else c + for k, c in channels.items() + }, + saved.checkpoint, + saved.config, + ) as channels, ManagedValuesManager( + managed_values_dict, ensure_config(saved.config) + ) as managed: + next_tasks = prepare_next_tasks( + saved.checkpoint, + nodes, + channels, + managed, + saved.config, + -1, + for_execution=False, + ) + return StateSnapshot( + values=read_channels(channels, select_channels), + next=tuple(t.name for t in next_tasks), + config=saved.config, + metadata=saved.metadata, + created_at=saved.checkpoint["ts"], + parent_config=saved.parent_config, + ) + + +async def _prepare_state_snapshot_async( + saved: CheckpointTuple, + nodes: Mapping[str, PregelNode], + channels: Mapping[str, BaseChannel], + managed_values_dict: dict[str, ManagedValueSpec], + select_channels: str | list[str], +) -> StateSnapshot: + async with AsyncChannelsManager( + { + k: LastValue(None) if isinstance(c, Context) else c + for k, c in channels.items() + }, + saved.checkpoint, + saved.config, + ) as channels, AsyncManagedValuesManager( + managed_values_dict, ensure_config(saved.config) + ) as managed: + next_tasks = prepare_next_tasks( + saved.checkpoint, + nodes, + channels, + managed, + saved.config, + -1, + for_execution=False, + ) + return StateSnapshot( + values=read_channels(channels, select_channels), + next=tuple(t.name for t in next_tasks), + config=saved.config, + metadata=saved.metadata, + created_at=saved.checkpoint["ts"], + parent_config=saved.parent_config, + ) + + class Pregel( RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]] ): @@ -361,144 +500,6 @@ class Pregel( if is_managed_value(v) } - def _get_nodes_and_channels( - self, checkpoint_ns: str - ) -> tuple[Mapping[str, PregelNode], Mapping[str, BaseChannel]]: - if checkpoint_ns == "": - return self.nodes, self.channels - - path = checkpoint_ns.split(CHECKPOINT_NAMESPACE_SEPARATOR) - nodes = self.nodes - channels = self.channels - for subgraph_node_name in path: - # if we have this separator it means we have a node that was triggered by Send - if SEND_CHECKPOINT_NAMESPACE_SEPARATOR in subgraph_node_name: - name_parts = subgraph_node_name.split( - SEND_CHECKPOINT_NAMESPACE_SEPARATOR - ) - if len(name_parts) != 2: - raise ValueError(f"Malformed node name '{subgraph_node_name}'") - - subgraph_node_name = name_parts[0] - if subgraph_node_name not in nodes: - raise ValueError(f"Couldn't find node '{subgraph_node_name}'.") - - subgraph_node = nodes[subgraph_node_name].get_node() - - if not isinstance(subgraph_node, RunnableSequence): - break - - first_step = subgraph_node.steps[0] - if isinstance(first_step, Pregel): - nodes = first_step.nodes - channels = first_step.channels - - return nodes, channels - - def _prepare_state_snapshot( - self, - saved: CheckpointTuple, - nodes: Mapping[str, PregelNode], - channels: Mapping[str, BaseChannel], - ) -> StateSnapshot: - with ChannelsManager( - { - k: LastValue(None) if isinstance(c, Context) else c - for k, c in channels.items() - }, - saved.checkpoint, - saved.config, - ) as channels, ManagedValuesManager( - self.managed_values_dict, ensure_config(saved.config) - ) as managed: - next_tasks = prepare_next_tasks( - saved.checkpoint, - nodes, - channels, - managed, - saved.config, - -1, - for_execution=False, - ) - return StateSnapshot( - values=read_channels(channels, self.stream_channels_asis), - next=tuple(t.name for t in next_tasks), - config=saved.config, - metadata=saved.metadata, - created_at=saved.checkpoint["ts"], - parent_config=saved.parent_config, - ) - - async def _prepare_state_snapshot_async( - self, - saved: CheckpointTuple, - nodes: Mapping[str, PregelNode], - channels: Mapping[str, BaseChannel], - ) -> StateSnapshot: - async with AsyncChannelsManager( - { - k: LastValue(None) if isinstance(c, Context) else c - for k, c in channels.items() - }, - saved.checkpoint, - saved.config, - ) as channels, AsyncManagedValuesManager( - self.managed_values_dict, ensure_config(saved.config) - ) as managed: - next_tasks = prepare_next_tasks( - saved.checkpoint, - nodes, - channels, - managed, - saved.config, - -1, - for_execution=False, - ) - return StateSnapshot( - values=read_channels(channels, self.stream_channels_asis), - next=tuple(t.name for t in next_tasks), - config=saved.config, - metadata=saved.metadata, - created_at=saved.checkpoint["ts"], - parent_config=saved.parent_config, - ) - - @staticmethod - 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 - def get_state( self, config: RunnableConfig, *, include_subgraph_state: bool = False ) -> StateSnapshot: @@ -541,13 +542,19 @@ class Pregel( if saved_checkpoint_ns not in checkpoint_ns_to_nodes_and_channels: checkpoint_ns_to_nodes_and_channels[ saved_checkpoint_ns - ] = self._get_nodes_and_channels(saved_checkpoint_ns) + ] = _get_nodes_and_channels( + self.nodes, self.channels, saved_checkpoint_ns + ) nodes, channels = checkpoint_ns_to_nodes_and_channels[ saved_checkpoint_ns ] - state_snapshot = self._prepare_state_snapshot( - checkpoint_tuple, nodes, channels + state_snapshot = _prepare_state_snapshot( + checkpoint_tuple, + nodes, + channels, + self.managed_values_dict, + self.stream_channels_asis, ) checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot checkpoint_ns_to_checkpoint_id[ @@ -559,7 +566,7 @@ class Pregel( values={}, next=(), config=config, metadata=None, created_at=None ) - state_snapshot = self._assemble_state_snapshot_hierarchy( + state_snapshot = _assemble_state_snapshot_hierarchy( checkpoint_ns, checkpoint_ns_to_state_snapshots ) return state_snapshot @@ -611,13 +618,19 @@ class Pregel( if saved_checkpoint_ns not in checkpoint_ns_to_nodes_and_channels: checkpoint_ns_to_nodes_and_channels[ saved_checkpoint_ns - ] = self._get_nodes_and_channels(saved_checkpoint_ns) + ] = _get_nodes_and_channels( + self.nodes, self.channels, saved_checkpoint_ns + ) nodes, channels = checkpoint_ns_to_nodes_and_channels[ saved_checkpoint_ns ] - state_snapshot = await self._prepare_state_snapshot_async( - checkpoint_tuple, nodes, channels + state_snapshot = await _prepare_state_snapshot_async( + checkpoint_tuple, + nodes, + channels, + self.managed_values_dict, + self.stream_channels_asis, ) checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot checkpoint_ns_to_checkpoint_id[ @@ -629,7 +642,7 @@ class Pregel( values={}, next=(), config=config, metadata=None, created_at=None ) - state_snapshot = self._assemble_state_snapshot_hierarchy( + state_snapshot = _assemble_state_snapshot_hierarchy( checkpoint_ns, checkpoint_ns_to_state_snapshots ) return state_snapshot @@ -669,10 +682,18 @@ class Pregel( ) yield state_snapshot else: - nodes, channels = self._get_nodes_and_channels( - checkpoint_tuple.config["configurable"]["checkpoint_ns"] + nodes, channels = _get_nodes_and_channels( + self.nodes, + self.channels, + checkpoint_tuple.config["configurable"]["checkpoint_ns"], + ) + yield _prepare_state_snapshot( + checkpoint_tuple, + nodes, + channels, + self.managed_values_dict, + self.stream_channels_asis, ) - yield self._prepare_state_snapshot(checkpoint_tuple, nodes, channels) async def aget_state_history( self, @@ -709,11 +730,17 @@ class Pregel( ) yield state_snapshot else: - nodes, channels = self._get_nodes_and_channels( - checkpoint_tuple.config["configurable"]["checkpoint_ns"] + nodes, channels = _get_nodes_and_channels( + self.nodes, + self.channels, + checkpoint_tuple.config["configurable"]["checkpoint_ns"], ) - yield await self._prepare_state_snapshot_async( - checkpoint_tuple, nodes, channels + yield await _prepare_state_snapshot_async( + checkpoint_tuple, + nodes, + channels, + self.managed_values_dict, + self.stream_channels_asis, ) def update_state( From 6531ec7669767f505e584021d33e1720af17c509 Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 14 Aug 2024 14:02:29 -0400 Subject: [PATCH 19/41] remove inherit checkpointer --- libs/langgraph/langgraph/graph/graph.py | 5 +- libs/langgraph/langgraph/graph/state.py | 4 +- libs/langgraph/langgraph/pregel/__init__.py | 67 +++++++++------------ libs/langgraph/tests/test_pregel.py | 23 +++---- libs/langgraph/tests/test_pregel_async.py | 23 +++---- 5 files changed, 50 insertions(+), 72 deletions(-) diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index 6fc4eb8cd..96ccda3b2 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -24,6 +24,7 @@ from langchain_core.runnables.graph import Graph as DrawableGraph from langchain_core.runnables.graph import Node as DrawableNode from langgraph.channels.ephemeral_value import EphemeralValue +from langgraph.checkpoint.base import BaseCheckpointSaver from langgraph.constants import ( CHECKPOINT_NAMESPACE_SEPARATOR, END, @@ -33,7 +34,7 @@ from langgraph.constants import ( Send, ) from langgraph.errors import InvalidUpdateError -from langgraph.pregel import Channel, CheckpointerType, Pregel +from langgraph.pregel import Channel, Pregel from langgraph.pregel.read import PregelNode from langgraph.pregel.types import All from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry @@ -374,7 +375,7 @@ class Graph: def compile( self, - checkpointer: Optional[CheckpointerType] = None, + checkpointer: Optional[BaseCheckpointSaver] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, debug: bool = False, diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 8e7298616..f998b568a 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -29,6 +29,7 @@ from langgraph.channels.dynamic_barrier_value import DynamicBarrierValue, WaitFo 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, SEND_CHECKPOINT_NAMESPACE_SEPARATOR, @@ -44,7 +45,6 @@ from langgraph.graph.graph import ( Send, ) from langgraph.managed.base import ManagedValue, is_managed_value -from langgraph.pregel import CheckpointerType from langgraph.pregel.read import ChannelRead, PregelNode from langgraph.pregel.types import All, RetryPolicy from langgraph.pregel.write import SKIP_WRITE, ChannelWrite, ChannelWriteEntry @@ -381,7 +381,7 @@ class StateGraph(Graph): def compile( self, - checkpointer: Optional[CheckpointerType] = None, + checkpointer: Optional[BaseCheckpointSaver] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, interrupt_after: Optional[Union[All, Sequence[str]]] = None, debug: bool = False, diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index ffa53da85..9f583b5b5 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -13,7 +13,6 @@ from typing import ( Callable, Dict, Iterator, - Literal, Mapping, Optional, Sequence, @@ -117,10 +116,6 @@ WriteValue = Union[ ] -INHERIT_CHECKPOINTER = "inherit_checkpointer" -CheckpointerType = Union[BaseCheckpointSaver, Literal["inherit_checkpointer"]] - - class Channel: @overload @classmethod @@ -334,6 +329,16 @@ async def _prepare_state_snapshot_async( ) +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]] ): @@ -363,7 +368,7 @@ class Pregel( debug: bool = Field(default_factory=get_debug) """Whether to print debug information during execution. Defaults to False.""" - checkpointer: Optional[CheckpointerType] = None + checkpointer: Optional[BaseCheckpointSaver] = None """Checkpointer used to save and load graph state. Defaults to None.""" retry_policy: Optional[RetryPolicy] = None @@ -420,7 +425,6 @@ class Pregel( + ( self.checkpointer.config_specs if self.checkpointer is not None - and self.checkpointer != INHERIT_CHECKPOINTER else [] ) + ( @@ -500,6 +504,18 @@ class Pregel( if is_managed_value(v) } + @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, *, include_subgraph_state: bool = False ) -> StateSnapshot: @@ -1009,44 +1025,19 @@ class Pregel( stream_mode = stream_mode if stream_mode is not None else self.stream_mode if not isinstance(stream_mode, list): stream_mode = [stream_mode] - - if config and config.get("configurable", {}).get(CONFIG_KEY_CHECKPOINTER): - parent_checkpointer = config["configurable"][CONFIG_KEY_CHECKPOINTER] - else: - parent_checkpointer = None - if config and config.get("configurable", {}).get(CONFIG_KEY_READ) is not None: # if being called as a node in another graph, always use values mode stream_mode = ["values"] - - if ( - (interrupt_before or interrupt_after) - and parent_checkpointer is not None - and self.checkpointer is None - ): - raise ValueError( - "Missing checkpointer for a subgraph with interrupts. " - "Please compile the subgraph graph with checkpointer=INHERIT_CHECKPOINTER (from langgraph.pregel import INHERIT_CHECKPOINTER)." - ) - - if ( - parent_checkpointer is not None - and self.checkpointer is not None - and self.checkpointer != INHERIT_CHECKPOINTER - ): - raise ValueError( - "Custom checkpointers for subgraphs are not allowed. " - "Please compile the subgraph graph with checkpointer=INHERIT_CHECKPOINTER (from langgraph.pregel import INHERIT_CHECKPOINTER)." - ) - if ( - parent_checkpointer is not None - and self.checkpointer == INHERIT_CHECKPOINTER + config is not None + and config.get("configurable", {}).get(CONFIG_KEY_CHECKPOINTER) + and (interrupt_before or interrupt_after or _has_nested_interrupts(self)) ): - checkpointer = parent_checkpointer + checkpointer: Optional[BaseCheckpointSaver] = config["configurable"][ + CONFIG_KEY_CHECKPOINTER + ] else: checkpointer = self.checkpointer - return ( debug, stream_mode, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index b64d97b57..aa1db664e 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -61,7 +61,6 @@ from langgraph.prebuilt.chat_agent_executor import ( ) from langgraph.prebuilt.tool_node import ToolNode from langgraph.pregel import ( - INHERIT_CHECKPOINTER, Channel, GraphRecursionError, Pregel, @@ -7745,7 +7744,7 @@ def test_nested_graph_interrupts( graph.add_node("outer_1", outer_1) graph.add_node( "inner", - inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER), + inner.compile(interrupt_before=["inner_2"]), ) graph.add_node("outer_2", outer_2) graph.set_entry_point("outer_1") @@ -8924,7 +8923,7 @@ def test_nested_graph_interrupts_parallel( graph = StateGraph(State) graph.add_node( "inner", - inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER), + inner.compile(interrupt_before=["inner_2"]), ) graph.add_node("outer_1", outer_1) graph.add_node("outer_2", outer_2) @@ -9045,9 +9044,7 @@ def test_doubly_nested_graph_interrupts( child = StateGraph(ChildState) child.add_node( "child_1", - grandchild.compile( - interrupt_before=["grandchild_2"], checkpointer=INHERIT_CHECKPOINTER - ), + grandchild.compile(interrupt_before=["grandchild_2"]), ) child.set_entry_point("child_1") child.set_finish_point("child_1") @@ -9060,7 +9057,7 @@ def test_doubly_nested_graph_interrupts( graph = StateGraph(State) graph.add_node("parent_1", parent_1) - graph.add_node("child", child.compile(checkpointer=INHERIT_CHECKPOINTER)) + 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") @@ -9154,7 +9151,7 @@ def test_nested_graph_state( graph.add_node("outer_1", outer_1) graph.add_node( "inner", - inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER), + inner.compile(interrupt_before=["inner_2"]), ) graph.add_node("outer_2", outer_2) graph.set_entry_point("outer_1") @@ -9632,9 +9629,7 @@ def test_doubly_nested_graph_state( child = StateGraph(ChildState) child.add_node( "child_1", - grandchild.compile( - interrupt_before=["grandchild_2"], checkpointer=INHERIT_CHECKPOINTER - ), + grandchild.compile(interrupt_before=["grandchild_2"]), ) child.set_entry_point("child_1") child.set_finish_point("child_1") @@ -9647,7 +9642,7 @@ def test_doubly_nested_graph_state( graph = StateGraph(State) graph.add_node("parent_1", parent_1) - graph.add_node("child", child.compile(checkpointer=INHERIT_CHECKPOINTER)) + 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") @@ -9917,9 +9912,7 @@ def test_send_to_nested_graphs( builder = StateGraph(OverallState) builder.add_node( "generate_joke", - subgraph.compile( - checkpointer=INHERIT_CHECKPOINTER, interrupt_before=["generate"] - ), + subgraph.compile(interrupt_before=["generate"]), ) builder.add_conditional_edges(START, continue_to_jokes) builder.add_edge("generate_joke", END) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 7b2e3222b..3c881e630 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -57,7 +57,6 @@ 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 ( - INHERIT_CHECKPOINTER, Channel, GraphRecursionError, Pregel, @@ -6240,7 +6239,7 @@ async def test_nested_graph_interrupts( graph.add_node("outer_1", outer_1) graph.add_node( "inner", - inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER), + inner.compile(interrupt_before=["inner_2"]), ) graph.add_node("outer_2", outer_2) graph.set_entry_point("outer_1") @@ -7426,7 +7425,7 @@ async def test_nested_graph_interrupts_parallel( graph = StateGraph(State) graph.add_node( "inner", - inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER), + inner.compile(interrupt_before=["inner_2"]), ) graph.add_node("outer_1", outer_1) graph.add_node("outer_2", outer_2) @@ -7549,9 +7548,7 @@ async def test_doubly_nested_graph_interrupts( child = StateGraph(ChildState) child.add_node( "child_1", - grandchild.compile( - interrupt_before=["grandchild_2"], checkpointer=INHERIT_CHECKPOINTER - ), + grandchild.compile(interrupt_before=["grandchild_2"]), ) child.set_entry_point("child_1") child.set_finish_point("child_1") @@ -7564,7 +7561,7 @@ async def test_doubly_nested_graph_interrupts( graph = StateGraph(State) graph.add_node("parent_1", parent_1) - graph.add_node("child", child.compile(checkpointer=INHERIT_CHECKPOINTER)) + 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") @@ -7661,7 +7658,7 @@ async def test_nested_graph_state( graph.add_node("outer_1", outer_1) graph.add_node( "inner", - inner.compile(interrupt_before=["inner_2"], checkpointer=INHERIT_CHECKPOINTER), + inner.compile(interrupt_before=["inner_2"]), ) graph.add_node("outer_2", outer_2) graph.set_entry_point("outer_1") @@ -8143,9 +8140,7 @@ async def test_doubly_nested_graph_state( child = StateGraph(ChildState) child.add_node( "child_1", - grandchild.compile( - interrupt_before=["grandchild_2"], checkpointer=INHERIT_CHECKPOINTER - ), + grandchild.compile(interrupt_before=["grandchild_2"]), ) child.set_entry_point("child_1") child.set_finish_point("child_1") @@ -8158,7 +8153,7 @@ async def test_doubly_nested_graph_state( graph = StateGraph(State) graph.add_node("parent_1", parent_1) - graph.add_node("child", child.compile(checkpointer=INHERIT_CHECKPOINTER)) + 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") @@ -8428,9 +8423,7 @@ async def test_send_to_nested_graphs( builder = StateGraph(OverallState) builder.add_node( "generate_joke", - subgraph.compile( - checkpointer=INHERIT_CHECKPOINTER, interrupt_before=["generate"] - ), + subgraph.compile(interrupt_before=["generate"]), ) builder.add_conditional_edges(START, continue_to_jokes) builder.add_edge("generate_joke", END) From 7fa97898aaacc0dd0155674800e76852c148f6de Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 14 Aug 2024 16:02:37 -0400 Subject: [PATCH 20/41] correctly propagate all subgraph attributes --- libs/langgraph/langgraph/pregel/__init__.py | 118 +++++++------------- libs/langgraph/tests/test_pregel.py | 72 ++++-------- libs/langgraph/tests/test_pregel_async.py | 74 +++++------- 3 files changed, 88 insertions(+), 176 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 9f583b5b5..5c89ed27f 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -190,15 +190,12 @@ class Channel: ) -def _get_nodes_and_channels( - nodes: Mapping[str, PregelNode], - channels: Mapping[str, BaseChannel], - checkpoint_ns: str, -) -> tuple[Mapping[str, PregelNode], Mapping[str, BaseChannel]]: +def _get_subgraph(graph: Pregel, checkpoint_ns: str) -> Pregel: if checkpoint_ns == "": - return nodes, channels + return graph path = checkpoint_ns.split(CHECKPOINT_NAMESPACE_SEPARATOR) + nodes = graph.nodes for subgraph_node_name in path: # if we have this separator it means we have a node that was triggered by Send if SEND_CHECKPOINT_NAMESPACE_SEPARATOR in subgraph_node_name: @@ -210,17 +207,17 @@ def _get_nodes_and_channels( if subgraph_node_name not in nodes: raise ValueError(f"Couldn't find node '{subgraph_node_name}'.") - subgraph_node = nodes[subgraph_node_name].get_node() - - if not isinstance(subgraph_node, RunnableSequence): - break - - first_step = subgraph_node.steps[0] - if isinstance(first_step, Pregel): - nodes = first_step.nodes - channels = first_step.channels - - return nodes, channels + subgraph_node = nodes[subgraph_node_name] + if isinstance(subgraph_node.bound, Pregel): + nodes = subgraph_node.bound.nodes + elif isinstance(subgraph_node.bound, RunnableSequence): + for runnable in subgraph_node.bound.steps: + if isinstance(runnable, Pregel): + nodes = runnable.nodes + break + else: + continue + return subgraph_node.bound def _assemble_state_snapshot_hierarchy( @@ -259,24 +256,21 @@ def _assemble_state_snapshot_hierarchy( def _prepare_state_snapshot( saved: CheckpointTuple, - nodes: Mapping[str, PregelNode], - channels: Mapping[str, BaseChannel], - managed_values_dict: dict[str, ManagedValueSpec], - select_channels: str | list[str], + graph: Pregel, ) -> StateSnapshot: with ChannelsManager( { k: LastValue(None) if isinstance(c, Context) else c - for k, c in channels.items() + for k, c in graph.channels.items() }, saved.checkpoint, saved.config, ) as channels, ManagedValuesManager( - managed_values_dict, ensure_config(saved.config) + graph.managed_values_dict, ensure_config(saved.config) ) as managed: next_tasks = prepare_next_tasks( saved.checkpoint, - nodes, + graph.nodes, channels, managed, saved.config, @@ -284,7 +278,7 @@ def _prepare_state_snapshot( for_execution=False, ) return StateSnapshot( - values=read_channels(channels, select_channels), + values=read_channels(channels, graph.stream_channels_asis), next=tuple(t.name for t in next_tasks), config=saved.config, metadata=saved.metadata, @@ -294,25 +288,21 @@ def _prepare_state_snapshot( async def _prepare_state_snapshot_async( - saved: CheckpointTuple, - nodes: Mapping[str, PregelNode], - channels: Mapping[str, BaseChannel], - managed_values_dict: dict[str, ManagedValueSpec], - select_channels: str | list[str], + saved: CheckpointTuple, graph: Pregel ) -> StateSnapshot: async with AsyncChannelsManager( { k: LastValue(None) if isinstance(c, Context) else c - for k, c in channels.items() + for k, c in graph.channels.items() }, saved.checkpoint, saved.config, ) as channels, AsyncManagedValuesManager( - managed_values_dict, ensure_config(saved.config) + graph.managed_values_dict, ensure_config(saved.config) ) as managed: next_tasks = prepare_next_tasks( saved.checkpoint, - nodes, + graph.nodes, channels, managed, saved.config, @@ -320,7 +310,7 @@ async def _prepare_state_snapshot_async( for_execution=False, ) return StateSnapshot( - values=read_channels(channels, select_channels), + values=read_channels(channels, graph.stream_channels_asis), next=tuple(t.name for t in next_tasks), config=saved.config, metadata=saved.metadata, @@ -534,9 +524,7 @@ class Pregel( 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_nodes_and_channels: dict[ - str, tuple[Mapping[str, PregelNode], Mapping[str, BaseChannel]] - ] = {} + checkpoint_ns_to_graph: dict[str, Pregel] = {} for checkpoint_tuple in checkpoint_tuples: saved_checkpoint_ns = checkpoint_tuple.config["configurable"][ "checkpoint_ns" @@ -555,22 +543,14 @@ class Pregel( existing_checkpoint_id is None or saved_checkpoint_id > existing_checkpoint_id ): - if saved_checkpoint_ns not in checkpoint_ns_to_nodes_and_channels: - checkpoint_ns_to_nodes_and_channels[ - saved_checkpoint_ns - ] = _get_nodes_and_channels( - self.nodes, self.channels, saved_checkpoint_ns + if saved_checkpoint_ns not in checkpoint_ns_to_graph: + checkpoint_ns_to_graph[saved_checkpoint_ns] = _get_subgraph( + self, saved_checkpoint_ns ) - nodes, channels = checkpoint_ns_to_nodes_and_channels[ - saved_checkpoint_ns - ] state_snapshot = _prepare_state_snapshot( checkpoint_tuple, - nodes, - channels, - self.managed_values_dict, - self.stream_channels_asis, + checkpoint_ns_to_graph[saved_checkpoint_ns], ) checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot checkpoint_ns_to_checkpoint_id[ @@ -610,9 +590,7 @@ class Pregel( 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_nodes_and_channels: dict[ - str, tuple[Mapping[str, PregelNode], Mapping[str, BaseChannel]] - ] = {} + checkpoint_ns_to_graph: dict[str, Pregel] = {} async for checkpoint_tuple in checkpoint_tuples: saved_checkpoint_ns = checkpoint_tuple.config["configurable"][ "checkpoint_ns" @@ -631,22 +609,14 @@ class Pregel( existing_checkpoint_id is None or saved_checkpoint_id > existing_checkpoint_id ): - if saved_checkpoint_ns not in checkpoint_ns_to_nodes_and_channels: - checkpoint_ns_to_nodes_and_channels[ - saved_checkpoint_ns - ] = _get_nodes_and_channels( - self.nodes, self.channels, saved_checkpoint_ns + if saved_checkpoint_ns not in checkpoint_ns_to_graph: + checkpoint_ns_to_graph[saved_checkpoint_ns] = _get_subgraph( + self, saved_checkpoint_ns ) - nodes, channels = checkpoint_ns_to_nodes_and_channels[ - saved_checkpoint_ns - ] state_snapshot = await _prepare_state_snapshot_async( checkpoint_tuple, - nodes, - channels, - self.managed_values_dict, - self.stream_channels_asis, + checkpoint_ns_to_graph[saved_checkpoint_ns], ) checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot checkpoint_ns_to_checkpoint_id[ @@ -698,17 +668,13 @@ class Pregel( ) yield state_snapshot else: - nodes, channels = _get_nodes_and_channels( - self.nodes, - self.channels, + graph = _get_subgraph( + self, checkpoint_tuple.config["configurable"]["checkpoint_ns"], ) yield _prepare_state_snapshot( checkpoint_tuple, - nodes, - channels, - self.managed_values_dict, - self.stream_channels_asis, + graph, ) async def aget_state_history( @@ -746,17 +712,13 @@ class Pregel( ) yield state_snapshot else: - nodes, channels = _get_nodes_and_channels( - self.nodes, - self.channels, + graph = _get_subgraph( + self, checkpoint_tuple.config["configurable"]["checkpoint_ns"], ) yield await _prepare_state_snapshot_async( checkpoint_tuple, - nodes, - channels, - self.managed_values_dict, - self.stream_channels_asis, + graph, ) def update_state( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index aa1db664e..7bdc357d7 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8602,7 +8602,7 @@ def test_nested_graph_interrupts( ] assert child_state_history == [ StateSnapshot( - values={"my_key": "hi my value here"}, + values={"my_key": "hi my value here", "my_other_key": "hi my value"}, next=("inner_2",), config={ "configurable": { @@ -9140,6 +9140,7 @@ def test_nested_graph_state( class State(TypedDict): my_key: str + other_parent_key: str def outer_1(state: State): return {"my_key": "hi " + state["my_key"]} @@ -9214,7 +9215,7 @@ def test_nested_graph_state( }, subgraph_state_snapshots={ "inner": StateSnapshot( - values={"my_key": "hi my value here"}, + values={"my_key": "hi my value here", "my_other_key": "hi my value"}, next=("inner_2",), config={ "configurable": { @@ -9271,7 +9272,10 @@ def test_nested_graph_state( }, subgraph_state_snapshots={ "inner": StateSnapshot( - values={"my_key": "hi my value here"}, + values={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, next=("inner_2",), config={ "configurable": { @@ -9376,7 +9380,10 @@ def test_nested_graph_state( {"configurable": {"thread_id": "1", "checkpoint_ns": "inner"}} ) assert child_snapshot == StateSnapshot( - values={"my_key": "hi my value here and there"}, + values={ + "my_key": "hi my value here and there", + "my_other_key": "hi my value here", + }, next=(), config={ "configurable": { @@ -9519,7 +9526,10 @@ def test_nested_graph_state( }, subgraph_state_snapshots={ "inner": StateSnapshot( - values={"my_key": "hi my value here and there"}, + values={ + "my_key": "hi my value here and there", + "my_other_key": "hi my value here", + }, next=(), config={ "configurable": { @@ -9931,6 +9941,13 @@ def test_send_to_nested_graphs( 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": []}, next=("generate_joke", "generate_joke"), @@ -9950,50 +9967,7 @@ def test_send_to_nested_graphs( "checkpoint_id": AnyStr(), } }, - subgraph_state_snapshots={ - subgraph_nodes[0]: StateSnapshot( - values={"jokes": []}, - next=("generate",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - subgraph_nodes[1]: StateSnapshot( - values={"jokes": []}, - next=("generate",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - }, + subgraph_state_snapshots=subgraph_state_snapshots, ) assert actual_snapshot == expected_snapshot diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 3c881e630..87d9240e0 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7104,7 +7104,7 @@ async def test_nested_graph_interrupts( ] assert child_state_history == [ StateSnapshot( - values={"my_key": "hi my value here"}, + values={"my_key": "hi my value here", "my_other_key": "hi my value"}, next=("inner_2",), config={ "configurable": { @@ -7647,6 +7647,7 @@ async def test_nested_graph_state( class State(TypedDict): my_key: str + other_parent_key: str async def outer_1(state: State): return {"my_key": "hi " + state["my_key"]} @@ -7721,7 +7722,10 @@ async def test_nested_graph_state( }, subgraph_state_snapshots={ "inner": StateSnapshot( - values={"my_key": "hi my value here"}, + values={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, next=("inner_2",), config={ "configurable": { @@ -7780,7 +7784,10 @@ async def test_nested_graph_state( }, subgraph_state_snapshots={ "inner": StateSnapshot( - values={"my_key": "hi my value here"}, + values={ + "my_key": "hi my value here", + "my_other_key": "hi my value", + }, next=("inner_2",), config={ "configurable": { @@ -7885,7 +7892,10 @@ async def test_nested_graph_state( {"configurable": {"thread_id": "1", "checkpoint_ns": "inner"}} ) assert child_snapshot == StateSnapshot( - values={"my_key": "hi my value here and there"}, + values={ + "my_key": "hi my value here and there", + "my_other_key": "hi my value here", + }, next=(), config={ "configurable": { @@ -8030,7 +8040,10 @@ async def test_nested_graph_state( }, subgraph_state_snapshots={ "inner": StateSnapshot( - values={"my_key": "hi my value here and there"}, + values={ + "my_key": "hi my value here and there", + "my_other_key": "hi my value here", + }, next=(), config={ "configurable": { @@ -8442,6 +8455,12 @@ async def test_send_to_nested_graphs( 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": []}, next=("generate_joke", "generate_joke"), @@ -8461,50 +8480,7 @@ async def test_send_to_nested_graphs( "checkpoint_id": AnyStr(), } }, - subgraph_state_snapshots={ - subgraph_nodes[0]: StateSnapshot( - values={"jokes": []}, - next=("generate",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[0], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - subgraph_nodes[1]: StateSnapshot( - values={"jokes": []}, - next=("generate",), - config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - metadata={"source": "loop", "writes": {"edit": None}, "step": 1}, - created_at=AnyStr(), - parent_config={ - "configurable": { - "thread_id": "1", - "checkpoint_ns": subgraph_nodes[1], - "checkpoint_id": AnyStr(), - } - }, - subgraph_state_snapshots=None, - ), - }, + subgraph_state_snapshots=subgraph_state_snapshots, ) assert actual_snapshot == expected_snapshot From 45054df71ae48d9d8a44a0427f824f9bac2f759e Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 14 Aug 2024 21:05:00 -0400 Subject: [PATCH 21/41] remove include_subgraph_state kwarg --- libs/langgraph/langgraph/pregel/__init__.py | 61 +-- libs/langgraph/tests/test_pregel.py | 489 ++++++++++++++++--- libs/langgraph/tests/test_pregel_async.py | 494 +++++++++++++++++--- 3 files changed, 862 insertions(+), 182 deletions(-) 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( From f51e7ea9a407d35e7853e1182878eeec1a997d5e Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 21 Aug 2024 15:38:30 -0400 Subject: [PATCH 22/41] pass pending writes in checkpointers --- .../langgraph/checkpoint/postgres/__init__.py | 1 + .../langgraph/checkpoint/postgres/aio.py | 1 + .../langgraph/checkpoint/sqlite/__init__.py | 16 +++++++++++++++- .../langgraph/checkpoint/sqlite/aio.py | 12 ++++++++++++ .../langgraph/checkpoint/memory/__init__.py | 4 ++++ 5 files changed, 33 insertions(+), 1 deletion(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py index 02d5880a0..1da25b35f 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/__init__.py @@ -150,6 +150,7 @@ class PostgresSaver(BasePostgresSaver): } if value["parent_checkpoint_id"] else None, + self._load_writes(value["pending_writes"]), ) def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py index 7ddb81237..7be1f36dc 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/aio.py @@ -135,6 +135,7 @@ class AsyncPostgresSaver(BasePostgresSaver): } if value["parent_checkpoint_id"] else None, + await asyncio.to_thread(self._load_writes, value["pending_writes"]), ) async def aget_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]: diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py index 9bf4139ad..d67454034 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/__init__.py @@ -318,7 +318,9 @@ class SqliteSaver(BaseCheckpointSaver): ORDER BY checkpoint_id DESC""" if limit: query += f" LIMIT {limit}" - with self.cursor(transaction=False) as cur: + with self.cursor(transaction=False) as cur, self.cursor( + transaction=False + ) as writes_cur: cur.execute(query, param_values) for ( thread_id, @@ -329,6 +331,14 @@ class SqliteSaver(BaseCheckpointSaver): checkpoint, metadata, ) in cur: + writes_cur.execute( + "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?", + ( + thread_id, + checkpoint_ns, + checkpoint_id, + ), + ) yield CheckpointTuple( { "configurable": { @@ -350,6 +360,10 @@ class SqliteSaver(BaseCheckpointSaver): if parent_checkpoint_id else None ), + [ + (task_id, channel, self.serde.loads_typed((type, value))) + for task_id, channel, type, value in writes_cur + ], ) def put( diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py index 56d14f613..364adf063 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/aio.py @@ -346,6 +346,14 @@ class AsyncSqliteSaver(BaseCheckpointSaver): checkpoint, metadata, ) in cursor: + writes_cur = await self.conn.execute( + "SELECT task_id, channel, type, value FROM writes WHERE thread_id = ? AND checkpoint_ns = ? AND checkpoint_id = ?", + ( + thread_id, + checkpoint_ns, + checkpoint_id, + ), + ) yield CheckpointTuple( { "configurable": { @@ -367,6 +375,10 @@ class AsyncSqliteSaver(BaseCheckpointSaver): if parent_checkpoint_id else None ), + [ + (task_id, channel, self.serde.loads_typed((type, value))) + async for task_id, channel, type, value in writes_cur + ], ) async def aput( diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 6cc1a3b14..6b3714f05 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -206,6 +206,7 @@ class MemorySaver( elif limit is not None: limit -= 1 + writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)] yield CheckpointTuple( config={ "configurable": { @@ -215,6 +216,9 @@ class MemorySaver( } }, checkpoint=self.serde.loads_typed(checkpoint), + pending_writes=[ + (id, c, self.serde.loads_typed(v)) for id, c, v in writes + ], metadata=metadata, parent_config={ "configurable": { From 0a87b9fa1c8e6385c088f665311a0fe0fb286e5f Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 21 Aug 2024 15:49:56 -0400 Subject: [PATCH 23/41] update more tests --- libs/langgraph/langgraph/pregel/__init__.py | 7 +- libs/langgraph/tests/test_pregel.py | 148 +++++++++++++--- libs/langgraph/tests/test_pregel_async.py | 183 ++++++++++++++++++-- 3 files changed, 296 insertions(+), 42 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 488334f63..cb8cb29fb 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -558,7 +558,12 @@ class Pregel( if not checkpoint_ns_to_state_snapshots: return StateSnapshot( - values={}, next=(), config=config, metadata=None, created_at=None, tasks=() + values={}, + next=(), + config=config, + metadata=None, + created_at=None, + tasks=(), ) state_snapshot = _assemble_state_snapshot_hierarchy( diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 8016f3911..6058ecf1d 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7986,7 +7986,13 @@ 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", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner",), config={ "configurable": { @@ -8014,7 +8020,13 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=(PregelTask(AnyStr(), "inner_2"),), + tasks=( + PregelTask( + AnyStr(), + "inner_2", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner_2",), config={ "configurable": { @@ -8288,7 +8300,13 @@ 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", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner",), config={ "configurable": { @@ -8357,7 +8375,13 @@ 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", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner",), config={ "configurable": { @@ -8385,7 +8409,13 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=(PregelTask(AnyStr(), "inner_2"),), + tasks=( + PregelTask( + AnyStr(), + "inner_2", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner_2",), config={ "configurable": { @@ -8635,7 +8665,13 @@ 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", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner",), config={ "configurable": { @@ -8663,7 +8699,13 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=(PregelTask(AnyStr(), "inner_2"),), + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner_2",), config={ "configurable": { @@ -9040,7 +9082,9 @@ def test_nested_graph_interrupts( assert state_history == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), + ), next=("inner",), config={ "configurable": { @@ -9068,7 +9112,11 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=(PregelTask(AnyStr(), "inner_2"),), + tasks=( + PregelTask( + AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) + ), + ), next=("inner_2",), config={ "configurable": { @@ -9152,7 +9200,9 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=(PregelTask(AnyStr(), "inner_2"),), + tasks=( + PregelTask(AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),)), + ), next=("inner_2",), config={ "configurable": { @@ -9201,7 +9251,9 @@ 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", interrupts=(Interrupt(when="before"),)), + ), next=("inner",), config={ "configurable": { @@ -9229,7 +9281,11 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=(PregelTask(AnyStr(), "inner_2"),), + tasks=( + PregelTask( + AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) + ), + ), next=("inner_2",), config={ "configurable": { @@ -9261,7 +9317,9 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), + ), next=("inner",), config={ "configurable": { @@ -9289,7 +9347,11 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=(PregelTask(AnyStr(), "inner_2"),), + tasks=( + PregelTask( + AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) + ), + ), next=("inner_2",), config={ "configurable": { @@ -9420,7 +9482,9 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), + ), next=("inner",), config={ "configurable": { @@ -9448,7 +9512,11 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=(PregelTask(AnyStr(), "inner_2"),), + tasks=( + PregelTask( + AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) + ), + ), next=("inner_2",), config={ "configurable": { @@ -9869,7 +9937,11 @@ def test_nested_graph_state( # 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"),), + tasks=( + PregelTask( + AnyStr(), "inner", interrupts=(Interrupt(when="before", value=None),) + ), + ), next=("inner",), config={ "configurable": { @@ -9894,7 +9966,14 @@ def test_nested_graph_state( subgraph_state_snapshots={ "inner": StateSnapshot( values={"my_key": "hi my value here", "my_other_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner_2"),), + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + error=None, + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner_2",), config={ "configurable": { @@ -9928,7 +10007,13 @@ def test_nested_graph_state( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask( + AnyStr(), + "inner", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner",), config={ "configurable": { @@ -9956,7 +10041,14 @@ def test_nested_graph_state( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=(PregelTask(AnyStr(), "inner_2"),), + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + error=None, + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner_2",), config={ "configurable": { @@ -10357,7 +10449,7 @@ def test_doubly_nested_graph_state( 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"),), + tasks=(PregelTask(AnyStr(), "child", interrupts=(Interrupt(when="before"),)),), next=("child",), config={ "configurable": { @@ -10382,7 +10474,11 @@ def test_doubly_nested_graph_state( subgraph_state_snapshots={ "child": StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "child_1"),), + tasks=( + PregelTask( + AnyStr(), "child_1", interrupts=(Interrupt(when="before"),) + ), + ), next=("child_1",), config={ "configurable": { @@ -10403,7 +10499,13 @@ def test_doubly_nested_graph_state( subgraph_state_snapshots={ "child_1": StateSnapshot( values={"my_key": "hi my value here"}, - tasks=(PregelTask(AnyStr(), "grandchild_2"),), + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + interrupts=(Interrupt(when="before"),), + ), + ), next=("grandchild_2",), config={ "configurable": { diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index eb6a1abe8..22178f103 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6456,7 +6456,13 @@ 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", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner",), config={ "configurable": { @@ -6484,6 +6490,13 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, + tasks=( + PregelTask( + AnyStr(), + "inner_2", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner_2",), config={ "configurable": { @@ -6642,6 +6655,7 @@ async def test_nested_graph_interrupts( "my_other_key": "hi my value here", }, next=(), + tasks=(), config={ "configurable": { "thread_id": "1", @@ -6761,7 +6775,13 @@ 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", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner",), config={ "configurable": { @@ -6830,7 +6850,13 @@ 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", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner",), config={ "configurable": { @@ -6858,6 +6884,13 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, + tasks=( + PregelTask( + AnyStr(), + "inner_2", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner_2",), config={ "configurable": { @@ -7020,6 +7053,7 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here and there", "my_other_key": "hi my value here", }, + tasks=(), next=(), config={ "configurable": { @@ -7109,7 +7143,13 @@ 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", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner",), config={ "configurable": { @@ -7137,6 +7177,13 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner_2",), config={ "configurable": { @@ -7270,6 +7317,7 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here and there", "my_other_key": "hi my value here", }, + tasks=(), next=(), config={ "configurable": { @@ -7429,6 +7477,7 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here and there", "my_other_key": "hi my value here", }, + tasks=(), next=(), config={ "configurable": { @@ -7510,7 +7559,9 @@ async def test_nested_graph_interrupts( assert state_history == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), + ), next=("inner",), config={ "configurable": { @@ -7538,6 +7589,11 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, + tasks=( + PregelTask( + AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) + ), + ), next=("inner_2",), config={ "configurable": { @@ -7618,9 +7674,14 @@ 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", interrupts=(Interrupt(when="before"),)), + ), + next=("inner_2",), config={ "configurable": { "thread_id": "6", @@ -7660,7 +7721,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" } @@ -7668,7 +7729,9 @@ 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", interrupts=(Interrupt(when="before"),)), + ), next=("inner",), config={ "configurable": { @@ -7696,6 +7759,11 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, + tasks=( + PregelTask( + AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) + ), + ), next=("inner_2",), config={ "configurable": { @@ -7727,7 +7795,9 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), + ), next=("inner",), config={ "configurable": { @@ -7755,6 +7825,11 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, + tasks=( + PregelTask( + AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) + ), + ), next=("inner_2",), config={ "configurable": { @@ -7825,9 +7900,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)] == [ @@ -7885,7 +7960,9 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, - tasks=(PregelTask(AnyStr(), "inner"),), + tasks=( + PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), + ), next=("inner",), config={ "configurable": { @@ -7913,6 +7990,11 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, + tasks=( + PregelTask( + AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) + ), + ), next=("inner_2",), config={ "configurable": { @@ -7973,6 +8055,7 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here and there", "my_other_key": "hi my value here", }, + tasks=(), next=(), config={ "configurable": { @@ -8337,6 +8420,11 @@ async def test_nested_graph_state( # 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", interrupts=(Interrupt(when="before", value=None),) + ), + ), next=("inner",), config={ "configurable": { @@ -8360,10 +8448,15 @@ async def test_nested_graph_state( }, subgraph_state_snapshots={ "inner": StateSnapshot( - values={ - "my_key": "hi my value here", - "my_other_key": "hi my value", - }, + values={"my_key": "hi my value here", "my_other_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + error=None, + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner_2",), config={ "configurable": { @@ -8397,6 +8490,13 @@ async def test_nested_graph_state( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), + "inner", + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner",), config={ "configurable": { @@ -8424,6 +8524,14 @@ async def test_nested_graph_state( "my_key": "hi my value here", "my_other_key": "hi my value", }, + tasks=( + PregelTask( + AnyStr(), + name="inner_2", + error=None, + interrupts=(Interrupt(when="before", value=None),), + ), + ), next=("inner_2",), config={ "configurable": { @@ -8456,6 +8564,7 @@ async def test_nested_graph_state( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -8477,6 +8586,7 @@ async def test_nested_graph_state( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -8499,6 +8609,7 @@ async def test_nested_graph_state( # 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": { @@ -8532,6 +8643,7 @@ async def test_nested_graph_state( "my_key": "hi my value here and there", "my_other_key": "hi my value here", }, + tasks=(), next=(), config={ "configurable": { @@ -8571,6 +8683,7 @@ async def test_nested_graph_state( }, ) == StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8598,6 +8711,7 @@ async def test_nested_graph_state( assert [s async for s in app.aget_state_history(config)] == [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, + tasks=(), next=(), config={ "configurable": { @@ -8625,6 +8739,7 @@ async def test_nested_graph_state( ), StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(PregelTask(AnyStr(), "outer_2"),), next=("outer_2",), config={ "configurable": { @@ -8650,6 +8765,7 @@ async def test_nested_graph_state( ), StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8677,6 +8793,7 @@ async def test_nested_graph_state( "my_key": "hi my value here and there", "my_other_key": "hi my value here", }, + tasks=(), next=(), config={ "configurable": { @@ -8709,6 +8826,7 @@ async def test_nested_graph_state( ), StateSnapshot( values={"my_key": "my value"}, + tasks=(PregelTask(AnyStr(), "outer_1"),), next=("outer_1",), config={ "configurable": { @@ -8730,6 +8848,7 @@ async def test_nested_graph_state( ), StateSnapshot( values={}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { @@ -8813,6 +8932,7 @@ async def test_doubly_nested_graph_state( 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", interrupts=(Interrupt(when="before"),)),), next=("child",), config={ "configurable": { @@ -8837,6 +8957,11 @@ async def test_doubly_nested_graph_state( subgraph_state_snapshots={ "child": StateSnapshot( values={"my_key": "hi my value"}, + tasks=( + PregelTask( + AnyStr(), "child_1", interrupts=(Interrupt(when="before"),) + ), + ), next=("child_1",), config={ "configurable": { @@ -8857,6 +8982,13 @@ async def test_doubly_nested_graph_state( subgraph_state_snapshots={ "child_1": StateSnapshot( values={"my_key": "hi my value here"}, + tasks=( + PregelTask( + AnyStr(), + "grandchild_2", + interrupts=(Interrupt(when="before"),), + ), + ), next=("grandchild_2",), config={ "configurable": { @@ -8887,6 +9019,7 @@ async def test_doubly_nested_graph_state( 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": { @@ -8917,6 +9050,7 @@ async def test_doubly_nested_graph_state( ) assert grandchild_snapshot == StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(), next=(), config={ "configurable": { @@ -8946,6 +9080,7 @@ async def test_doubly_nested_graph_state( ) assert child_snapshot == StateSnapshot( values={"my_key": "hi my value here and there"}, + tasks=(), next=(), config={ "configurable": { @@ -8979,6 +9114,7 @@ async def test_doubly_nested_graph_state( }, ) == StateSnapshot( values={"my_key": "hi my value"}, + tasks=(PregelTask(AnyStr(), "child"),), next=("child",), config={ "configurable": { @@ -9069,6 +9205,10 @@ async def test_send_to_nested_graphs( } expected_snapshot = StateSnapshot( values={"subjects": ["cats", "dogs"], "jokes": []}, + tasks=( + PregelTask(AnyStr(), "generate_joke"), + PregelTask(AnyStr(), "generate_joke"), + ), next=("generate_joke", "generate_joke"), config={ "configurable": { @@ -9102,6 +9242,7 @@ async def test_send_to_nested_graphs( "subjects": ["cats", "dogs"], "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], }, + tasks=(), next=(), config={ "configurable": { @@ -9146,6 +9287,7 @@ async def test_send_to_nested_graphs( "subjects": ["cats", "dogs"], "jokes": ["Joke about cats - hohoho", "Joke about dogs - hohoho"], }, + tasks=(), next=(), config={ "configurable": { @@ -9177,6 +9319,10 @@ async def test_send_to_nested_graphs( 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", @@ -9197,6 +9343,7 @@ async def test_send_to_nested_graphs( ), StateSnapshot( values={"jokes": []}, + tasks=(PregelTask(AnyStr(), "__start__"),), next=("__start__",), config={ "configurable": { From e7bc74e9186ca4d74bdcb14c9bf426f1d1acb723 Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 21 Aug 2024 18:06:18 -0400 Subject: [PATCH 24/41] remove refactors --- libs/langgraph/langgraph/pregel/__init__.py | 163 +++++++++----------- 1 file changed, 70 insertions(+), 93 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index cb8cb29fb..5c54d5b7e 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -258,73 +258,6 @@ def _assemble_state_snapshot_hierarchy( return state_snapshot -def _prepare_state_snapshot( - saved: CheckpointTuple, - graph: Pregel, -) -> StateSnapshot: - with ChannelsManager( - { - k: LastValue(None) if isinstance(c, Context) else c - for k, c in graph.channels.items() - }, - saved.checkpoint, - saved.config, - ) as channels, ManagedValuesManager( - graph.managed_values_dict, ensure_config(saved.config) - ) as managed: - next_tasks = prepare_next_tasks( - saved.checkpoint, - graph.nodes, - channels, - managed, - saved.config, - saved.metadata.get("step", -1) + 1, - for_execution=False, - ) - return StateSnapshot( - values=read_channels(channels, graph.stream_channels_asis), - next=tuple(t.name for t in next_tasks), - config=saved.config, - metadata=saved.metadata, - created_at=saved.checkpoint["ts"], - parent_config=saved.parent_config, - tasks=tasks_w_writes(next_tasks, saved.pending_writes), - ) - - -async def _prepare_state_snapshot_async( - saved: CheckpointTuple, graph: Pregel -) -> StateSnapshot: - async with AsyncChannelsManager( - { - k: LastValue(None) if isinstance(c, Context) else c - for k, c in graph.channels.items() - }, - saved.checkpoint, - saved.config, - ) as channels, AsyncManagedValuesManager( - graph.managed_values_dict, ensure_config(saved.config) - ) as managed: - next_tasks = prepare_next_tasks( - saved.checkpoint, - graph.nodes, - channels, - managed, - saved.config, - saved.metadata.get("step", -1) + 1, - for_execution=False, - ) - return StateSnapshot( - values=read_channels(channels, graph.stream_channels_asis), - next=tuple(t.name for t in next_tasks), - config=saved.config, - metadata=saved.metadata, - created_at=saved.checkpoint["ts"], - parent_config=saved.parent_config, - tasks=tasks_w_writes(next_tasks, saved.pending_writes), - ) - - def _has_nested_interrupts( graph: Pregel, ) -> bool: @@ -517,20 +450,16 @@ class Pregel( if not self.checkpointer: raise ValueError("No checkpointer set") - checkpoint_tuple = self.checkpointer.get_tuple(config) - checkpoint_config = checkpoint_tuple.config if checkpoint_tuple else config + saved = self.checkpointer.get_tuple(config) + checkpoint_config = saved.config if saved 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 self.checkpointer.list(config): - saved_checkpoint_ns = checkpoint_tuple.config["configurable"][ - "checkpoint_ns" - ] - saved_checkpoint_id = checkpoint_tuple.config["configurable"][ - "checkpoint_id" - ] + for saved in self.checkpointer.list(config): + saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"] + saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"] if checkpoint_id != saved_checkpoint_id: continue @@ -547,10 +476,36 @@ class Pregel( self, saved_checkpoint_ns ) - state_snapshot = _prepare_state_snapshot( - checkpoint_tuple, - checkpoint_ns_to_graph[saved_checkpoint_ns], - ) + graph = checkpoint_ns_to_graph[saved_checkpoint_ns] + with ChannelsManager( + { + k: LastValue(None) if isinstance(c, Context) else c + for k, c in graph.channels.items() + }, + saved.checkpoint, + saved.config, + ) as channels, ManagedValuesManager( + graph.managed_values_dict, ensure_config(saved.config) + ) as 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 checkpoint_ns_to_checkpoint_id[ saved_checkpoint_ns @@ -576,20 +531,16 @@ class Pregel( if not self.checkpointer: raise ValueError("No checkpointer set") - checkpoint_tuple = await self.checkpointer.aget_tuple(config) - checkpoint_config = checkpoint_tuple.config if checkpoint_tuple else config + saved = await self.checkpointer.aget_tuple(config) + checkpoint_config = saved.config if saved 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 self.checkpointer.alist(config): - saved_checkpoint_ns = checkpoint_tuple.config["configurable"][ - "checkpoint_ns" - ] - saved_checkpoint_id = checkpoint_tuple.config["configurable"][ - "checkpoint_id" - ] + async for saved in self.checkpointer.alist(config): + saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"] + saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"] if checkpoint_id != saved_checkpoint_id: continue @@ -606,10 +557,36 @@ class Pregel( self, saved_checkpoint_ns ) - state_snapshot = await _prepare_state_snapshot_async( - checkpoint_tuple, - checkpoint_ns_to_graph[saved_checkpoint_ns], - ) + graph = checkpoint_ns_to_graph[saved_checkpoint_ns] + async with AsyncChannelsManager( + { + k: LastValue(None) if isinstance(c, Context) else c + for k, c in graph.channels.items() + }, + saved.checkpoint, + saved.config, + ) as channels, AsyncManagedValuesManager( + graph.managed_values_dict, ensure_config(saved.config) + ) as 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 checkpoint_ns_to_checkpoint_id[ saved_checkpoint_ns From 0b6088f91386d885cb661a5bbcf68218a253bc4c Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 21 Aug 2024 18:26:54 -0400 Subject: [PATCH 25/41] remove futures.clear --- libs/langgraph/langgraph/pregel/__init__.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index ed3ec185a..4ef53b40f 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1185,8 +1185,6 @@ class Pregel( else: loop.put_writes(task.id, [(ERROR, exc)]) - futures.clear() - else: # save task writes to checkpointer loop.put_writes(task.id, task.writes) @@ -1447,7 +1445,6 @@ class Pregel( else: loop.put_writes(task.id, [(ERROR, exc)]) - futures.clear() else: # save task writes to checkpointer loop.put_writes(task.id, task.writes) From 6c7d9c35bc34a18e132f96aa181fa8acbc822f6f Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 22 Aug 2024 15:34:26 -0400 Subject: [PATCH 26/41] remove interrupts --- libs/langgraph/tests/test_pregel.py | 69 ++++------------------- libs/langgraph/tests/test_pregel_async.py | 61 ++++++-------------- 2 files changed, 30 insertions(+), 100 deletions(-) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 2474f6790..df6ecf8ec 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -8179,7 +8179,6 @@ def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner",), @@ -8213,7 +8212,6 @@ def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner_2", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner_2",), @@ -8493,7 +8491,6 @@ def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner",), @@ -8568,7 +8565,6 @@ def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner",), @@ -8602,7 +8598,6 @@ def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner_2", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner_2",), @@ -8858,7 +8853,6 @@ def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner",), @@ -8892,7 +8886,6 @@ def test_nested_graph_interrupts( PregelTask( AnyStr(), name="inner_2", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner_2",), @@ -9271,9 +9264,7 @@ def test_nested_graph_interrupts( assert state_history == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=( - PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), - ), + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -9301,11 +9292,7 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=( - PregelTask( - AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) - ), - ), + tasks=(PregelTask(AnyStr(), "inner_2"),), next=("inner_2",), config={ "configurable": { @@ -9389,9 +9376,7 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=( - PregelTask(AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),)), - ), + tasks=(PregelTask(AnyStr(), "inner_2"),), next=("inner_2",), config={ "configurable": { @@ -9440,9 +9425,7 @@ def test_nested_graph_interrupts( assert list(app.get_state_history(config)) == [ StateSnapshot( values={"my_key": "hi my value"}, - tasks=( - PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), - ), + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -9470,11 +9453,7 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=( - PregelTask( - AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) - ), - ), + tasks=(PregelTask(AnyStr(), "inner_2"),), next=("inner_2",), config={ "configurable": { @@ -9506,9 +9485,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, - tasks=( - PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), - ), + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -9536,11 +9513,7 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=( - PregelTask( - AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) - ), - ), + tasks=(PregelTask(AnyStr(), "inner_2"),), next=("inner_2",), config={ "configurable": { @@ -9671,9 +9644,7 @@ def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, - tasks=( - PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), - ), + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -9701,11 +9672,7 @@ def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=( - PregelTask( - AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) - ), - ), + tasks=(PregelTask(AnyStr(), "inner_2"),), next=("inner_2",), config={ "configurable": { @@ -10126,11 +10093,7 @@ def test_nested_graph_state( # 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", interrupts=(Interrupt(when="before", value=None),) - ), - ), + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -10160,7 +10123,6 @@ def test_nested_graph_state( AnyStr(), name="inner_2", error=None, - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner_2",), @@ -10200,7 +10162,6 @@ def test_nested_graph_state( PregelTask( AnyStr(), "inner", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner",), @@ -10235,7 +10196,6 @@ def test_nested_graph_state( AnyStr(), name="inner_2", error=None, - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner_2",), @@ -10638,7 +10598,7 @@ def test_doubly_nested_graph_state( 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", interrupts=(Interrupt(when="before"),)),), + tasks=(PregelTask(AnyStr(), "child"),), next=("child",), config={ "configurable": { @@ -10663,11 +10623,7 @@ def test_doubly_nested_graph_state( subgraph_state_snapshots={ "child": StateSnapshot( values={"my_key": "hi my value"}, - tasks=( - PregelTask( - AnyStr(), "child_1", interrupts=(Interrupt(when="before"),) - ), - ), + tasks=(PregelTask(AnyStr(), "child_1"),), next=("child_1",), config={ "configurable": { @@ -10692,7 +10648,6 @@ def test_doubly_nested_graph_state( PregelTask( AnyStr(), "grandchild_2", - interrupts=(Interrupt(when="before"),), ), ), next=("grandchild_2",), diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 409448a9d..0d6f21bf7 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -6801,7 +6801,6 @@ async def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner",), @@ -6835,7 +6834,6 @@ async def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner_2", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner_2",), @@ -7120,7 +7118,6 @@ async def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner",), @@ -7195,7 +7192,6 @@ async def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner",), @@ -7229,7 +7225,6 @@ async def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner_2", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner_2",), @@ -7488,7 +7483,6 @@ async def test_nested_graph_interrupts( PregelTask( AnyStr(), "inner", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner",), @@ -7522,7 +7516,6 @@ async def test_nested_graph_interrupts( PregelTask( AnyStr(), name="inner_2", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner_2",), @@ -7901,7 +7894,10 @@ async def test_nested_graph_interrupts( StateSnapshot( values={"my_key": "hi my value"}, tasks=( - PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), + PregelTask( + AnyStr(), + "inner", + ), ), next=("inner",), config={ @@ -7932,7 +7928,8 @@ async def test_nested_graph_interrupts( }, tasks=( PregelTask( - AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) + AnyStr(), + "inner_2", ), ), next=("inner_2",), @@ -8019,9 +8016,7 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=( - PregelTask(AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),)), - ), + tasks=(PregelTask(AnyStr(), "inner_2"),), next=("inner_2",), config={ "configurable": { @@ -8070,9 +8065,7 @@ 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", interrupts=(Interrupt(when="before"),)), - ), + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8100,11 +8093,7 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=( - PregelTask( - AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) - ), - ), + tasks=(PregelTask(AnyStr(), "inner_2"),), next=("inner_2",), config={ "configurable": { @@ -8136,9 +8125,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, - tasks=( - PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), - ), + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8166,11 +8153,7 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=( - PregelTask( - AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) - ), - ), + tasks=(PregelTask(AnyStr(), "inner_2"),), next=("inner_2",), config={ "configurable": { @@ -8301,9 +8284,7 @@ async def test_nested_graph_interrupts( ), StateSnapshot( values={"my_key": "hi my value"}, - tasks=( - PregelTask(AnyStr(), "inner", interrupts=(Interrupt(when="before"),)), - ), + tasks=(PregelTask(AnyStr(), "inner"),), next=("inner",), config={ "configurable": { @@ -8331,11 +8312,7 @@ async def test_nested_graph_interrupts( "my_key": "hi my value here", "my_other_key": "hi my value", }, - tasks=( - PregelTask( - AnyStr(), "inner_2", interrupts=(Interrupt(when="before"),) - ), - ), + tasks=(PregelTask(AnyStr(), "inner_2"),), next=("inner_2",), config={ "configurable": { @@ -8763,7 +8740,8 @@ async def test_nested_graph_state( values={"my_key": "hi my value"}, tasks=( PregelTask( - AnyStr(), "inner", interrupts=(Interrupt(when="before", value=None),) + AnyStr(), + "inner", ), ), next=("inner",), @@ -8795,7 +8773,6 @@ async def test_nested_graph_state( AnyStr(), name="inner_2", error=None, - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner_2",), @@ -8835,7 +8812,6 @@ async def test_nested_graph_state( PregelTask( AnyStr(), "inner", - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner",), @@ -8870,7 +8846,6 @@ async def test_nested_graph_state( AnyStr(), name="inner_2", error=None, - interrupts=(Interrupt(when="before", value=None),), ), ), next=("inner_2",), @@ -9273,7 +9248,7 @@ async def test_doubly_nested_graph_state( 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", interrupts=(Interrupt(when="before"),)),), + tasks=(PregelTask(AnyStr(), "child"),), next=("child",), config={ "configurable": { @@ -9300,7 +9275,8 @@ async def test_doubly_nested_graph_state( values={"my_key": "hi my value"}, tasks=( PregelTask( - AnyStr(), "child_1", interrupts=(Interrupt(when="before"),) + AnyStr(), + "child_1", ), ), next=("child_1",), @@ -9327,7 +9303,6 @@ async def test_doubly_nested_graph_state( PregelTask( AnyStr(), "grandchild_2", - interrupts=(Interrupt(when="before"),), ), ), next=("grandchild_2",), From acd8acf23708065026b6d9e6d447fdedbec38fac Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 22 Aug 2024 16:47:12 -0400 Subject: [PATCH 27/41] lint --- libs/checkpoint/langgraph/checkpoint/memory/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 132b05c15..5d73f4ba5 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -210,7 +210,9 @@ class MemorySaver( elif limit is not None: limit -= 1 - writes = self.writes[(thread_id, checkpoint_ns, checkpoint_id)].values() + writes = self.writes[ + (thread_id, checkpoint_ns, checkpoint_id) + ].values() yield CheckpointTuple( config={ From 4935cf52bf841e216511cb631365d872910b1511 Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 22 Aug 2024 16:52:16 -0400 Subject: [PATCH 28/41] lint --- libs/langgraph/langgraph/pregel/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 0096f7f8b..5ee96dde5 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -67,7 +67,6 @@ from langgraph.constants import ( ERROR, INTERRUPT, SEND_CHECKPOINT_NAMESPACE_SEPARATOR, - Interrupt, ) from langgraph.errors import GraphInterrupt, GraphRecursionError, InvalidUpdateError from langgraph.managed.base import ManagedValueSpec From 72893d9abb9a813bd882d84f9f64cbe89227a938 Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 22 Aug 2024 17:33:17 -0400 Subject: [PATCH 29/41] code review --- libs/langgraph/langgraph/pregel/__init__.py | 57 ++++----------------- libs/langgraph/langgraph/pregel/types.py | 2 +- libs/langgraph/langgraph/pregel/utils.py | 36 +++++++++++++ 3 files changed, 47 insertions(+), 48 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 5ee96dde5..7f15d01a4 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -88,7 +88,10 @@ from langgraph.pregel.types import ( StateSnapshot, StreamMode, ) -from langgraph.pregel.utils import get_new_channel_versions +from langgraph.pregel.utils import ( + assemble_state_snapshot_hierarchy, + 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 @@ -205,40 +208,6 @@ def _get_subgraph(graph: Pregel, checkpoint_ns: str) -> Pregel: return subgraph_node.bound -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 - - def _has_nested_interrupts( graph: Pregel, ) -> bool: @@ -497,7 +466,7 @@ class Pregel( tasks=(), ) - state_snapshot = _assemble_state_snapshot_hierarchy( + state_snapshot = assemble_state_snapshot_hierarchy( checkpoint_ns, checkpoint_ns_to_state_snapshots ) return state_snapshot @@ -573,7 +542,7 @@ class Pregel( tasks=(), ) - state_snapshot = _assemble_state_snapshot_hierarchy( + state_snapshot = assemble_state_snapshot_hierarchy( checkpoint_ns, checkpoint_ns_to_state_snapshots ) return state_snapshot @@ -1113,7 +1082,8 @@ class Pregel( ) if not done: break # timed out - for fut, task in zip(done, [futures.pop(fut) for fut in done]): + for fut in done: + task = futures.pop(fut) if exc := _exception(fut): # save error to checkpointer if isinstance(exc, GraphInterrupt): @@ -1364,7 +1334,8 @@ class Pregel( if not done: break # timed out - for fut, task in zip(done, [futures.pop(fut) for fut in done]): + for fut in done: + task = futures.pop(fut) if exc := _exception(fut): # save error to checkpointer if isinstance(exc, GraphInterrupt): @@ -1581,11 +1552,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/types.py b/libs/langgraph/langgraph/pregel/types.py index a8d66cb20..d4881452b 100644 --- a/libs/langgraph/langgraph/pregel/types.py +++ b/libs/langgraph/langgraph/pregel/types.py @@ -93,7 +93,7 @@ class StateSnapshot(NamedTuple): 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 thread ID suffix to snapshot.""" + """State snapshots of subgraphs represented as a mapping from checkpoint namespace (`checkpoint_ns`) to snapshot.""" All = Literal["*"] diff --git a/libs/langgraph/langgraph/pregel/utils.py b/libs/langgraph/langgraph/pregel/utils.py index d3d0d989f..9c76e5f7e 100644 --- a/libs/langgraph/langgraph/pregel/utils.py +++ b/libs/langgraph/langgraph/pregel/utils.py @@ -1,4 +1,6 @@ from langgraph.checkpoint.base import ChannelVersions +from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR +from langgraph.pregel.types import StateSnapshot def get_new_channel_versions( @@ -17,3 +19,37 @@ def get_new_channel_versions( new_versions = current_versions return new_versions + + +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 From 9f6e57d2a7e7c8ff506ecf9f296002e75cfc785a Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 22 Aug 2024 17:36:24 -0400 Subject: [PATCH 30/41] more code review --- libs/langgraph/langgraph/pregel/__init__.py | 2 +- libs/langgraph/langgraph/pregel/algo.py | 19 ++--------- libs/langgraph/langgraph/pregel/get_state.py | 36 ++++++++++++++++++++ libs/langgraph/langgraph/pregel/utils.py | 36 -------------------- 4 files changed, 39 insertions(+), 54 deletions(-) create mode 100644 libs/langgraph/langgraph/pregel/get_state.py diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 7f15d01a4..477f5a6d1 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -77,6 +77,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 @@ -89,7 +90,6 @@ from langgraph.pregel.types import ( StreamMode, ) from langgraph.pregel.utils import ( - assemble_state_snapshot_hierarchy, get_new_channel_versions, ) from langgraph.pregel.validate import validate_graph, validate_keys diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 27777bbab..e3e699518 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -283,9 +283,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))) @@ -293,21 +293,6 @@ def prepare_next_tasks( if for_execution: proc = processes[packet.node] if node := proc.get_node(): - triggers = [TASKS] - metadata = { - "langgraph_step": step, - "langgraph_node": packet.node, - "langgraph_triggers": triggers, - "langgraph_task_idx": len(tasks), - } - checkpoint_ns = ( - f"{parent_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{packet.node}:{packet.id}" - if parent_ns - else f"{packet.node}:{packet.id}" - ) - task_id = str( - uuid5(UUID(checkpoint["id"]), json.dumps((checkpoint_ns, metadata))) - ) writes = deque() tasks.append( PregelExecutableTask( 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/utils.py b/libs/langgraph/langgraph/pregel/utils.py index 9c76e5f7e..d3d0d989f 100644 --- a/libs/langgraph/langgraph/pregel/utils.py +++ b/libs/langgraph/langgraph/pregel/utils.py @@ -1,6 +1,4 @@ from langgraph.checkpoint.base import ChannelVersions -from langgraph.constants import CHECKPOINT_NAMESPACE_SEPARATOR -from langgraph.pregel.types import StateSnapshot def get_new_channel_versions( @@ -19,37 +17,3 @@ def get_new_channel_versions( new_versions = current_versions return new_versions - - -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 From 4162be81196c8468b8177b08ca590b6faacd62e3 Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 22 Aug 2024 20:03:36 -0400 Subject: [PATCH 31/41] optimize subgraph state lookups --- .../langgraph/checkpoint/postgres/base.py | 3 + .../langgraph/checkpoint/sqlite/utils.py | 3 + .../langgraph/checkpoint/memory/__init__.py | 7 +- libs/langgraph/langgraph/pregel/__init__.py | 218 +++++++++--------- 4 files changed, 116 insertions(+), 115 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index 4b665fa64..ff2ec6681 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -255,6 +255,9 @@ class BasePostgresSaver(BaseCheckpointSaver): if config: wheres.append("thread_id = %s ") param_values.append(config["configurable"]["thread_id"]) + 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-sqlite/langgraph/checkpoint/sqlite/utils.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py index 56034ea34..26b8594d6 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py @@ -70,6 +70,9 @@ def search_where( if config is not None: wheres.append("thread_id = ?") param_values.append(config["configurable"]["thread_id"]) + 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/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 5d73f4ba5..0a84326d7 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -177,6 +177,7 @@ class MemorySaver( Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. """ thread_ids = (config["configurable"]["thread_id"],) if config else self.storage + config_checkpoint_id = get_checkpoint_id(config) if config else None for thread_id in thread_ids: for checkpoint_ns in self.storage[thread_id].keys(): for checkpoint_id, ( @@ -188,7 +189,11 @@ class MemorySaver( key=lambda x: x[0], reverse=True, ): - # filter by checkpoint ID + # filter by checkpoint ID from config + if config_checkpoint_id and checkpoint_id != config_checkpoint_id: + continue + + # filter by checkpoint ID from `before` config if ( before and (before_checkpoint_id := get_checkpoint_id(before)) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 477f5a6d1..039d3bbef 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -178,34 +178,31 @@ class Channel: ) -def _get_subgraph(graph: Pregel, checkpoint_ns: str) -> Pregel: - if checkpoint_ns == "": - return graph - - path = checkpoint_ns.split(CHECKPOINT_NAMESPACE_SEPARATOR) - nodes = graph.nodes - for subgraph_node_name in path: - # if we have this separator it means we have a node that was triggered by Send - if SEND_CHECKPOINT_NAMESPACE_SEPARATOR in subgraph_node_name: - name_parts = subgraph_node_name.split(SEND_CHECKPOINT_NAMESPACE_SEPARATOR) - if len(name_parts) != 2: - raise ValueError(f"Malformed node name '{subgraph_node_name}'") - - subgraph_node_name = name_parts[0] - if subgraph_node_name not in nodes: - raise ValueError(f"Couldn't find node '{subgraph_node_name}'.") - - subgraph_node = nodes[subgraph_node_name] - if isinstance(subgraph_node.bound, Pregel): - nodes = subgraph_node.bound.nodes - elif isinstance(subgraph_node.bound, RunnableSequence): - for runnable in subgraph_node.bound.steps: +def _get_checkpoint_ns_to_graph( + graph: Pregel, checkpoint_ns_to_graph: dict[str, Pregel] = {}, checkpoint_ns="" +) -> Pregel: + for node_name, node in graph.nodes.items(): + if isinstance(node.bound, Pregel): + _get_checkpoint_ns_to_graph( + node.bound, + checkpoint_ns_to_graph, + f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}" + if checkpoint_ns + else node_name, + ) + elif isinstance(node.bound, RunnableSequence): + for runnable in node.bound.steps: if isinstance(runnable, Pregel): - nodes = runnable.nodes - break - else: - continue - return subgraph_node.bound + _get_checkpoint_ns_to_graph( + node.bound, + checkpoint_ns_to_graph, + f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}" + if checkpoint_ns + else node_name, + ) + + checkpoint_ns_to_graph[checkpoint_ns] = graph + return checkpoint_ns_to_graph def _has_nested_interrupts( @@ -401,59 +398,54 @@ class Pregel( saved = self.checkpointer.get_tuple(config) checkpoint_config = saved.config if saved 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 saved in self.checkpointer.list(config): + checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self) + + # 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) + + for saved in checkpoint_tuples: saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"] saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"] - if checkpoint_id != saved_checkpoint_id: + + 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 - 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 + with ChannelsManager( + graph.channels, saved.checkpoint, saved.config, skip_context=True + ) as ( + channels, + managed, ): - if saved_checkpoint_ns not in checkpoint_ns_to_graph: - checkpoint_ns_to_graph[saved_checkpoint_ns] = _get_subgraph( - self, saved_checkpoint_ns - ) - - graph = checkpoint_ns_to_graph[saved_checkpoint_ns] - with ChannelsManager( - graph.channels, saved.checkpoint, saved.config, skip_context=True - ) as ( + next_tasks = prepare_next_tasks( + saved.checkpoint, + graph.nodes, 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), - ) + 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 - checkpoint_ns_to_checkpoint_id[ - saved_checkpoint_ns - ] = saved_checkpoint_id + checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot + checkpoint_ns_to_checkpoint_id[saved_checkpoint_ns] = saved_checkpoint_id if not checkpoint_ns_to_state_snapshots: return StateSnapshot( @@ -479,57 +471,55 @@ class Pregel( saved = await self.checkpointer.aget_tuple(config) checkpoint_config = saved.config if saved 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 saved in self.checkpointer.alist(config): + checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self) + + # we only lookup subgraph checkpoints if we actually have subgraphs + if len(set(checkpoint_ns_to_graph)) == 1: + + async def list_checkpoints(): + yield saved + + checkpoint_tuples = list_checkpoints() + else: + checkpoint_tuples = self.checkpointer.alist(saved.config) + + 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 != saved_checkpoint_id: + + 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 - existing_checkpoint_id = checkpoint_ns_to_checkpoint_id.get( - saved_checkpoint_ns - ) + 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), + ) - # keep only most recent checkpoint_id - if ( - existing_checkpoint_id is None - or saved_checkpoint_id > existing_checkpoint_id - ): - if saved_checkpoint_ns not in checkpoint_ns_to_graph: - checkpoint_ns_to_graph[saved_checkpoint_ns] = _get_subgraph( - self, saved_checkpoint_ns - ) - - graph = checkpoint_ns_to_graph[saved_checkpoint_ns] - 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 - checkpoint_ns_to_checkpoint_id[ - saved_checkpoint_ns - ] = saved_checkpoint_id + checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot + checkpoint_ns_to_checkpoint_id[saved_checkpoint_ns] = saved_checkpoint_id if not checkpoint_ns_to_state_snapshots: return StateSnapshot( From 065055e5870fd1120eed6a98dfb893ad5108421c Mon Sep 17 00:00:00 2001 From: vbarda Date: Thu, 22 Aug 2024 20:37:48 -0400 Subject: [PATCH 32/41] small change --- libs/langgraph/langgraph/pregel/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 039d3bbef..2ee00bb29 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -194,7 +194,7 @@ def _get_checkpoint_ns_to_graph( for runnable in node.bound.steps: if isinstance(runnable, Pregel): _get_checkpoint_ns_to_graph( - node.bound, + runnable, checkpoint_ns_to_graph, f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}" if checkpoint_ns From 1333d8b4781ecad144b220c6427dbe1f8fb27d50 Mon Sep 17 00:00:00 2001 From: vbarda Date: Fri, 23 Aug 2024 14:26:49 -0400 Subject: [PATCH 34/41] cleanup --- libs/langgraph/langgraph/pregel/__init__.py | 24 +++++++++++---------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 2ee00bb29..f6ce5aefb 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -179,26 +179,28 @@ class Channel: def _get_checkpoint_ns_to_graph( - graph: Pregel, checkpoint_ns_to_graph: dict[str, Pregel] = {}, checkpoint_ns="" + graph: Pregel, + checkpoint_ns_to_graph: Optional[dict[str, Pregel]] = None, + checkpoint_ns: str = "", ) -> Pregel: + if checkpoint_ns_to_graph is None: + checkpoint_ns_to_graph = {} + 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, - f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}" - if checkpoint_ns - else node_name, + node.bound, checkpoint_ns_to_graph, new_checkpoint_ns ) 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, - f"{checkpoint_ns}{CHECKPOINT_NAMESPACE_SEPARATOR}{node_name}" - if checkpoint_ns - else node_name, + runnable, checkpoint_ns_to_graph, new_checkpoint_ns ) checkpoint_ns_to_graph[checkpoint_ns] = graph From 1f29925034d788b3b8b8b8f26bcd0eadee3da286 Mon Sep 17 00:00:00 2001 From: vbarda Date: Fri, 23 Aug 2024 14:44:42 -0400 Subject: [PATCH 35/41] add max recursion depth --- libs/langgraph/langgraph/pregel/__init__.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index f6ce5aefb..33bae3330 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -182,10 +182,16 @@ 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( + f"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}" @@ -194,13 +200,16 @@ def _get_checkpoint_ns_to_graph( ) if isinstance(node.bound, Pregel): _get_checkpoint_ns_to_graph( - node.bound, checkpoint_ns_to_graph, new_checkpoint_ns + 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 + runnable, + checkpoint_ns_to_graph, + new_checkpoint_ns, + max_depth - 1, ) checkpoint_ns_to_graph[checkpoint_ns] = graph From 71442916e56247199ff6e15bf6b435d1fb5a30cd Mon Sep 17 00:00:00 2001 From: vbarda Date: Fri, 23 Aug 2024 14:52:31 -0400 Subject: [PATCH 36/41] lint --- libs/langgraph/langgraph/pregel/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 33bae3330..79f7afa05 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -189,7 +189,7 @@ def _get_checkpoint_ns_to_graph( if max_depth <= 0: raise RecursionError( - f"Reached maximum recursion depth while building checkpoint NS -> graph mapping." + "Reached maximum recursion depth while building checkpoint NS -> graph mapping." ) for node_name, node in graph.nodes.items(): From 15692acef9527525b5ad9c997c76fdba50645938 Mon Sep 17 00:00:00 2001 From: vbarda Date: Fri, 23 Aug 2024 17:56:28 -0400 Subject: [PATCH 37/41] refactor to remove nested DB calls --- libs/langgraph/langgraph/pregel/__init__.py | 228 ++++++++++++++------ 1 file changed, 160 insertions(+), 68 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 79f7afa05..20410b21a 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -54,6 +54,7 @@ from langgraph.channels.base import ( from langgraph.channels.context import Context from langgraph.checkpoint.base import ( BaseCheckpointSaver, + CheckpointTuple, copy_checkpoint, create_checkpoint, empty_checkpoint, @@ -216,6 +217,131 @@ def _get_checkpoint_ns_to_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_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 saved_checkpoint_id != config["configurable"]["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_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 saved_checkpoint_id != config["configurable"]["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: @@ -409,7 +535,6 @@ class Pregel( saved = self.checkpointer.get_tuple(config) checkpoint_config = saved.config if saved else config checkpoint_ns = checkpoint_config["configurable"].get("checkpoint_ns", "") - checkpoint_ns_to_checkpoint_id: dict[str, str] = {} checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self) @@ -421,7 +546,6 @@ class Pregel( for saved in checkpoint_tuples: saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"] - saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"] graph_checkpoint_ns = saved_checkpoint_ns.split( SEND_CHECKPOINT_NAMESPACE_SEPARATOR @@ -456,7 +580,6 @@ class Pregel( ) checkpoint_ns_to_state_snapshots[saved_checkpoint_ns] = state_snapshot - checkpoint_ns_to_checkpoint_id[saved_checkpoint_ns] = saved_checkpoint_id if not checkpoint_ns_to_state_snapshots: return StateSnapshot( @@ -481,72 +604,21 @@ class Pregel( saved = await self.checkpointer.aget_tuple(config) checkpoint_config = saved.config if saved else config - checkpoint_ns = checkpoint_config["configurable"].get("checkpoint_ns", "") - checkpoint_ns_to_checkpoint_id: dict[str, str] = {} - checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self) # we only lookup subgraph checkpoints if we actually have subgraphs if len(set(checkpoint_ns_to_graph)) == 1: - async def list_checkpoints(): + async def alist_checkpoints(): yield saved - checkpoint_tuples = list_checkpoints() + checkpoint_tuples = alist_checkpoints() else: checkpoint_tuples = self.checkpointer.alist(saved.config) - async for saved in checkpoint_tuples: - saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"] - saved_checkpoint_id = saved.config["configurable"]["checkpoint_id"] - - 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 - checkpoint_ns_to_checkpoint_id[saved_checkpoint_ns] = saved_checkpoint_id - - 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 await _prepare_state_snapshot_async( + checkpoint_config, checkpoint_ns_to_graph, checkpoint_tuples ) - return state_snapshot def get_state_history( self, @@ -566,17 +638,24 @@ class Pregel( raise ValueError("Checkpointer does not support filtering") checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - for checkpoint_tuple in self.checkpointer.list( - config, before=before, limit=limit, filter=filter - ): + 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( + config, before=before, limit=limit, filter=filter + ) + ] + for checkpoint_tuple in checkpoint_tuples: if ( checkpoint_tuple.config["configurable"]["checkpoint_ns"] != checkpoint_ns ): - # only list root checkpoints here continue - state_snapshot = self.get_state(checkpoint_tuple.config) + state_snapshot = _prepare_state_snapshot( + checkpoint_tuple.config, checkpoint_ns_to_graph, iter(checkpoint_tuples) + ) yield state_snapshot async def aget_state_history( @@ -597,17 +676,30 @@ class Pregel( raise ValueError("Checkpointer does not support filtering") checkpoint_ns = config["configurable"].get("checkpoint_ns", "") - async for checkpoint_tuple in self.checkpointer.alist( - config, before=before, limit=limit, filter=filter - ): + 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( + 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 ): - # only list root checkpoints here continue - state_snapshot = await self.aget_state(checkpoint_tuple.config) + state_snapshot = await _prepare_state_snapshot_async( + checkpoint_tuple.config, checkpoint_ns_to_graph, alist_checkpoints() + ) yield state_snapshot def update_state( From cb30f686428b4c95d85a408525106e8fd785d811 Mon Sep 17 00:00:00 2001 From: vbarda Date: Fri, 23 Aug 2024 18:01:35 -0400 Subject: [PATCH 38/41] remove more reused code --- libs/langgraph/langgraph/pregel/__init__.py | 55 +-------------------- 1 file changed, 2 insertions(+), 53 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 20410b21a..2173cc7d5 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -534,8 +534,6 @@ class Pregel( saved = self.checkpointer.get_tuple(config) checkpoint_config = saved.config if saved else config - checkpoint_ns = checkpoint_config["configurable"].get("checkpoint_ns", "") - checkpoint_ns_to_state_snapshots: dict[str, StateSnapshot] = {} checkpoint_ns_to_graph: dict[str, Pregel] = _get_checkpoint_ns_to_graph(self) # we only lookup subgraph checkpoints if we actually have subgraphs @@ -544,58 +542,9 @@ class Pregel( else: checkpoint_tuples = self.checkpointer.list(saved.config) - for saved in checkpoint_tuples: - saved_checkpoint_ns = saved.config["configurable"]["checkpoint_ns"] - - 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 _prepare_state_snapshot( + checkpoint_config, checkpoint_ns_to_graph, checkpoint_tuples ) - return state_snapshot async def aget_state(self, config: RunnableConfig) -> StateSnapshot: """Get the current state of the graph.""" From 904a1a34716a9b4c4e78468c5266b1d135b8bd72 Mon Sep 17 00:00:00 2001 From: vbarda Date: Fri, 23 Aug 2024 18:09:10 -0400 Subject: [PATCH 39/41] extra paranoia --- libs/langgraph/langgraph/pregel/__init__.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 2173cc7d5..496ec1b30 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -223,11 +223,12 @@ def _prepare_state_snapshot( 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 saved_checkpoint_id != config["configurable"]["checkpoint_id"]: + if checkpoint_id and saved_checkpoint_id != checkpoint_id: continue graph_checkpoint_ns = saved_checkpoint_ns.split( @@ -287,11 +288,12 @@ async def _prepare_state_snapshot_async( 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 saved_checkpoint_id != config["configurable"]["checkpoint_id"]: + if checkpoint_id and saved_checkpoint_id != checkpoint_id: continue graph_checkpoint_ns = saved_checkpoint_ns.split( From 85e698e20b535a2a783874847a9ca80c73f79123 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 26 Aug 2024 19:17:59 -0400 Subject: [PATCH 40/41] filter on checkpoint NS --- .../langgraph/checkpoint/postgres/base.py | 4 + .../langgraph/checkpoint/sqlite/utils.py | 4 + .../langgraph/checkpoint/memory/__init__.py | 4 + libs/langgraph/tests/test_pregel.py | 210 ++++++------------ libs/langgraph/tests/test_pregel_async.py | 209 ++++++----------- 5 files changed, 145 insertions(+), 286 deletions(-) diff --git a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py index ff2ec6681..91b49a162 100644 --- a/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/checkpoint/postgres/base.py @@ -255,6 +255,10 @@ class BasePostgresSaver(BaseCheckpointSaver): if config: wheres.append("thread_id = %s ") param_values.append(config["configurable"]["thread_id"]) + 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) diff --git a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py index 26b8594d6..0e1e06fcc 100644 --- a/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py +++ b/libs/checkpoint-sqlite/langgraph/checkpoint/sqlite/utils.py @@ -70,6 +70,10 @@ def search_where( if config is not None: wheres.append("thread_id = ?") param_values.append(config["configurable"]["thread_id"]) + 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) diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index 0a84326d7..dee64a327 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -177,9 +177,13 @@ class MemorySaver( Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. """ thread_ids = (config["configurable"]["thread_id"],) if config else self.storage + 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_ns in self.storage[thread_id].keys(): + if config_checkpoint_ns and checkpoint_ns != config_checkpoint_ns: + continue + for checkpoint_id, ( checkpoint, metadata_b, diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a7642851d..685eb7c87 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -10442,81 +10442,9 @@ def test_nested_graph_state( } }, ) - # test loading inner snapshot - child_snapshot = app.get_state( - {"configurable": {"thread_id": "1", "checkpoint_ns": "inner"}} - ) - assert child_snapshot == 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, - ) - # test looking up parent state by checkpoint ID - assert app.get_state( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], - } - }, - ) == 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": child_snapshot}, - ) # test full history at the end - assert list(app.get_state_history(config)) == [ + actual_history = list(app.get_state_history(config)) + expected_history = [ StateSnapshot( values={"my_key": "hi my value here and there and back again"}, tasks=(), @@ -10675,6 +10603,10 @@ def test_nested_graph_state( 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( @@ -10847,75 +10779,11 @@ def test_doubly_nested_graph_state( } }, ) + + # test getting snapshot by ID + config = list(app.get_state_history(config))[2].config # test getting grandchild snapshot - grandchild_snapshot = app.get_state( - {"configurable": {"thread_id": "1", "checkpoint_ns": "child|child_1"}} - ) - assert grandchild_snapshot == 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, - ) - # test getting child snapshot - child_snapshot = app.get_state( - {"configurable": {"thread_id": "1", "checkpoint_ns": "child"}}, - ) - assert child_snapshot == 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": grandchild_snapshot}, - ) - # test getting parent snapshot for a checkpoint ID - assert app.get_state( - { - "configurable": { - "thread_id": "1", - "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], - } - }, - ) == StateSnapshot( + assert app.get_state(config) == StateSnapshot( values={"my_key": "hi my value"}, tasks=(PregelTask(AnyStr(), "child"),), next=("child",), @@ -10939,7 +10807,63 @@ def test_doubly_nested_graph_state( "checkpoint_id": AnyStr(), } }, - subgraph_state_snapshots={"child": child_snapshot}, + 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, + ) + }, + ) + }, ) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 426bfd730..50aafae86 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -8946,81 +8946,9 @@ async def test_nested_graph_state( } }, ) - # test loading inner snapshot - child_snapshot = await app.aget_state( - {"configurable": {"thread_id": "1", "checkpoint_ns": "inner"}} - ) - assert child_snapshot == 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, - ) - # test looking up parent state by checkpoint ID - assert await app.aget_state( - { - "configurable": { - "thread_id": "1", - "checkpoint_ns": "", - "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], - } - }, - ) == 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": child_snapshot}, - ) # test full history at the end - assert [s async for s in app.aget_state_history(config)] == [ + 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=(), @@ -9179,6 +9107,10 @@ async def test_nested_graph_state( 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( @@ -9356,75 +9288,10 @@ async def test_doubly_nested_graph_state( } }, ) + # test getting snapshot by ID + config = [s async for s in app.aget_state_history(config)][2].config # test getting grandchild snapshot - grandchild_snapshot = await app.aget_state( - {"configurable": {"thread_id": "1", "checkpoint_ns": "child|child_1"}} - ) - assert grandchild_snapshot == 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, - ) - # test getting child snapshot - child_snapshot = await app.aget_state( - {"configurable": {"thread_id": "1", "checkpoint_ns": "child"}}, - ) - assert child_snapshot == 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": grandchild_snapshot}, - ) - # test getting parent snapshot for a checkpoint ID - assert await app.aget_state( - { - "configurable": { - "thread_id": "1", - "checkpoint_id": child_snapshot.config["configurable"]["checkpoint_id"], - } - }, - ) == StateSnapshot( + assert await app.aget_state(config) == StateSnapshot( values={"my_key": "hi my value"}, tasks=(PregelTask(AnyStr(), "child"),), next=("child",), @@ -9448,7 +9315,63 @@ async def test_doubly_nested_graph_state( "checkpoint_id": AnyStr(), } }, - subgraph_state_snapshots={"child": child_snapshot}, + 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, + ) + }, + ) + }, ) From 507930e5b71f371d463da6d1b68d2e1934640147 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 26 Aug 2024 19:29:05 -0400 Subject: [PATCH 41/41] lint --- libs/checkpoint/langgraph/checkpoint/memory/__init__.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py index dee64a327..6918de87a 100644 --- a/libs/checkpoint/langgraph/checkpoint/memory/__init__.py +++ b/libs/checkpoint/langgraph/checkpoint/memory/__init__.py @@ -177,7 +177,9 @@ class MemorySaver( Iterator[CheckpointTuple]: An iterator of matching checkpoint tuples. """ thread_ids = (config["configurable"]["thread_id"],) if config else self.storage - config_checkpoint_ns = config["configurable"].get("checkpoint_ns") if config else None + 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_ns in self.storage[thread_id].keys():