From 0c37a1717263dfa5ccc2d5868f4b16054a87402c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 14 May 2024 10:37:01 -0700 Subject: [PATCH] Fix memory checkpointer list/alist - add tests for memory checkpointer list and alist - fix async step counting --- langgraph/checkpoint/memory.py | 26 ++++++--- langgraph/pregel/__init__.py | 1 - tests/test_pregel.py | 100 ++++++++++++++++++++++++++++++++ tests/test_pregel_async.py | 102 +++++++++++++++++++++++++++++++++ 4 files changed, 221 insertions(+), 8 deletions(-) diff --git a/langgraph/checkpoint/memory.py b/langgraph/checkpoint/memory.py index 04f94afd5..ec8316e99 100644 --- a/langgraph/checkpoint/memory.py +++ b/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, diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 9549c2bad..4927b7596 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -1066,7 +1066,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( diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 4529e595c..379b127c6 100644 --- a/tests/test_pregel.py +++ b/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) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index c2a52ab11..32ba8a849 100644 --- a/tests/test_pregel_async.py +++ b/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)