From b43ef6440f2cc6a2f5801131d2aa783dc3c6244a Mon Sep 17 00:00:00 2001 From: vbarda Date: Tue, 23 Jul 2024 20:38:20 -0400 Subject: [PATCH] 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.