Add search and asearch APIs to BaseCheckpointerSaver class. Implement search and asearch in MemorySaver.

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,

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,

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