From 85e698e20b535a2a783874847a9ca80c73f79123 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 26 Aug 2024 19:17:59 -0400 Subject: [PATCH] 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, + ) + }, + ) + }, )