From dcb7278c5678db08e84b6ec2593cc9b7bdc2a7ec Mon Sep 17 00:00:00 2001 From: Andrew Nguonly Date: Mon, 13 May 2024 15:55:38 -0700 Subject: [PATCH] Add search and asearch APIs to BaseCheckpointerSaver class. Implement search and asearch in MemorySaver. --- langgraph/checkpoint/base.py | 24 ++++++++ langgraph/checkpoint/memory.py | 73 ++++++++++++++++++++++ tests/checkpoint/__init__.py | 0 tests/checkpoint/test_memory.py | 106 ++++++++++++++++++++++++++++++++ 4 files changed, 203 insertions(+) create mode 100644 tests/checkpoint/__init__.py create mode 100644 tests/checkpoint/test_memory.py diff --git a/langgraph/checkpoint/base.py b/langgraph/checkpoint/base.py index 695eb3d96..8306511bc 100644 --- a/langgraph/checkpoint/base.py +++ b/langgraph/checkpoint/base.py @@ -36,6 +36,11 @@ class CheckpointMetadata(TypedDict, total=False): Mapping from node name to writes emitted by that node. """ + score: Optional[int] + """The score of the checkpoint. + + The score can be used to mark a checkpoint as "good". + """ class Checkpoint(TypedDict): @@ -148,6 +153,15 @@ class BaseCheckpointSaver(ABC): ) -> Iterator[CheckpointTuple]: raise NotImplementedError + def search( + self, + metadata: CheckpointMetadata, + *, + before: Optional[RunnableConfig] = None, + limit: Optional[int] = None, + ) -> Iterator[CheckpointTuple]: + raise NotImplementedError + def put( self, config: RunnableConfig, @@ -173,6 +187,16 @@ class BaseCheckpointSaver(ABC): raise NotImplementedError yield + def asearch( + self, + metadata: CheckpointMetadata, + *, + before: Optional[RunnableConfig] = None, + limit: Optional[int] = None, + ) -> AsyncIterator[CheckpointTuple]: + raise NotImplementedError + yield + async def aput( self, config: RunnableConfig, diff --git a/langgraph/checkpoint/memory.py b/langgraph/checkpoint/memory.py index 04f94afd5..f1751835d 100644 --- a/langgraph/checkpoint/memory.py +++ b/langgraph/checkpoint/memory.py @@ -119,6 +119,55 @@ class MemorySaver(BaseCheckpointSaver): metadata=self.serde.loads(metadata), ) + def search( + self, + metadata_query: CheckpointMetadata, + *, + before: Optional[RunnableConfig] = None, + limit: Optional[int] = None, + ) -> Iterator[CheckpointTuple]: + """Search for checkpoints by metadata. + + This method retrieves a list of checkpoint tuples from the in-memory + storage based on the provided metadata query. The metadata query does + not need to contain all keys defined in the CheckpointMetadata class. + The checkpoints are ordered by timestamp in descending order. + + Args: + metadata_query (CheckpointMetadata): The metadata query to use for searching the checkpoints. + before (Optional[RunnableConfig]): If provided, only checkpoints before the specified timestamp are returned. Defaults to None. + limit (Optional[int]): The maximum number of checkpoints to return. Defaults to None. + + Yields: + Iterator[CheckpointTuple]: An iterator of checkpoint tuples. + """ + for thread_id, checkpoints in self.storage.items(): + for ts, (checkpoint_bytes, metadata_bytes) in checkpoints.items(): + # filter by thread_ts + if before and ts >= before["configurable"]["thread_ts"]: + continue + + # check if all query key/value pairs match the metadata + metadata = self.serde.loads(metadata_bytes) + all_keys_match = all( + query_value == metadata[query_key] + for query_key, query_value in metadata_query.items() + ) + + # if all query key/value pairs match, yield the checkpoint + if all_keys_match: + # limit search results + if limit is not None: + if limit <= 0: + break + limit -= 1 + + yield CheckpointTuple( + config={"configurable": {"thread_id": thread_id, "thread_ts": ts}}, + checkpoint=self.serde.loads(checkpoint_bytes), + metadata=metadata, + ) + def put( self, config: RunnableConfig, @@ -188,6 +237,30 @@ class MemorySaver(BaseCheckpointSaver): except StopIteration: return + async def asearch( + self, + metadata_query: CheckpointMetadata, + ) -> AsyncIterator[CheckpointTuple]: + """Asynchronous version of search. + + This method is an asynchronous wrapper around search that runs the synchronous + method in a separate thread using asyncio. + """ + loop = asyncio.get_running_loop() + iter = await loop.run_in_executor(None, self.search, metadata_query) + + def next_item(iter: Iterator[CheckpointTuple]) -> CheckpointTuple: + try: + return next(iter) + except StopIteration: + return None + + while True: + result = await loop.run_in_executor(None, next_item, iter) + if result is None: + break + yield result + async def aput( self, config: RunnableConfig, diff --git a/tests/checkpoint/__init__.py b/tests/checkpoint/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/checkpoint/test_memory.py b/tests/checkpoint/test_memory.py new file mode 100644 index 000000000..9b5bf5660 --- /dev/null +++ b/tests/checkpoint/test_memory.py @@ -0,0 +1,106 @@ +import pytest +from typing import AsyncIterator + +from langchain_core.runnables import RunnableConfig + +from langgraph.checkpoint.base import Checkpoint, CheckpointMetadata +from langgraph.checkpoint.memory import MemorySaver + + +class TestMemorySaver: + @pytest.fixture(autouse=True) + def setup(self): + self.memory_saver = MemorySaver() + + # objects for test setup + self.config_1: RunnableConfig = {"configurable": {"thread_id": "thread-1", "thread_ts": "1"}} + self.config_2: RunnableConfig = {"configurable": {"thread_id": "thread-2", "thread_ts": "2"}} + + self.chkpnt_1: Checkpoint = { + "v": 1, + "ts": "1", + "channel_values": {}, + "channel_versions": {}, + "versions_seen": {} + } + self.chkpnt_2: Checkpoint = { + "v": 2, + "ts": "2", + "channel_values": {}, + "channel_versions": {}, + "versions_seen": {} + } + + self.metadata_1: CheckpointMetadata = { + "source": "input", + "step": 2, + "writes": {}, + "score": 1, + } + self.metadata_2: CheckpointMetadata = { + "source": "loop", + "step": 1, + "writes": {"foo": "bar"}, + "score": None, + } + + async def _async_iterator_to_list(self, async_iterator: AsyncIterator): + result = [] + async for item in async_iterator: + result.append(item) + return result + + async def test_search(self): + # set up test + # 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) + + # call method / assertions + query_1: CheckpointMetadata = {"source": "input"} # search by 1 key + query_2: CheckpointMetadata = {"step": 1, "writes": {"foo": "bar"}} # search by multiple keys + query_3: CheckpointMetadata = {} # search by no keys, return all checkpoints + query_4: CheckpointMetadata = {"source": "update", "step": 1} # no match + + search_results_1 = list(self.memory_saver.search(query_1)) + assert len(search_results_1) == 1 + assert search_results_1[0].metadata == self.metadata_1 + + search_results_2 = list(self.memory_saver.search(query_2)) + assert len(search_results_2) == 1 + assert search_results_2[0].metadata == self.metadata_2 + + search_results_3 = list(self.memory_saver.search(query_3)) + assert len(search_results_3) == 2 + + search_results_4 = list(self.memory_saver.search(query_4)) + assert len(search_results_4) == 0 + + # TODO: test before and limit params + + async def test_asearch(self): + # set up test + # 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) + + # call method / assertions + query_1: CheckpointMetadata = {"source": "input"} # search by 1 key + query_2: CheckpointMetadata = {"step": 1, "writes": {"foo": "bar"}} # search by multiple keys + query_3: CheckpointMetadata = {} # search by no keys, return all checkpoints + query_4: CheckpointMetadata = {"source": "update", "step": 1} # no match + + + search_results_1 = [c async for c in self.memory_saver.asearch(query_1)] + assert len(search_results_1) == 1 + assert search_results_1[0].metadata == self.metadata_1 + + search_results_2 = [c async for c in self.memory_saver.asearch(query_2)] + assert len(search_results_2) == 1 + assert search_results_2[0].metadata == self.metadata_2 + + search_results_3 = [c async for c in self.memory_saver.asearch(query_3)] + assert len(search_results_3) == 2 + + search_results_4 = [c async for c in self.memory_saver.asearch(query_4)] + assert len(search_results_4) == 0