Merge pull request #16 from langchain-ai/nc/checkpoints

Nc/checkpoints

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

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,8 +14,11 @@ 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")
 
 
 class EmptyChannelError(Exception):
@@ -30,7 +34,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 +47,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,20 +74,22 @@ class BaseChannel(Generic[Value, Update], ABC):
         Raises EmptyChannelError if the channel is empty (never updated yet)."""
 
     @abstractmethod
-    def checkpoint(self) -> str | None:
-        """Return a string representation of the channel's current state,
-        or None if the channel doesn't support checkpoints.
+    def checkpoint(self) -> Checkpoint | None:
+        """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
 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:
@@ -91,12 +99,28 @@ 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 = {
+        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

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
@@ -20,6 +19,10 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[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]:
@@ -32,10 +35,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 +63,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()

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.")
 
@@ -107,4 +107,4 @@ class Context(Generic[Value], BaseChannel[Value, None]):
             raise EmptyChannelError()
 
     def checkpoint(self) -> None:
-        return None
+        raise EmptyChannelError()

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()

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)

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 await asyncio.get_running_loop().run_in_executor(None, self.get, config)
+
+    async def aput(self, config: RunnableConfig, checkpoint: Mapping[str, Any]) -> None:
+        return await asyncio.get_running_loop().run_in_executor(
+            None, self.put, config, checkpoint
+        )

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})

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"

permchain/pregel/__init__.py

@@ -32,14 +32,20 @@ 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,
     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
@@ -111,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
 
@@ -121,6 +129,15 @@ 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]
+            + self.checkpoint.config_specs
+            if self.checkpoint is not None
+            else []
+        )
+
     @property
     def InputType(self) -> Any:
         if isinstance(self.input, str):
@@ -163,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) 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,
@@ -229,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],
@@ -242,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) 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,
@@ -309,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,

permchain/pregel/constants.py

@@ -1,2 +0,0 @@
-CONFIG_KEY_SEND = "__pregel_send"
-CONFIG_KEY_READ = "__pregel_read"

permchain/pregel/read.py

@@ -5,16 +5,20 @@ 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
-from permchain.pregel.constants import CONFIG_KEY_READ
+from permchain.constants import CONFIG_KEY_READ
 
 
 class ChannelRead(RunnableLambda):
@@ -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)
@@ -85,6 +89,9 @@ class ChannelInvoke(RunnableBinding):
         )
 
     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}
         )

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")

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]
 

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"

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]

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:
@@ -154,13 +204,15 @@ 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
         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:
@@ -168,13 +220,15 @@ 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
         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 +271,9 @@ def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
         with pytest.raises(InvalidUpdateError):
             channel.update([5])  # type: ignore
 
+        with pytest.raises(EmptyChannelError):
+            channel.checkpoint()
+
 
 async def test_ctx_manager_async(mocker: MockerFixture) -> None:
     setup = mocker.Mock()

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))

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))