From 56a38ee85964b834613854af01e7ade1021d7af2 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Tue, 14 Nov 2023 11:17:39 +0000
Subject: [PATCH 1/6] Change Channel.checkpoint() to return any python value,
 add checkpoint tests

---
 permchain/channels/base.py       |  9 +++--
 permchain/channels/binop.py      | 11 +++---
 permchain/channels/context.py    |  4 +-
 permchain/channels/last_value.py | 11 +++---
 permchain/channels/topic.py      | 19 +++++----
 tests/test_channels.py           | 68 ++++++++++++++++++++++++++++++++
 6 files changed, 97 insertions(+), 25 deletions(-)

diff --git a/permchain/channels/base.py b/permchain/channels/base.py
index 7a32a76d4..ee53993f5 100644
--- a/permchain/channels/base.py
+++ b/permchain/channels/base.py
@@ -15,6 +15,7 @@ from typing_extensions import Self
 
 Value = TypeVar("Value")
 Update = TypeVar("Update")
+Checkpoint = TypeVar("Checkpoint")
 
 
 class EmptyChannelError(Exception):
@@ -30,7 +31,7 @@ class InvalidUpdateError(Exception):
     pass
 
 
-class BaseChannel(Generic[Value, Update], ABC):
+class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
     @property
     @abstractmethod
     def ValueType(self) -> Any:
@@ -43,7 +44,9 @@ class BaseChannel(Generic[Value, Update], ABC):
 
     @contextmanager
     @abstractmethod
-    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+    def empty(
+        self, checkpoint: Optional[Checkpoint] = None
+    ) -> Generator[Self, None, None]:
         """Return a new identical channel, optionally initialized from a checkpoint."""
 
     @asynccontextmanager
@@ -68,7 +71,7 @@ class BaseChannel(Generic[Value, Update], ABC):
         Raises EmptyChannelError if the channel is empty (never updated yet)."""
 
     @abstractmethod
-    def checkpoint(self) -> str | None:
+    def checkpoint(self) -> Checkpoint | None:
         """Return a string representation of the channel's current state,
         or None if the channel doesn't support checkpoints.
 
diff --git a/permchain/channels/binop.py b/permchain/channels/binop.py
index 981322b46..322f99c3e 100644
--- a/permchain/channels/binop.py
+++ b/permchain/channels/binop.py
@@ -1,4 +1,3 @@
-import json
 from contextlib import contextmanager
 from typing import Callable, Generator, Generic, Optional, Sequence, Type
 
@@ -7,7 +6,7 @@ from typing_extensions import Self
 from permchain.channels.base import BaseChannel, EmptyChannelError, Value
 
 
-class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]):
+class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
     """Stores the result of applying a binary operator to the current value and each new value.
 
     ```python
@@ -32,10 +31,10 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]):
         return self.typ
 
     @contextmanager
-    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+    def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
         empty = self.__class__(self.typ, self.operator)
         if checkpoint is not None:
-            empty.value = json.loads(checkpoint)
+            empty.value = checkpoint
         try:
             yield empty
         finally:
@@ -60,8 +59,8 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value]):
         except AttributeError:
             raise EmptyChannelError()
 
-    def checkpoint(self) -> str:
+    def checkpoint(self) -> Value:
         try:
-            return json.dumps(self.value)
+            return self.value
         except AttributeError:
             raise EmptyChannelError()
diff --git a/permchain/channels/context.py b/permchain/channels/context.py
index 183fdafd6..97a95727d 100644
--- a/permchain/channels/context.py
+++ b/permchain/channels/context.py
@@ -22,7 +22,7 @@ from permchain.channels.base import (
 )
 
 
-class Context(Generic[Value], BaseChannel[Value, None]):
+class Context(Generic[Value], BaseChannel[Value, None, None]):
     """Exposes the value of a context manager, for the duration of an invocation.
     Context manager is entered before the first step, and exited after the last step.
     Optionally, provide an equivalent async context manager, which will be used
@@ -66,7 +66,7 @@ class Context(Generic[Value], BaseChannel[Value, None]):
         raise InvalidUpdateError()
 
     @contextmanager
-    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+    def empty(self, checkpoint: None = None) -> Generator[Self, None, None]:
         if self.ctx is None:
             raise ValueError("Cannot enter sync context manager.")
 
diff --git a/permchain/channels/last_value.py b/permchain/channels/last_value.py
index 6dc60e138..33f5cb0b8 100644
--- a/permchain/channels/last_value.py
+++ b/permchain/channels/last_value.py
@@ -1,4 +1,3 @@
-import json
 from contextlib import contextmanager
 from typing import Generator, Generic, Optional, Sequence, Type
 
@@ -12,7 +11,7 @@ from permchain.channels.base import (
 )
 
 
-class LastValue(Generic[Value], BaseChannel[Value, Value]):
+class LastValue(Generic[Value], BaseChannel[Value, Value, Value]):
     """Stores the last value received, can receive at most one value per step."""
 
     def __init__(self, typ: Type[Value]) -> None:
@@ -29,10 +28,10 @@ class LastValue(Generic[Value], BaseChannel[Value, Value]):
         return self.typ
 
     @contextmanager
-    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+    def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]:
         empty = self.__class__(self.typ)
         if checkpoint is not None:
-            empty.value = json.loads(checkpoint)
+            empty.value = checkpoint
         try:
             yield empty
         finally:
@@ -55,8 +54,8 @@ class LastValue(Generic[Value], BaseChannel[Value, Value]):
         except AttributeError:
             raise EmptyChannelError()
 
-    def checkpoint(self) -> str:
+    def checkpoint(self) -> Value:
         try:
-            return json.dumps(self.value)
+            return self.value
         except AttributeError:
             raise EmptyChannelError()
diff --git a/permchain/channels/topic.py b/permchain/channels/topic.py
index 71b821efe..03907725d 100644
--- a/permchain/channels/topic.py
+++ b/permchain/channels/topic.py
@@ -1,4 +1,3 @@
-import json
 from contextlib import contextmanager
 from typing import Any, Generator, Generic, Iterator, Optional, Sequence, Type, Union
 
