From 4fd90f8ad628740e0d3dd70f5dc6c844958db3c2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sun, 31 Mar 2024 09:25:08 -0700 Subject: [PATCH] Rename BaseChannel.empty() to BaseChannel.from_checkpoint() - clearer name - channels with complex data structures in checkpoint (eg a set or list) now copy the checkpointed data structures before creating the new channel --- langgraph/channels/any_value.py | 16 +++--- langgraph/channels/base.py | 37 ++++++++----- langgraph/channels/binop.py | 16 +++--- langgraph/channels/context.py | 12 ++--- langgraph/channels/ephemeral_value.py | 16 +++--- langgraph/channels/last_value.py | 16 +++--- langgraph/channels/named_barrier_value.py | 16 +++--- langgraph/channels/topic.py | 12 ++--- tests/test_channels.py | 66 +++++++++++++---------- 9 files changed, 119 insertions(+), 88 deletions(-) diff --git a/langgraph/channels/any_value.py b/langgraph/channels/any_value.py index 73edb3494..d6a2d4e52 100644 --- a/langgraph/channels/any_value.py +++ b/langgraph/channels/any_value.py @@ -23,8 +23,16 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]): """The type of the update received by the channel.""" return self.typ + def checkpoint(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() + @contextmanager - def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: + def from_checkpoint( + self, checkpoint: Optional[Value] = None + ) -> Generator[Self, None, None]: empty = self.__class__(self.typ) if checkpoint is not None: empty.value = checkpoint @@ -51,9 +59,3 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]): return self.value except AttributeError: raise EmptyChannelError() - - def checkpoint(self) -> Value: - try: - return self.value - except AttributeError: - raise EmptyChannelError() diff --git a/langgraph/channels/base.py b/langgraph/channels/base.py index 71680c8ea..c244e750c 100644 --- a/langgraph/channels/base.py +++ b/langgraph/channels/base.py @@ -45,19 +45,33 @@ class BaseChannel(Generic[Value, Update, C], ABC): def UpdateType(self) -> Any: """The type of the update received by the channel.""" + # ser/de methods + + @abstractmethod + def checkpoint(self) -> Optional[C]: + """Return a serializable representation of the channel's current state. + Raises EmptyChannelError if the channel is empty (never updated yet), + or doesn't support checkpoints.""" + @contextmanager @abstractmethod - def empty(self, checkpoint: Optional[C] = None) -> Generator[Self, None, None]: - """Return a new identical channel, optionally initialized from a checkpoint.""" + def from_checkpoint( + self, checkpoint: Optional[C] = None + ) -> Generator[Self, None, None]: + """Return a new identical channel, optionally initialized from a checkpoint. + If the checkpoint contains complex data structures, they should be copied.""" @asynccontextmanager - async def aempty( + async def afrom_checkpoint( self, checkpoint: Optional[C] = None ) -> AsyncGenerator[Self, None]: - """Return a new identical channel, optionally initialized from a checkpoint.""" - with self.empty(checkpoint) as value: + """Return a new identical channel, optionally initialized from a checkpoint. + If the checkpoint contains complex data structures, they should be copied.""" + with self.from_checkpoint(checkpoint) as value: yield value + # state methods + @abstractmethod def update(self, values: Sequence[Update]) -> None: """Update the channel's value with the given sequence of updates. @@ -71,13 +85,6 @@ class BaseChannel(Generic[Value, Update, C], ABC): Raises EmptyChannelError if the channel is empty (never updated yet).""" - @abstractmethod - def checkpoint(self) -> Optional[C]: - """Return a string representation of the channel's current state. - - Raises EmptyChannelError if the channel is empty (never updated yet), - or doesn't supportcheckpoints.""" - @contextmanager def ChannelsManager( @@ -87,7 +94,8 @@ def ChannelsManager( """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(checkpoint["channel_values"].get(k)) for k, v in channels.items() + k: v.from_checkpoint(checkpoint["channel_values"].get(k)) + for k, v in channels.items() } try: yield {k: v.__enter__() for k, v in empty.items()} @@ -103,7 +111,8 @@ async def AsyncChannelsManager( ) -> AsyncGenerator[Mapping[str, BaseChannel], None]: """Manage channels for the lifetime of a Pregel invocation (multiple steps).""" empty = { - k: v.aempty(checkpoint["channel_values"].get(k)) for k, v in channels.items() + k: v.afrom_checkpoint(checkpoint["channel_values"].get(k)) + for k, v in channels.items() } try: yield {k: await v.__aenter__() for k, v in empty.items()} diff --git a/langgraph/channels/binop.py b/langgraph/channels/binop.py index 6ec62fcb9..53a6c9f5d 100644 --- a/langgraph/channels/binop.py +++ b/langgraph/channels/binop.py @@ -34,8 +34,16 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): """The type of the update received by the channel.""" return self.typ + def checkpoint(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() + @contextmanager - def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: + def from_checkpoint( + self, checkpoint: Optional[Value] = None + ) -> Generator[Self, None, None]: empty = self.__class__(self.typ, self.operator) if checkpoint is not None: empty.value = checkpoint @@ -62,9 +70,3 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): return self.value except AttributeError: raise EmptyChannelError() - - def checkpoint(self) -> Value: - try: - return self.value - except AttributeError: - raise EmptyChannelError() diff --git a/langgraph/channels/context.py b/langgraph/channels/context.py index d37e96305..fca05d950 100644 --- a/langgraph/channels/context.py +++ b/langgraph/channels/context.py @@ -65,8 +65,11 @@ class Context(Generic[Value], BaseChannel[Value, None, None]): """The type of the update received by the channel.""" raise InvalidUpdateError() + def checkpoint(self) -> None: + raise EmptyChannelError() + @contextmanager - def empty(self, checkpoint: None = None) -> Generator[Self, None, None]: + def from_checkpoint(self, checkpoint: None = None) -> Generator[Self, None, None]: if self.ctx is None: raise ValueError("Cannot enter sync context manager.") @@ -80,7 +83,7 @@ class Context(Generic[Value], BaseChannel[Value, None, None]): ctx.__exit__(None, None, None) @asynccontextmanager - async def aempty( + async def afrom_checkpoint( self, checkpoint: Optional[str] = None ) -> AsyncGenerator[Self, None]: if self.actx is not None: @@ -93,7 +96,7 @@ class Context(Generic[Value], BaseChannel[Value, None, None]): finally: await actx.__aexit__(None, None, None) else: - with self.empty() as empty: + with self.from_checkpoint() as empty: yield empty def update(self, values: Sequence[None]) -> None: @@ -105,6 +108,3 @@ class Context(Generic[Value], BaseChannel[Value, None, None]): return self.value except AttributeError: raise EmptyChannelError() - - def checkpoint(self) -> None: - raise EmptyChannelError() diff --git a/langgraph/channels/ephemeral_value.py b/langgraph/channels/ephemeral_value.py index 2baa2f461..fb0a27e2f 100644 --- a/langgraph/channels/ephemeral_value.py +++ b/langgraph/channels/ephemeral_value.py @@ -28,8 +28,16 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): """The type of the update received by the channel.""" return self.typ + def checkpoint(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() + @contextmanager - def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: + def from_checkpoint( + self, checkpoint: Optional[Value] = None + ) -> Generator[Self, None, None]: empty = self.__class__(self.typ, self.guard) if checkpoint is not None: empty.value = checkpoint @@ -59,9 +67,3 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): return self.value except AttributeError: raise EmptyChannelError() - - def checkpoint(self) -> Value: - try: - return self.value - except AttributeError: - raise EmptyChannelError() diff --git a/langgraph/channels/last_value.py b/langgraph/channels/last_value.py index 858943c91..fd7495e02 100644 --- a/langgraph/channels/last_value.py +++ b/langgraph/channels/last_value.py @@ -27,8 +27,16 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): """The type of the update received by the channel.""" return self.typ + def checkpoint(self) -> Value: + try: + return self.value + except AttributeError: + raise EmptyChannelError() + @contextmanager - def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: + def from_checkpoint( + self, checkpoint: Optional[Value] = None + ) -> Generator[Self, None, None]: empty = self.__class__(self.typ) if checkpoint is not None: empty.value = checkpoint @@ -53,9 +61,3 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): return self.value except AttributeError: raise EmptyChannelError() - - def checkpoint(self) -> Value: - try: - return self.value - except AttributeError: - raise EmptyChannelError() diff --git a/langgraph/channels/named_barrier_value.py b/langgraph/channels/named_barrier_value.py index 7f55de31b..0b42686ed 100644 --- a/langgraph/channels/named_barrier_value.py +++ b/langgraph/channels/named_barrier_value.py @@ -11,10 +11,10 @@ from langgraph.channels.base import ( ) -class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, Value]): +class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]): """A channel that waits until all named values are received before making the value available.""" - def __init__(self, typ: Type[Value], names: set[str]) -> None: + def __init__(self, typ: Type[Value], names: set[Value]) -> None: self.typ = typ self.names = names self.seen = set() @@ -29,11 +29,16 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, Value]): """The type of the update received by the channel.""" return self.typ + def checkpoint(self) -> set[Value]: + return self.seen + @contextmanager - def empty(self, checkpoint: Optional[Value] = None) -> Generator[Self, None, None]: + def from_checkpoint( + self, checkpoint: Optional[set[Value]] = None + ) -> Generator[Self, None, None]: empty = self.__class__(self.typ, self.names) if checkpoint is not None: - empty.seen = checkpoint + empty.seen = checkpoint.copy() try: yield empty @@ -53,6 +58,3 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, Value]): if self.seen != self.names: raise EmptyChannelError() return None - - def checkpoint(self) -> Value: - return self.seen diff --git a/langgraph/channels/topic.py b/langgraph/channels/topic.py index 1a7cec7f3..f6ea87709 100644 --- a/langgraph/channels/topic.py +++ b/langgraph/channels/topic.py @@ -49,14 +49,17 @@ class Topic( """The type of the update received by the channel.""" return Union[self.typ, list[self.typ]] # type: ignore[name-defined] + def checkpoint(self) -> tuple[set[Value], list[Value]]: + return (self.seen, self.values) + @contextmanager - def empty( + def from_checkpoint( 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: - empty.seen = checkpoint[0] - empty.values = checkpoint[1] + empty.seen = checkpoint[0].copy() + empty.values = checkpoint[1].copy() try: yield empty finally: @@ -76,6 +79,3 @@ class Topic( def get(self) -> Sequence[Value]: return list(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 51b70eb04..6f24ba765 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -14,7 +14,7 @@ from langgraph.channels.topic import Topic def test_last_value() -> None: - with LastValue(int).empty() as channel: + with LastValue(int).from_checkpoint() as channel: assert channel.ValueType is int assert channel.UpdateType is int @@ -28,12 +28,12 @@ def test_last_value() -> None: channel.update([4]) assert channel.get() == 4 checkpoint = channel.checkpoint() - with LastValue(int).empty(checkpoint) as channel: + with LastValue(int).from_checkpoint(checkpoint) as channel: assert channel.get() == 4 async def test_last_value_async() -> None: - async with LastValue(int).aempty() as channel: + async with LastValue(int).afrom_checkpoint() as channel: assert channel.ValueType is int assert channel.UpdateType is int @@ -47,12 +47,12 @@ async def test_last_value_async() -> None: channel.update([4]) assert channel.get() == 4 checkpoint = channel.checkpoint() - async with LastValue(int).aempty(checkpoint) as channel: + async with LastValue(int).afrom_checkpoint(checkpoint) as channel: assert channel.get() == 4 def test_topic() -> None: - with Topic(str).empty() as channel: + with Topic(str).from_checkpoint() as channel: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] @@ -65,12 +65,16 @@ def test_topic() -> None: channel.update(["e"]) assert channel.get() == ["e"] checkpoint = channel.checkpoint() - with Topic(str).empty(checkpoint) as channel: + with Topic(str).from_checkpoint(checkpoint) as channel: assert channel.get() == ["e"] + with Topic(str).from_checkpoint(checkpoint) as channel_copy: + channel_copy.update(["f"]) + assert channel_copy.get() == ["f"] + assert channel.get() == ["e"] async def test_topic_async() -> None: - async with Topic(str).aempty() as channel: + async with Topic(str).afrom_checkpoint() as channel: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] @@ -83,12 +87,12 @@ async def test_topic_async() -> None: channel.update(["e"]) assert channel.get() == ["e"] checkpoint = channel.checkpoint() - async with Topic(str).aempty(checkpoint) as channel: + async with Topic(str).afrom_checkpoint(checkpoint) as channel: assert channel.get() == ["e"] def test_topic_unique() -> None: - with Topic(str, unique=True).empty() as channel: + with Topic(str, unique=True).from_checkpoint() as channel: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] @@ -101,14 +105,14 @@ def test_topic_unique() -> None: channel.update(["e"]) assert channel.get() == ["e"] checkpoint = channel.checkpoint() - with Topic(str, unique=True).empty(checkpoint) as channel: + with Topic(str, unique=True).from_checkpoint(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: - async with Topic(str, unique=True).aempty() as channel: + async with Topic(str, unique=True).afrom_checkpoint() as channel: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] @@ -121,14 +125,14 @@ async def test_topic_unique_async() -> None: channel.update(["e"]) assert channel.get() == ["e"] checkpoint = channel.checkpoint() - async with Topic(str, unique=True).aempty(checkpoint) as channel: + async with Topic(str, unique=True).afrom_checkpoint(checkpoint) as channel: assert channel.get() == ["e"] channel.update(["d", "f"]) assert channel.get() == ["f"], "de-dupes from checkpoint" def test_topic_accumulate() -> None: - with Topic(str, accumulate=True).empty() as channel: + with Topic(str, accumulate=True).from_checkpoint() as channel: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] @@ -139,14 +143,14 @@ def test_topic_accumulate() -> None: channel.update([]) assert channel.get() == ["a", "b", "b", "c", "d", "d"] checkpoint = channel.checkpoint() - with Topic(str, accumulate=True).empty(checkpoint) as channel: + with Topic(str, accumulate=True).from_checkpoint(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: - async with Topic(str, accumulate=True).aempty() as channel: + async with Topic(str, accumulate=True).afrom_checkpoint() as channel: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] @@ -157,14 +161,14 @@ async def test_topic_accumulate_async() -> None: channel.update([]) assert channel.get() == ["a", "b", "b", "c", "d", "d"] checkpoint = channel.checkpoint() - async with Topic(str, accumulate=True).aempty(checkpoint) as channel: + async with Topic(str, accumulate=True).afrom_checkpoint(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: - with Topic(str, unique=True, accumulate=True).empty() as channel: + with Topic(str, unique=True, accumulate=True).from_checkpoint() as channel: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] @@ -175,14 +179,16 @@ def test_topic_unique_accumulate() -> None: channel.update([]) assert channel.get() == ["a", "b", "c", "d"] checkpoint = channel.checkpoint() - with Topic(str, unique=True, accumulate=True).empty(checkpoint) as channel: + with Topic(str, unique=True, accumulate=True).from_checkpoint( + 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: - async with Topic(str, unique=True, accumulate=True).aempty() as channel: + async with Topic(str, unique=True, accumulate=True).afrom_checkpoint() as channel: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] @@ -193,14 +199,16 @@ async def test_topic_unique_accumulate_async() -> None: channel.update([]) assert channel.get() == ["a", "b", "c", "d"] checkpoint = channel.checkpoint() - async with Topic(str, unique=True, accumulate=True).aempty(checkpoint) as channel: + async with Topic(str, unique=True, accumulate=True).afrom_checkpoint( + 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: - with BinaryOperatorAggregate(int, operator.add).empty() as channel: + with BinaryOperatorAggregate(int, operator.add).from_checkpoint() as channel: assert channel.ValueType is int assert channel.UpdateType is int @@ -211,12 +219,14 @@ def test_binop() -> None: channel.update([4]) assert channel.get() == 10 checkpoint = channel.checkpoint() - with BinaryOperatorAggregate(int, operator.add).empty(checkpoint) as channel: + with BinaryOperatorAggregate(int, operator.add).from_checkpoint( + checkpoint + ) as channel: assert channel.get() == 10 async def test_binop_async() -> None: - async with BinaryOperatorAggregate(int, operator.add).aempty() as channel: + async with BinaryOperatorAggregate(int, operator.add).afrom_checkpoint() as channel: assert channel.ValueType is int assert channel.UpdateType is int @@ -227,7 +237,9 @@ async def test_binop_async() -> None: channel.update([4]) assert channel.get() == 10 checkpoint = channel.checkpoint() - async with BinaryOperatorAggregate(int, operator.add).aempty(checkpoint) as channel: + async with BinaryOperatorAggregate(int, operator.add).afrom_checkpoint( + checkpoint + ) as channel: assert channel.get() == 10 @@ -243,7 +255,7 @@ def test_ctx_manager(mocker: MockerFixture) -> None: finally: cleanup() - with Context(an_int, None, int).empty() as channel: + with Context(an_int, None, int).from_checkpoint() as channel: assert setup.call_count == 1 assert cleanup.call_count == 0 @@ -261,7 +273,7 @@ def test_ctx_manager(mocker: MockerFixture) -> None: def test_ctx_manager_ctx(mocker: MockerFixture) -> None: - with Context(httpx.Client).empty() as channel: + with Context(httpx.Client).from_checkpoint() as channel: assert channel.ValueType is httpx.Client with pytest.raises(InvalidUpdateError): assert channel.UpdateType is None @@ -294,7 +306,7 @@ async def test_ctx_manager_async(mocker: MockerFixture) -> None: finally: cleanup() - async with Context(an_int_sync, an_int, int).aempty() as channel: + async with Context(an_int_sync, an_int, int).afrom_checkpoint() as channel: assert setup.call_count == 1 assert cleanup.call_count == 0