Merge pull request #447 from langchain-ai/nc/14may/small-fixes

Fix memory checkpointer list/alist

langgraph/checkpoint/memory.py

@@ -1,5 +1,6 @@
 import asyncio
 from collections import defaultdict
+from functools import partial
 from typing import AsyncIterator, Iterator, Optional
 
 from langchain_core.runnables import RunnableConfig
@@ -112,7 +113,8 @@ class MemorySaver(BaseCheckpointSaver):
                 continue
             if limit is not None and limit <= 0:
                 break
-            limit -= 1
+            elif limit is not None:
+                limit -= 1
             yield CheckpointTuple(
                 config={"configurable": {"thread_id": thread_id, "thread_ts": ts}},
                 checkpoint=self.serde.loads(checkpoint),
@@ -168,7 +170,13 @@ class MemorySaver(BaseCheckpointSaver):
             None, self.get_tuple, config
         )
 
-    async def alist(self, config: RunnableConfig) -> AsyncIterator[CheckpointTuple]:
+    async def alist(
+        self,
+        config: RunnableConfig,
+        *,
+        before: Optional[RunnableConfig] = None,
+        limit: Optional[int] = None,
+    ) -> AsyncIterator[CheckpointTuple]:
         """Asynchronous version of list.
 
         This method is an asynchronous wrapper around list that runs the synchronous
@@ -181,12 +189,16 @@ class MemorySaver(BaseCheckpointSaver):
             AsyncIterator[CheckpointTuple]: An asynchronous iterator of checkpoint tuples.
         """
         loop = asyncio.get_running_loop()
-        iter = loop.run_in_executor(None, self.list, config)
+        iter = await loop.run_in_executor(
+            None, partial(self.list, before=before, limit=limit), config
+        )
         while True:
-            try:
-                yield await loop.run_in_executor(None, next, iter)
-            except StopIteration:
-                return
+            # handling StopIteration exception inside coroutine won't work
+            # as expected, so using next() with default value to break the loop
+            if item := await loop.run_in_executor(None, next, iter, None):
+                yield item
+            else:
+                break
 
     async def aput(
         self,

langgraph/pregel/__init__.py

@@ -1069,7 +1069,6 @@ class Pregel(
                 # channel updates from step N are only visible in step N+1,
                 # channels are guaranteed to be immutable for the duration of the step,
                 # channel updates being applied only at the transition between steps
-                start = saved.metadata.get("step", -1) + 1 if saved else 0
                 stop = start + config["recursion_limit"] + 1
                 for step in range(start, stop):
                     next_checkpoint, next_tasks = _prepare_next_tasks(

tests/test_pregel.py

@@ -351,6 +351,106 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None:
     snapshot = app.get_state({"configurable": {"thread_id": 2}})
     assert snapshot.next == ()
 
+    # list history
+    assert [c for c in app.get_state_history({"configurable": {"thread_id": 1}})] == [
+        StateSnapshot(
+            values={"input": 2},
+            next=("one",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "input", "step": -1, "writes": 2},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 3, "input": 2},
+            next=("two",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "loop", "step": 0, "writes": None},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 3, "output": 4, "input": 2},
+            next=(),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "loop", "step": 1, "writes": 4},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 3, "output": 4, "input": 20},
+            next=("one",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "input", "step": 2, "writes": 20},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 21, "output": 4, "input": 20},
+            next=("two",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "loop", "step": 3, "writes": None},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 21, "output": 4, "input": 3},
+            next=("one",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "input", "step": 4, "writes": 3},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 4, "output": 4, "input": 3},
+            next=("two",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "loop", "step": 5, "writes": None},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 4, "output": 5, "input": 3},
+            next=(),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "loop", "step": 6, "writes": 5},
+            parent_config=None,
+        ),
+    ]
+
 
 def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)

tests/test_pregel_async.py

@@ -329,6 +329,108 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N
     snapshot = await app.aget_state({"configurable": {"thread_id": 2}})
     assert snapshot.next == ()
 
+    # list history
+    assert [
+        c async for c in app.aget_state_history({"configurable": {"thread_id": 1}})
+    ] == [
+        StateSnapshot(
+            values={"input": 2},
+            next=("one",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "input", "step": -1, "writes": 2},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 3, "input": 2},
+            next=("two",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "loop", "step": 0, "writes": None},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 3, "output": 4, "input": 2},
+            next=(),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "loop", "step": 1, "writes": 4},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 3, "output": 4, "input": 20},
+            next=("one",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "input", "step": 2, "writes": 20},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 21, "output": 4, "input": 20},
+            next=("two",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "loop", "step": 3, "writes": None},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 21, "output": 4, "input": 3},
+            next=("one",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "input", "step": 4, "writes": 3},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 4, "output": 4, "input": 3},
+            next=("two",),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "loop", "step": 5, "writes": None},
+            parent_config=None,
+        ),
+        StateSnapshot(
+            values={"inbox": 4, "output": 5, "input": 3},
+            next=(),
+            config={
+                "configurable": {
+                    "thread_id": 1,
+                    "thread_ts": AnyStr(),
+                }
+            },
+            metadata={"source": "loop", "step": 6, "writes": 5},
+            parent_config=None,
+        ),
+    ]
+
 
 async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)