@@ -15,7 +14,10 @@ def flatten(values: Sequence[Value | list[Value]]) -> Iterator[Value]:
             yield value
 
 
-class Topic(Generic[Value], BaseChannel[Sequence[Value], Value | list[Value]]):
+class Topic(
+    Generic[Value],
+    BaseChannel[Sequence[Value], Value | list[Value], tuple[set[Value], list[Value]]],
+):
     """A configurable PubSub Topic.
 
     Args:
@@ -46,12 +48,13 @@ class Topic(Generic[Value], BaseChannel[Sequence[Value], Value | list[Value]]):
         return Union[self.typ, list[self.typ]]  # type: ignore[name-defined]
 
     @contextmanager
-    def empty(self, checkpoint: Optional[str] = None) -> Generator[Self, None, None]:
+    def empty(
+        self, checkpoint: Optional[tuple[set[Value], list[Value]]] = None
+    ) -> Generator[Self, None, None]:
         empty = self.__class__(self.typ, self.unique, self.accumulate)
         if checkpoint is not None:
-            parsed = json.loads(checkpoint)
-            empty.seen = set(parsed["seen"])
-            empty.values = list(parsed["values"])
+            empty.seen = checkpoint[0]
+            empty.values = checkpoint[1]
         try:
             yield empty
         finally:
@@ -72,5 +75,5 @@ class Topic(Generic[Value], BaseChannel[Sequence[Value], Value | list[Value]]):
     def get(self) -> Sequence[Value]:
         return list(self.values)
 
-    def checkpoint(self) -> str:
-        return json.dumps({"seen": list(self.seen), "values": self.values})
+    def checkpoint(self) -> tuple[set[Value], list[Value]]:
+        return (self.seen, self.values)
diff --git a/tests/test_channels.py b/tests/test_channels.py
index 2b562da62..ed4e28028 100644
--- a/tests/test_channels.py
+++ b/tests/test_channels.py
@@ -27,6 +27,9 @@ def test_last_value() -> None:
         assert channel.get() == 3
         channel.update([4])
         assert channel.get() == 4
+        checkpoint = channel.checkpoint()
+    with LastValue(int).empty(checkpoint) as channel:
+        assert channel.get() == 4
 
 
 async def test_last_value_async() -> None:
@@ -43,6 +46,9 @@ async def test_last_value_async() -> None:
         assert channel.get() == 3
         channel.update([4])
         assert channel.get() == 4
+        checkpoint = channel.checkpoint()
+    async with LastValue(int).aempty(checkpoint) as channel:
+        assert channel.get() == 4
 
 
 def test_topic() -> None:
@@ -56,6 +62,11 @@ def test_topic() -> None:
         assert channel.get() == ["c", "d", "d"]
         channel.update([])
         assert channel.get() == []
+        channel.update(["e"])
+        assert channel.get() == ["e"]
+        checkpoint = channel.checkpoint()
+    with Topic(str).empty(checkpoint) as channel:
+        assert channel.get() == ["e"]
 
 
 async def test_topic_async() -> None:
@@ -69,6 +80,11 @@ async def test_topic_async() -> None:
         assert channel.get() == ["b", "c", "d", "d"]
         channel.update([])
         assert channel.get() == []
+        channel.update(["e"])
+        assert channel.get() == ["e"]
+        checkpoint = channel.checkpoint()
+    async with Topic(str).aempty(checkpoint) as channel:
+        assert channel.get() == ["e"]
 
 
 def test_topic_unique() -> None:
@@ -82,6 +98,13 @@ def test_topic_unique() -> None:
         assert channel.get() == ["c", "d"], "de-dupes from current and previous steps"
         channel.update([])
         assert channel.get() == []
+        channel.update(["e"])
+        assert channel.get() == ["e"]
+        checkpoint = channel.checkpoint()
+    with Topic(str, unique=True).empty(checkpoint) as channel:
+        assert channel.get() == ["e"]
+        channel.update(["d", "f"])
+        assert channel.get() == ["f"], "de-dupes from checkpoint"
 
 
 async def test_topic_unique_async() -> None:
@@ -95,6 +118,13 @@ async def test_topic_unique_async() -> None:
         assert channel.get() == ["c", "d"], "de-dupes from current and previous steps"
         channel.update([])
         assert channel.get() == []
+        channel.update(["e"])
+        assert channel.get() == ["e"]
+        checkpoint = channel.checkpoint()
+    async with Topic(str, unique=True).aempty(checkpoint) as channel:
+        assert channel.get() == ["e"]
+        channel.update(["d", "f"])
+        assert channel.get() == ["f"], "de-dupes from checkpoint"
 
 
 def test_topic_accumulate() -> None:
@@ -108,6 +138,11 @@ def test_topic_accumulate() -> None:
         assert channel.get() == ["a", "b", "b", "c", "d", "d"]
         channel.update([])
         assert channel.get() == ["a", "b", "b", "c", "d", "d"]
+        checkpoint = channel.checkpoint()
+    with Topic(str, accumulate=True).empty(checkpoint) as channel:
+        assert channel.get() == ["a", "b", "b", "c", "d", "d"]
+        channel.update(["e"])
+        assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"]
 
 
 async def test_topic_accumulate_async() -> None:
@@ -121,6 +156,11 @@ async def test_topic_accumulate_async() -> None:
         assert channel.get() == ["a", "b", "b", "c", "d", "d"]
         channel.update([])
         assert channel.get() == ["a", "b", "b", "c", "d", "d"]
+        checkpoint = channel.checkpoint()
+    async with Topic(str, accumulate=True).aempty(checkpoint) as channel:
+        assert channel.get() == ["a", "b", "b", "c", "d", "d"]
+        channel.update(["e"])
+        assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"]
 
 
 def test_topic_unique_accumulate() -> None:
@@ -134,6 +174,11 @@ def test_topic_unique_accumulate() -> None:
         assert channel.get() == ["a", "b", "c", "d"]
         channel.update([])
         assert channel.get() == ["a", "b", "c", "d"]
+        checkpoint = channel.checkpoint()
+    with Topic(str, unique=True, accumulate=True).empty(checkpoint) as channel:
+        assert channel.get() == ["a", "b", "c", "d"]
+        channel.update(["d", "e"])
+        assert channel.get() == ["a", "b", "c", "d", "e"]
 
 
 async def test_topic_unique_accumulate_async() -> None:
@@ -147,6 +192,11 @@ async def test_topic_unique_accumulate_async() -> None:
         assert channel.get() == ["a", "b", "c", "d"]
         channel.update([])
         assert channel.get() == ["a", "b", "c", "d"]
+        checkpoint = channel.checkpoint()
+    async with Topic(str, unique=True, accumulate=True).aempty(checkpoint) as channel:
+        assert channel.get() == ["a", "b", "c", "d"]
+        channel.update(["d", "e"])
+        assert channel.get() == ["a", "b", "c", "d", "e"]
 
 
 def test_binop() -> None:
@@ -161,6 +211,9 @@ def test_binop() -> None:
         assert channel.get() == 6
         channel.update([4])
         assert channel.get() == 10
+        checkpoint = channel.checkpoint()
+    with BinaryOperatorAggregate(int, operator.add).empty(checkpoint) as channel:
+        assert channel.get() == 10
 
 
 async def test_binop_async() -> None:
@@ -175,6 +228,9 @@ async def test_binop_async() -> None:
         assert channel.get() == 6
         channel.update([4])
         assert channel.get() == 10
+        checkpoint = channel.checkpoint()
+    async with BinaryOperatorAggregate(int, operator.add).aempty(checkpoint) as channel:
+        assert channel.get() == 10
 
 
 def test_ctx_manager(mocker: MockerFixture) -> None:
@@ -217,6 +273,18 @@ def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
         with pytest.raises(InvalidUpdateError):
             channel.update([5])  # type: ignore
 
+        checkpoint = channel.checkpoint()
+        assert checkpoint is None
+    with Context(httpx.Client).empty(checkpoint) as channel:
+        assert channel.ValueType is httpx.Client
+        with pytest.raises(InvalidUpdateError):
+            assert channel.UpdateType is None
+
+        assert isinstance(channel.get(), httpx.Client)
+
+        with pytest.raises(InvalidUpdateError):
+            channel.update([5])
+
 
 async def test_ctx_manager_async(mocker: MockerFixture) -> None:
     setup = mocker.Mock()

From b266b73ab90e92b1ea748ff5804e3802d594fec4 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Tue, 14 Nov 2023 11:27:31 +0000
Subject: [PATCH 2/6] Make channel managers aware of checkpoints

---
 permchain/channels/base.py   | 18 ++++++++++++++----
 permchain/pregel/__init__.py |  4 ++--
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/permchain/channels/base.py b/permchain/channels/base.py
index ee53993f5..1983147f5 100644
--- a/permchain/channels/base.py
+++ b/permchain/channels/base.py
@@ -80,11 +80,13 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
 
 @contextmanager
 def ChannelsManager(
-    channels: Mapping[str, BaseChannel]
+    channels: Mapping[str, BaseChannel],
+    checkpoint: Optional[Mapping[str, Any]],
 ) -> Generator[Mapping[str, BaseChannel], None, None]:
     """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
     # TODO use https://docs.python.org/3/library/contextlib.html#contextlib.ExitStack
