From a7d48465dafdc328e79f591edb956ad5e1a11bfd Mon Sep 17 00:00:00 2001 From: vbarda Date: Wed, 24 Jul 2024 12:44:35 -0400 Subject: [PATCH] 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)