From 503304458721134eb927a3dc057d0061e76c8e37 Mon Sep 17 00:00:00 2001 From: vbarda Date: Mon, 12 Aug 2024 16:12:47 -0400 Subject: [PATCH] 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)