-    empty = {k: v.empty() for k, v in channels.items()}
+    checkpoint = checkpoint or {}
+    empty = {k: v.empty(checkpoint.get(k)) for k, v in channels.items()}
     try:
         yield {k: v.__enter__() for k, v in empty.items()}
     finally:
@@ -94,12 +96,20 @@ def ChannelsManager(
 
 @asynccontextmanager
 async def AsyncChannelsManager(
-    channels: Mapping[str, BaseChannel]
+    channels: Mapping[str, BaseChannel],
+    checkpoint: Optional[Mapping[str, Any]],
 ) -> AsyncGenerator[Mapping[str, BaseChannel], None]:
     """Manage channels for the lifetime of a Pregel invocation (multiple steps)."""
-    empty = {k: v.aempty() for k, v in channels.items()}
+    checkpoint = checkpoint or {}
+    empty = {k: v.aempty(checkpoint.get(k)) for k, v in channels.items()}
     try:
         yield {k: await v.__aenter__() for k, v in empty.items()}
     finally:
         for v in empty.values():
             await v.__aexit__(None, None, None)
+
+
+def create_checkpoint(channels: Mapping[str, BaseChannel]) -> Mapping[str, Any]:
+    """Create a checkpoint for the given channels."""
+    checkpoint = {k: v.checkpoint() for k, v in channels.items()}
+    return {k: v for k, v in checkpoint.items() if v is not None}
diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py
index dcf48c73c..374d390c1 100644
--- a/permchain/pregel/__init__.py
+++ b/permchain/pregel/__init__.py
@@ -164,7 +164,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
     ) -> Iterator[dict[str, Any] | Any]:
         processes = {**self.chains}
         # TODO this is where we'd restore from checkpoint
-        with ChannelsManager(self.channels) as channels, get_executor_for_config(
+        with ChannelsManager(self.channels, None) as channels, get_executor_for_config(
             config
         ) as executor:
             next_tasks = _apply_writes_and_prepare_next_tasks(
@@ -243,7 +243,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
     ) -> AsyncIterator[dict[str, Any] | Any]:
         processes = {**self.chains}
         # TODO this is where we'd restore from checkpoint
-        async with AsyncChannelsManager(self.channels) as channels:
+        async with AsyncChannelsManager(self.channels, None) as channels:
             next_tasks = _apply_writes_and_prepare_next_tasks(
                 processes,
                 channels,

From de4a5ed418923fa8e7dc0c7d6cbfde39970c6830 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Tue, 14 Nov 2023 11:36:30 +0000
Subject: [PATCH 3/6] Update langchain, implement config_specs

---
 permchain/pregel/__init__.py | 10 ++++++++++
 permchain/pregel/read.py     | 10 +++++++---
 poetry.lock                  | 16 ++++++++--------
 pyproject.toml               |  2 +-
 4 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py
index 374d390c1..3148e0667 100644
--- a/permchain/pregel/__init__.py
+++ b/permchain/pregel/__init__.py
@@ -32,6 +32,10 @@ from langchain.schema.runnable.config import (
     get_executor_for_config,
     patch_config,
 )
+from langchain.schema.runnable.utils import (
+    ConfigurableFieldSpec,
+    get_unique_config_specs,
+)
 
 from permchain.channels.base import (
     AsyncChannelsManager,
@@ -121,6 +125,12 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         )
         return values
 
+    @property
+    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+        return get_unique_config_specs(
+            spec for chain in self.chains.values() for spec in chain.config_specs
+        )
+
     @property
     def InputType(self) -> Any:
         if isinstance(self.input, str):
diff --git a/permchain/pregel/read.py b/permchain/pregel/read.py
index 88c5b4522..4793dee67 100644
--- a/permchain/pregel/read.py
+++ b/permchain/pregel/read.py
@@ -5,12 +5,16 @@ from typing import Any, Callable, Mapping, Optional, Sequence
 from langchain.pydantic_v1 import Field
 from langchain.schema.runnable import (
     Runnable,
-    RunnableBinding,
     RunnableConfig,
     RunnableLambda,
     RunnablePassthrough,
 )
-from langchain.schema.runnable.base import Other, RunnableEach, coerce_to_runnable
+from langchain.schema.runnable.base import (
+    Other,
+    RunnableBindingBase,
+    RunnableEach,
+    coerce_to_runnable,
+)
 from langchain.schema.runnable.utils import ConfigurableFieldSpec
 
 from permchain.channels.base import BaseChannel
@@ -60,7 +64,7 @@ class ChannelRead(RunnableLambda):
 default_bound = RunnablePassthrough()
 
 
-class ChannelInvoke(RunnableBinding):
+class ChannelInvoke(RunnableBindingBase):
     channels: Mapping[None, str] | Mapping[str, str]
 
     bound: Runnable[Any, Any] = Field(default=default_bound)
diff --git a/poetry.lock b/poetry.lock
index 93fbf093e..d4bca125c 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1471,13 +1471,13 @@ files = [
 
 [[package]]
 name = "langchain"
-version = "0.0.330"
+version = "0.0.335"
 description = "Building applications with LLMs through composability"
 optional = false
 python-versions = ">=3.8.1,<4.0"
 files = [
-    {file = "langchain-0.0.330-py3-none-any.whl", hash = "sha256:ed557f4d680e02d9a05b175cae1ba146b7239d4d429d1d5271d4578b4956dbd6"},
-    {file = "langchain-0.0.330.tar.gz", hash = "sha256:5bed52769b63d76eb63589193e2efb66f5c7c429726af608e658f635335bd46a"},
+    {file = "langchain-0.0.335-py3-none-any.whl", hash = "sha256:f74c98366070a46953c071c69f6c01671a9437569c08406cace256ccaabdfcaf"},
+    {file = "langchain-0.0.335.tar.gz", hash = "sha256:93136fe6cc9ac06a80ccf7cf581e58af5cfcc31fef1083b30165df9a9bc53f5d"},
 ]
 
 [package.dependencies]
@@ -1486,7 +1486,7 @@ anyio = "<4.0"
 async-timeout = {version = ">=4.0.0,<5.0.0", markers = "python_version < \"3.11\""}
 dataclasses-json = ">=0.5.7,<0.7"
 jsonpatch = ">=1.33,<2.0"
-langsmith = ">=0.0.52,<0.1.0"
+langsmith = ">=0.0.63,<0.1.0"
 numpy = ">=1,<2"
 pydantic = ">=1,<3"
 PyYAML = ">=5.3"
@@ -1511,13 +1511,13 @@ text-helpers = ["chardet (>=5.1.0,<6.0.0)"]
 
 [[package]]
 name = "langsmith"
-version = "0.0.57"
+version = "0.0.64"
 description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform."
 optional = false
 python-versions = ">=3.8.1,<4.0"
 files = [
-    {file = "langsmith-0.0.57-py3-none-any.whl", hash = "sha256:d9d466cc45ce5224096ffb820d019b6f83678fc1f1021076ed75728aba60ec2b"},
-    {file = "langsmith-0.0.57.tar.gz", hash = "sha256:34929afd84cbfd46a8469229e3befc14c7e89186a0bee8ce9d084c7b8b271005"},
+    {file = "langsmith-0.0.64-py3-none-any.whl", hash = "sha256:461acdcd8332d1325c16dc57e8a2d5ec9d1578490a4eaabe14db74db74ceaf21"},
+    {file = "langsmith-0.0.64.tar.gz", hash = "sha256:e78c02501c2cff24fff7bd2d28ff3765b21675c7f0fcf6a09932bc218603c36e"},
 ]
 
 [package.dependencies]
@@ -3481,4 +3481,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.8.1,<4.0"
-content-hash = "ce9a7fe6e1972d14c6fd8aec807d4146098f6e4bc0efcbb0e8f10d922ff4901f"
+content-hash = "39ce01bbdc6757d98568030111d92e4c163667ebe674d8312f9f3b2c8b20f98a"
diff --git a/pyproject.toml b/pyproject.toml
index 0e58b144b..fbf453cd7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/permchain"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1,<4.0"
-langchain = ">=0.0.313"
+langchain = "^0.0.335"
 
 
 [tool.poetry.group.test.dependencies]

From f96ebb3357f6c102780199baf8941ed3fc3a61db Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Tue, 14 Nov 2023 12:14:18 +0000
Subject: [PATCH 4/6] Implement checkpoints

---
 permchain/channels/base.py       | 21 +++++++++---
 permchain/channels/context.py    |  2 +-
 permchain/checkpoint/__init__.py |  0
 permchain/checkpoint/base.py     | 44 ++++++++++++++++++++++++
 permchain/constants.py           |  4 +++
 permchain/pregel/__init__.py     | 59 ++++++++++++++++++++++++++------
 permchain/pregel/constants.py    |  2 --
 permchain/pregel/read.py         |  2 +-
 permchain/pregel/write.py        |  2 +-
 tests/test_channels.py           | 13 ++-----
 10 files changed, 118 insertions(+), 31 deletions(-)
 create mode 100644 permchain/checkpoint/__init__.py
 create mode 100644 permchain/checkpoint/base.py
 create mode 100644 permchain/constants.py
 delete mode 100644 permchain/pregel/constants.py

diff --git a/permchain/channels/base.py b/permchain/channels/base.py
index 1983147f5..eac51deac 100644
--- a/permchain/channels/base.py
+++ b/permchain/channels/base.py
@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager, contextmanager
+from datetime import datetime
 from typing import (
     Any,
     AsyncGenerator,
@@ -13,6 +14,8 @@ from typing import (
 
 from typing_extensions import Self
 
+from permchain.constants import CHECKPOINT_KEY_TS, CHECKPOINT_KEY_VERSION
+
 Value = TypeVar("Value")
 Update = TypeVar("Update")
 Checkpoint = TypeVar("Checkpoint")
@@ -72,10 +75,10 @@ class BaseChannel(Generic[Value, Update, Checkpoint], ABC):
 
     @abstractmethod
     def checkpoint(self) -> Checkpoint | None:
-        """Return a string representation of the channel's current state,
-        or None if the channel doesn't support checkpoints.
+        """Return a string representation of the channel's current state.
 
-        Raises EmptyChannelError if the channel is empty (never updated yet)."""
+        Raises EmptyChannelError if the channel is empty (never updated yet),
+        or doesn't supportcheckpoints."""
 
 
 @contextmanager
@@ -111,5 +114,13 @@ async def AsyncChannelsManager(
 
 def create_checkpoint(channels: Mapping[str, BaseChannel]) -> Mapping[str, Any]:
     """Create a checkpoint for the given channels."""
-    checkpoint = {k: v.checkpoint() for k, v in channels.items()}
-    return {k: v for k, v in checkpoint.items() if v is not None}
+    checkpoint = {
+        CHECKPOINT_KEY_VERSION: 1,
+        CHECKPOINT_KEY_TS: datetime.utcnow().isoformat(),
+    }
+    for k, v in channels.items():
+        try:
+            checkpoint[k] = v.checkpoint()
+        except EmptyChannelError:
+            pass
+    return checkpoint
diff --git a/permchain/channels/context.py b/permchain/channels/context.py
index 97a95727d..13098409f 100644
--- a/permchain/channels/context.py
+++ b/permchain/channels/context.py
@@ -107,4 +107,4 @@ class Context(Generic[Value], BaseChannel[Value, None, None]):
             raise EmptyChannelError()
 
     def checkpoint(self) -> None:
-        return None
+        raise EmptyChannelError()
diff --git a/permchain/checkpoint/__init__.py b/permchain/checkpoint/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/permchain/checkpoint/base.py b/permchain/checkpoint/base.py
new file mode 100644
index 000000000..ef6ae43e5
--- /dev/null
+++ b/permchain/checkpoint/base.py
@@ -0,0 +1,44 @@
+import asyncio
+import enum
+from abc import ABC, abstractmethod
+from typing import Any, Mapping, Sequence
+
+from langchain.load.serializable import Serializable
+from langchain.schema.runnable import RunnableConfig
+from langchain.schema.runnable.utils import ConfigurableFieldSpec
+
+
+# Before Python 3.11 native StrEnum is not available
+class StrEnum(str, enum.Enum):
+    """A string enum."""
+
+    pass
+
+
+class CheckpointAt(StrEnum):
+    END_OF_STEP = "end_of_step"
+    END_OF_RUN = "end_of_run"
+
+
+class BaseCheckpointAdapter(Serializable, ABC):
+    at: CheckpointAt = CheckpointAt.END_OF_RUN
+
+    @property
+    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+        return []
+
+    @abstractmethod
+    def get(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+        ...
+
+    @abstractmethod
+    def put(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+        ...
+
+    async def aget(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+        return asyncio.get_running_loop().run_in_executor(None, self.get, config)
+
+    async def aput(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+        return asyncio.get_running_loop().run_in_executor(
+            None, self.put, config, checkpoint
+        )
diff --git a/permchain/constants.py b/permchain/constants.py
new file mode 100644
index 000000000..bb96e6361
--- /dev/null
+++ b/permchain/constants.py
@@ -0,0 +1,4 @@
+CONFIG_KEY_SEND = "__pregel_send"
+CONFIG_KEY_READ = "__pregel_read"
+CHECKPOINT_KEY_VERSION = "__pregel_version"
+CHECKPOINT_KEY_TS = "__pregel_ts"
diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py
index 3148e0667..49ad344fc 100644
--- a/permchain/pregel/__init__.py
+++ b/permchain/pregel/__init__.py
@@ -42,8 +42,10 @@ from permchain.channels.base import (
     BaseChannel,
     ChannelsManager,
     EmptyChannelError,
+    create_checkpoint,
 )
-from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND
+from permchain.checkpoint.base import BaseCheckpointAdapter, CheckpointAt
+from permchain.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND
 from permchain.pregel.debug import print_checkpoint, print_step_start
 from permchain.pregel.io import map_input, map_output
 from permchain.pregel.log import logger
@@ -115,6 +117,8 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
 
     debug: bool = Field(default_factory=get_debug)
 
+    checkpoint: Optional[BaseCheckpointAdapter] = None
+
     class Config:
         arbitrary_types_allowed = True
 
@@ -128,7 +132,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
     @property
     def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
         return get_unique_config_specs(
-            spec for chain in self.chains.values() for spec in chain.config_specs
+            [spec for chain in self.chains.values() for spec in chain.config_specs]
+            + self.checkpoint.config_specs
+            if self.checkpoint is not None
+            else []
         )
 
     @property
@@ -173,10 +180,12 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         config: RunnableConfig,
     ) -> Iterator[dict[str, Any] | Any]:
         processes = {**self.chains}
-        # TODO this is where we'd restore from checkpoint
-        with ChannelsManager(self.channels, None) as channels, get_executor_for_config(
-            config
-        ) as executor:
+        checkpoint = (
+            self.checkpoint.get(config) if self.checkpoint is not None else None
+        )
+        with ChannelsManager(
+            self.channels, checkpoint
+        ) as channels, get_executor_for_config(config) as executor:
             next_tasks = _apply_writes_and_prepare_next_tasks(
                 processes,
                 channels,
@@ -239,12 +248,26 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 for output in map_output(self.output, pending_writes, channels):
                     yield output
 
-                # TODO this is where we'd save checkpoint
+                # save end of step checkpoint
+                if (
+                    self.checkpoint is not None
+                    and self.checkpoint.at == CheckpointAt.END_OF_STEP
+                ):
+                    checkpoint = create_checkpoint(channels)
+                    self.checkpoint.put(config, checkpoint)
 
                 # if no more tasks, we're done
                 if not next_tasks:
                     break
 
+            # save end of run checkpoint
+            if (
+                self.checkpoint is not None
+                and self.checkpoint.at == CheckpointAt.END_OF_RUN
+            ):
+                checkpoint = create_checkpoint(channels)
+                self.checkpoint.put(config, checkpoint)
+
     async def _atransform(
         self,
         input: AsyncIterator[dict[str, Any] | Any],
@@ -252,8 +275,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         config: RunnableConfig,
     ) -> AsyncIterator[dict[str, Any] | Any]:
         processes = {**self.chains}
-        # TODO this is where we'd restore from checkpoint
-        async with AsyncChannelsManager(self.channels, None) as channels:
+        checkpoint = (
+            await self.checkpoint.aget(config) if self.checkpoint is not None else None
+        )
+        async with AsyncChannelsManager(self.channels, checkpoint) as channels:
             next_tasks = _apply_writes_and_prepare_next_tasks(
                 processes,
                 channels,
@@ -319,12 +344,26 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
                 for output in map_output(self.output, pending_writes, channels):
                     yield output
 
-                # TODO this is where we'd save checkpoint
+                # save end of step checkpoint
+                if (
+                    self.checkpoint is not None
+                    and self.checkpoint.at == CheckpointAt.END_OF_STEP
+                ):
+                    checkpoint = create_checkpoint(channels)
+                    await self.checkpoint.aput(config, checkpoint)
 
                 # if no more tasks, we're done
                 if not next_tasks:
                     break
 
+            # save end of run checkpoint
+            if (
+                self.checkpoint is not None
+                and self.checkpoint.at == CheckpointAt.END_OF_RUN
+            ):
+                checkpoint = create_checkpoint(channels)
+                await self.checkpoint.aput(config, checkpoint)
+
     def invoke(
         self,
         input: dict[str, Any] | Any,
diff --git a/permchain/pregel/constants.py b/permchain/pregel/constants.py
deleted file mode 100644
index 6e5f9c37f..000000000
--- a/permchain/pregel/constants.py
+++ /dev/null
@@ -1,2 +0,0 @@
-CONFIG_KEY_SEND = "__pregel_send"
-CONFIG_KEY_READ = "__pregel_read"
diff --git a/permchain/pregel/read.py b/permchain/pregel/read.py
index 4793dee67..abb693e63 100644
--- a/permchain/pregel/read.py
+++ b/permchain/pregel/read.py
@@ -18,7 +18,7 @@ from langchain.schema.runnable.base import (
 from langchain.schema.runnable.utils import ConfigurableFieldSpec
 
 from permchain.channels.base import BaseChannel
-from permchain.pregel.constants import CONFIG_KEY_READ
+from permchain.constants import CONFIG_KEY_READ
 
 
 class ChannelRead(RunnableLambda):
diff --git a/permchain/pregel/write.py b/permchain/pregel/write.py
index 1cf9a5ccb..27480aa33 100644
--- a/permchain/pregel/write.py
+++ b/permchain/pregel/write.py
@@ -9,7 +9,7 @@ from langchain.schema.runnable import (
 )
 from langchain.schema.runnable.utils import ConfigurableFieldSpec
 
-from permchain.pregel.constants import CONFIG_KEY_SEND
+from permchain.constants import CONFIG_KEY_SEND
 
 TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None]
 
diff --git a/tests/test_channels.py b/tests/test_channels.py
index ed4e28028..104206c10 100644
--- a/tests/test_channels.py
+++ b/tests/test_channels.py
@@ -273,17 +273,8 @@ def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
         with pytest.raises(InvalidUpdateError):
             channel.update([5])  # type: ignore
 
-        checkpoint = channel.checkpoint()
-        assert checkpoint is None
-    with Context(httpx.Client).empty(checkpoint) as channel:
-        assert channel.ValueType is httpx.Client
-        with pytest.raises(InvalidUpdateError):
-            assert channel.UpdateType is None
-
-        assert isinstance(channel.get(), httpx.Client)
-
-        with pytest.raises(InvalidUpdateError):
-            channel.update([5])
+        with pytest.raises(EmptyChannelError):
+            channel.checkpoint()
 
 
 async def test_ctx_manager_async(mocker: MockerFixture) -> None:

From cec6ba4963d8746c577c3e68d395ef23bc725c4f Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Tue, 14 Nov 2023 12:38:17 +0000
Subject: [PATCH 5/6] Change binop to init value if possible

---
 permchain/channels/binop.py | 4 ++++
 tests/test_channels.py      | 6 ++----
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/permchain/channels/binop.py b/permchain/channels/binop.py
index 322f99c3e..8cbe092cc 100644
--- a/permchain/channels/binop.py
+++ b/permchain/channels/binop.py
@@ -19,6 +19,10 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]):
     def __init__(self, typ: Type[Value], operator: Callable[[Value, Value], Value]):
         self.typ = typ
         self.operator = operator
+        try:
+            self.value = typ()
+        except Exception:
+            pass
 
     @property
     def ValueType(self) -> Type[Value]:
diff --git a/tests/test_channels.py b/tests/test_channels.py
index 104206c10..e183ece2e 100644
--- a/tests/test_channels.py
+++ b/tests/test_channels.py
@@ -204,8 +204,7 @@ def test_binop() -> None:
         assert channel.ValueType is int
         assert channel.UpdateType is int
 
-        with pytest.raises(EmptyChannelError):
-            channel.get()
+        assert channel.get() == 0
 
         channel.update([1, 2, 3])
         assert channel.get() == 6
@@ -221,8 +220,7 @@ async def test_binop_async() -> None:
         assert channel.ValueType is int
         assert channel.UpdateType is int
 
-        with pytest.raises(EmptyChannelError):
-            channel.get()
+        assert channel.get() == 0
 
         channel.update([1, 2, 3])
         assert channel.get() == 6

From 682e7d175d303362803164a27c273b60ac767b00 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Tue, 14 Nov 2023 12:54:23 +0000
Subject: [PATCH 6/6] Add tests for checkpointing

---
 README.md                      | 11 ++++----
 permchain/checkpoint/base.py   |  4 +--
 permchain/checkpoint/memory.py | 29 +++++++++++++++++++
 permchain/pregel/read.py       |  3 ++
 permchain/pregel/validate.py   | 10 +++++++
 tests/test_pregel.py           | 41 +++++++++++++++++++++++++++
 tests/test_pregel_async.py     | 51 ++++++++++++++++++++++++++++++++++
 7 files changed, 142 insertions(+), 7 deletions(-)
 create mode 100644 permchain/checkpoint/memory.py

diff --git a/README.md b/README.md
index 60611839c..e74f0b270 100644
--- a/README.md
+++ b/README.md
@@ -81,11 +81,11 @@ Check `examples` for more examples.
   - [ ] Add tests for subscribe_to_each().join()
 - [x] Add optional debug logging
 - [ ] Add an optional Diff value for Channels that implements `__add__`, returned by update(), yielded by Pregel for output channels. Add replacing_keys set to AddableDict. use an addabledict for yielding values. channels that dont implement it get marked with replacing_keys
-- [ ] Implement checkpointing
-  - [ ] Use langchain.load dumps/loads functions (or use pickle?)
-  - [ ] Save checkpoints at end of each step
-  - [ ] Load checkpoint at start of invocation
-  - [ ] API to specify storage backend and save key
+- [x] Implement checkpointing
+  - [x] Save checkpoints at end of each step/run
+  - [x] Load checkpoint at start of invocation
+  - [x] API to specify storage backend and save key
+  - [x] Tests
 - [ ] Add more examples
   - [ ] multi agent simulation
   - [ ] human in the loop
@@ -93,6 +93,7 @@ Check `examples` for more examples.
   - [ ] agent executor (add current v total iterations info to read/write steps to enable doing a final update at the end)
   - [ ] run over dataset
 - [ ] Fault tolerance
+  - [ ] Expose a unique id to each step, hash of (app, chain, checkpoint) (include input updates for first step)
   - [ ] Retry individual processes in a step
   - [ ] Retry entire step?
 - [ ] Pregel.stream_log to contain additional keys specific to Pregel
diff --git a/permchain/checkpoint/base.py b/permchain/checkpoint/base.py
index ef6ae43e5..c735d32c8 100644
--- a/permchain/checkpoint/base.py
+++ b/permchain/checkpoint/base.py
@@ -36,9 +36,9 @@ class BaseCheckpointAdapter(Serializable, ABC):
         ...
 
     async def aget(self, config: RunnableConfig) -> Mapping[str, Any] | None:
-        return asyncio.get_running_loop().run_in_executor(None, self.get, config)
+        return await asyncio.get_running_loop().run_in_executor(None, self.get, config)
 
     async def aput(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
-        return asyncio.get_running_loop().run_in_executor(
+        return await asyncio.get_running_loop().run_in_executor(
             None, self.put, config, checkpoint
         )
diff --git a/permchain/checkpoint/memory.py b/permchain/checkpoint/memory.py
new file mode 100644
index 000000000..12d3a2932
--- /dev/null
+++ b/permchain/checkpoint/memory.py
@@ -0,0 +1,29 @@
+from typing import Any, Dict, Mapping, Sequence
+
+from langchain.pydantic_v1 import Field
+from langchain.schema.runnable import RunnableConfig
+from langchain.schema.runnable.utils import ConfigurableFieldSpec
+
+from permchain.checkpoint.base import BaseCheckpointAdapter
+
+
+class MemoryCheckpoint(BaseCheckpointAdapter):
+    storage: Dict[str, Mapping[str, Any]] = Field(default_factory=dict)
+
+    @property
+    def config_specs(self) -> Sequence[ConfigurableFieldSpec]:
+        return [
+            ConfigurableFieldSpec(
+                id="thread_id",
+                annotation=str,
+                name="Thread ID",
+                description=None,
+                default="",
+            ),
+        ]
+
+    def get(self, config: RunnableConfig) -> Mapping[str, Any] | None:
+        return self.storage.get(config["configurable"]["thread_id"], None)
+
+    def put(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+        return self.storage.update({config["configurable"]["thread_id"]: checkpoint})
diff --git a/permchain/pregel/read.py b/permchain/pregel/read.py
index abb693e63..66a20367c 100644
--- a/permchain/pregel/read.py
+++ b/permchain/pregel/read.py
@@ -89,6 +89,9 @@ class ChannelInvoke(RunnableBindingBase):
         )
 
     def join(self, channels: Sequence[str]) -> ChannelInvoke:
+        assert isinstance(channels, list) or isinstance(
+            channels, tuple
+        ), "channels must be a list or tuple"
         joiner = RunnablePassthrough.assign(
             **{chan: ChannelRead(chan) for chan in channels}
         )
diff --git a/permchain/pregel/validate.py b/permchain/pregel/validate.py
index afa51834e..927a78021 100644
--- a/permchain/pregel/validate.py
+++ b/permchain/pregel/validate.py
@@ -2,8 +2,14 @@ from typing import Any, Mapping, Sequence
 
 from permchain.channels.base import BaseChannel
 from permchain.channels.last_value import LastValue
+from permchain.constants import CHECKPOINT_KEY_TS, CHECKPOINT_KEY_VERSION
 from permchain.pregel.read import ChannelBatch, ChannelInvoke
 
+FORBIDDEN_CHANNEL_NAMES = {
+    CHECKPOINT_KEY_TS,
+    CHECKPOINT_KEY_VERSION,
+}
+
 
 def validate_chains_channels(
     chains: Mapping[str, ChannelInvoke | ChannelBatch],
@@ -47,3 +53,7 @@ def validate_chains_channels(
         for chan in output:
             if chan not in channels:
                 channels[chan] = LastValue(Any)
+
+    for name in FORBIDDEN_CHANNEL_NAMES:
+        if name in channels:
+            raise ValueError(f"Channel name {name} is reserved")
diff --git a/tests/test_pregel.py b/tests/test_pregel.py
index 619221270..d129c016e 100644
--- a/tests/test_pregel.py
+++ b/tests/test_pregel.py
@@ -1,3 +1,4 @@
+import operator
 import time
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
@@ -9,9 +10,11 @@ from pytest_mock import MockerFixture
 
 from permchain import Channel, Pregel
 from permchain.channels.base import InvalidUpdateError
+from permchain.channels.binop import BinaryOperatorAggregate
 from permchain.channels.context import Context
 from permchain.channels.last_value import LastValue
 from permchain.channels.topic import Topic
+from permchain.checkpoint.memory import MemoryCheckpoint
 
 
 def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
@@ -221,6 +224,44 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non
     assert app.invoke(2) == [3, 3]
 
 
+def test_invoke_checkpoint(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
+
+    def raise_if_above_10(input: int) -> int:
+        if input > 10:
+            raise ValueError("Input is too large")
+        return input
+
+    chain_one = (
+        Channel.subscribe_to(["input"]).join(["total"])
+        | add_one
+        | Channel.write_to("output", "total")
+        | raise_if_above_10
+    )
+
+    app = Pregel(
+        chains={"chain_one": chain_one},
+        channels={"total": BinaryOperatorAggregate(int, operator.add)},
+        checkpoint=MemoryCheckpoint(),
+    )
+
+    # total starts out as 0, so output is 0+2=2
+    assert app.invoke(2, {"configurable": {"thread_id": "1"}}) == 2
+    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 2
+    # total is now 2, so output is 2+3=5
+    assert app.invoke(3, {"configurable": {"thread_id": "1"}}) == 5
+    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
+    # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
+    with pytest.raises(ValueError):
+        app.invoke(4, {"configurable": {"thread_id": "1"}})
+    # checkpoint is not updated
+    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
+    # on a new thread, total starts out as 0, so output is 0+5=5
+    assert app.invoke(5, {"configurable": {"thread_id": "2"}}) == 5
+    assert app.checkpoint.get({"configurable": {"thread_id": "1"}}).get("total") == 7
+    assert app.checkpoint.get({"configurable": {"thread_id": "2"}}).get("total") == 5
+
+
 def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py
index b191b7438..091a11234 100644
--- a/tests/test_pregel_async.py
+++ b/tests/test_pregel_async.py
@@ -1,4 +1,5 @@
 import asyncio
+import operator
 from contextlib import asynccontextmanager, contextmanager
 from typing import Any, AsyncGenerator, AsyncIterator, Generator
 
@@ -8,9 +9,11 @@ from pytest_mock import MockerFixture
 
 from permchain import Channel, Pregel
 from permchain.channels.base import InvalidUpdateError
+from permchain.channels.binop import BinaryOperatorAggregate
 from permchain.channels.context import Context
 from permchain.channels.last_value import LastValue
 from permchain.channels.topic import Topic
+from permchain.checkpoint.memory import MemoryCheckpoint
 
 
 async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
@@ -230,6 +233,54 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture)
     assert await app.ainvoke(2) == [3, 3]
 
 
+async def test_invoke_checkpoint(mocker: MockerFixture) -> None:
+    add_one = mocker.Mock(side_effect=lambda x: x["total"] + x["input"])
+
+    def raise_if_above_10(input: int) -> int:
+        if input > 10:
+            raise ValueError("Input is too large")
+        return input
+
+    chain_one = (
+        Channel.subscribe_to(["input"]).join(["total"])
+        | add_one
+        | Channel.write_to("output", "total")
+        | raise_if_above_10
+    )
+
+    app = Pregel(
+        chains={"chain_one": chain_one},
+        channels={"total": BinaryOperatorAggregate(int, operator.add)},
+        checkpoint=MemoryCheckpoint(),
+    )
+
+    # total starts out as 0, so output is 0+2=2
+    assert await app.ainvoke(2, {"configurable": {"thread_id": "1"}}) == 2
+    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
+        "total"
+    ) == 2
+    # total is now 2, so output is 2+3=5
+    assert await app.ainvoke(3, {"configurable": {"thread_id": "1"}}) == 5
+    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
+        "total"
+    ) == 7
+    # total is now 2+5=7, so output would be 7+4=11, but raises ValueError
+    with pytest.raises(ValueError):
+        await app.ainvoke(4, {"configurable": {"thread_id": "1"}})
+    # checkpoint is not updated
+    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
+        "total"
+    ) == 7
+    # on a new thread, total starts out as 0, so output is 0+5=5
+    assert await app.ainvoke(5, {"configurable": {"thread_id": "2"}}) == 5
+    assert (await app.checkpoint.aget({"configurable": {"thread_id": "1"}})).get(
+        "total"
+    ) == 7
+    assert (await app.checkpoint.aget({"configurable": {"thread_id": "2"}})).get(
+        "total"
+    ) == 5
+
+
 async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None:
     add_one = mocker.Mock(side_effect=lambda x: x + 1)
     add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))