From 758d62fd22ac14e8ebd208eeec80ea9ad378d903 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 14 Jun 2024 11:10:11 -0700 Subject: [PATCH 1/2] Version all channel changes - previously when channels were cleared via update([]) or consume() that would not bump the version - that behavior is not compatible with checkpointers that store blobs separately, which join checkpoint and blobs using version - now update() and consume() return bool to indicate if version should be bumped - this also avoids bumping version when nothing changed --- langgraph/channels/any_value.py | 7 +- langgraph/channels/base.py | 6 +- langgraph/channels/binop.py | 6 +- langgraph/channels/context.py | 3 +- langgraph/channels/dynamic_barrier_value.py | 13 +++- langgraph/channels/ephemeral_value.py | 8 +-- langgraph/channels/last_value.py | 5 +- langgraph/channels/named_barrier_value.py | 12 +++- langgraph/channels/topic.py | 8 ++- langgraph/checkpoint/sqlite.py | 8 ++- langgraph/pregel/__init__.py | 36 ++++++++-- tests/test_channels.py | 77 ++++++++++++--------- tests/test_pregel.py | 4 +- tests/test_pregel_async.py | 4 +- 14 files changed, 129 insertions(+), 68 deletions(-) diff --git a/langgraph/channels/any_value.py b/langgraph/channels/any_value.py index caddc05bf..23602bcfe 100644 --- a/langgraph/channels/any_value.py +++ b/langgraph/channels/any_value.py @@ -45,15 +45,16 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]): except AttributeError: pass - def update(self, values: Sequence[Value]) -> None: + def update(self, values: Sequence[Value]) -> bool: if len(values) == 0: try: del self.value + return True except AttributeError: - pass - return + return False self.value = values[-1] + return True def get(self) -> Value: try: diff --git a/langgraph/channels/base.py b/langgraph/channels/base.py index b45d1548c..8133edf70 100644 --- a/langgraph/channels/base.py +++ b/langgraph/channels/base.py @@ -58,7 +58,7 @@ class BaseChannel(Generic[Value, Update, C], ABC): # state methods @abstractmethod - def update(self, values: Sequence[Update]) -> None: + def update(self, values: Sequence[Update]) -> bool: """Update the channel's value with the given sequence of updates. The order of the updates in the sequence is arbitrary. This method is called by Pregel for all channels at the end of each step. @@ -71,12 +71,12 @@ class BaseChannel(Generic[Value, Update, C], ABC): Raises EmptyChannelError if the channel is empty (never updated yet).""" - def consume(self) -> None: + def consume(self) -> bool: """Mark the current value of the channel as consumed. By default, no-op. This is called by Pregel before the start of the next step, for all channels that triggered a node. """ - pass + return False __all__ = [ diff --git a/langgraph/channels/binop.py b/langgraph/channels/binop.py index 03b4b62d8..b51b8efe1 100644 --- a/langgraph/channels/binop.py +++ b/langgraph/channels/binop.py @@ -85,15 +85,15 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): except AttributeError: pass - def update(self, values: Sequence[Value]) -> None: + def update(self, values: Sequence[Value]) -> bool: if not values: - return + return False if not hasattr(self, "value"): self.value = values[0] values = values[1:] - for value in values: self.value = self.operator(self.value, value) + return True def get(self) -> Value: try: diff --git a/langgraph/channels/context.py b/langgraph/channels/context.py index 205f9d6d3..dbcfa1627 100644 --- a/langgraph/channels/context.py +++ b/langgraph/channels/context.py @@ -95,9 +95,10 @@ class Context(Generic[Value], BaseChannel[Value, None, None]): with self.from_checkpoint() as empty: yield empty - def update(self, values: Sequence[None]) -> None: + def update(self, values: Sequence[None]) -> bool: if values: raise InvalidUpdateError() + return False def get(self) -> Value: try: diff --git a/langgraph/channels/dynamic_barrier_value.py b/langgraph/channels/dynamic_barrier_value.py index 45cd83a7c..875c55c45 100644 --- a/langgraph/channels/dynamic_barrier_value.py +++ b/langgraph/channels/dynamic_barrier_value.py @@ -59,27 +59,34 @@ class DynamicBarrierValue( finally: pass - def update(self, values: Sequence[Union[Value, WaitForNames]]) -> None: + def update(self, values: Sequence[Union[Value, WaitForNames]]) -> bool: if wait_for_names := [v for v in values if isinstance(v, WaitForNames)]: if len(wait_for_names) > 1: raise InvalidUpdateError( "Received multiple WaitForNames updates in the same step." ) self.names = wait_for_names[0].names + return True elif self.names is not None: + updated = False for value in values: assert not isinstance(value, WaitForNames) if value in self.names: - self.seen.add(value) + if value not in self.seen: + self.seen.add(value) + updated = True else: raise InvalidUpdateError(f"Value {value} not in {self.names}") + return updated def get(self) -> Value: if self.seen != self.names: raise EmptyChannelError() return None - def consume(self) -> None: + def consume(self) -> bool: if self.seen == self.names: self.seen = set() self.names = None + return True + return False diff --git a/langgraph/channels/ephemeral_value.py b/langgraph/channels/ephemeral_value.py index 17a11f552..eb1dd9d36 100644 --- a/langgraph/channels/ephemeral_value.py +++ b/langgraph/channels/ephemeral_value.py @@ -45,20 +45,20 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): except AttributeError: pass - def update(self, values: Sequence[Value]) -> None: + def update(self, values: Sequence[Value]) -> bool: if len(values) == 0: try: del self.value + return True except AttributeError: - pass - finally: - return + return False if len(values) != 1 and self.guard: raise InvalidUpdateError( "EphemeralValue can only receive one value per step." ) self.value = values[-1] + return True def get(self) -> Value: try: diff --git a/langgraph/channels/last_value.py b/langgraph/channels/last_value.py index c3748b7f9..dbae1306d 100644 --- a/langgraph/channels/last_value.py +++ b/langgraph/channels/last_value.py @@ -44,13 +44,14 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): except AttributeError: pass - def update(self, values: Sequence[Value]) -> None: + def update(self, values: Sequence[Value]) -> bool: if len(values) == 0: - return + return False if len(values) != 1: raise InvalidUpdateError("LastValue can only receive one value per step.") self.value = values[-1] + return True def get(self) -> Value: try: diff --git a/langgraph/channels/named_barrier_value.py b/langgraph/channels/named_barrier_value.py index 8181d0a02..bd97754bf 100644 --- a/langgraph/channels/named_barrier_value.py +++ b/langgraph/channels/named_barrier_value.py @@ -41,18 +41,24 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]): finally: pass - def update(self, values: Sequence[Value]) -> None: + def update(self, values: Sequence[Value]) -> bool: + updated = False for value in values: if value in self.names: - self.seen.add(value) + if value not in self.seen: + self.seen.add(value) + updated = True else: raise InvalidUpdateError(f"Value {value} not in {self.names}") + return updated def get(self) -> Value: if self.seen != self.names: raise EmptyChannelError() return None - def consume(self) -> None: + def consume(self) -> bool: if self.seen == self.names: self.seen = set() + return True + return False diff --git a/langgraph/channels/topic.py b/langgraph/channels/topic.py index c9bebdbfd..960cd0dfa 100644 --- a/langgraph/channels/topic.py +++ b/langgraph/channels/topic.py @@ -4,6 +4,7 @@ from typing import Any, Generator, Generic, Iterator, Optional, Sequence, Type, from typing_extensions import Self from langgraph.channels.base import BaseChannel, Value +from langgraph.errors import EmptyChannelError def flatten(values: Sequence[Union[Value, list[Value]]]) -> Iterator[Value]: @@ -66,6 +67,7 @@ class Topic( pass def update(self, values: Sequence[Union[Value, list[Value]]]) -> None: + current = list(self.values) if not self.accumulate: self.values = list[Value]() if flat_values := flatten(values): @@ -76,6 +78,10 @@ class Topic( self.values.append(value) else: self.values.extend(flat_values) + return self.values != current def get(self) -> Sequence[Value]: - return list(self.values) + if self.values: + return list(self.values) + else: + raise EmptyChannelError diff --git a/langgraph/checkpoint/sqlite.py b/langgraph/checkpoint/sqlite.py index d7ce4e9c8..e13f14bf9 100644 --- a/langgraph/checkpoint/sqlite.py +++ b/langgraph/checkpoint/sqlite.py @@ -18,6 +18,7 @@ from langgraph.checkpoint.base import ( CheckpointTuple, SerializerProtocol, ) +from langgraph.errors import EmptyChannelError from langgraph.serde.jsonplus import JsonPlusSerializer @@ -435,11 +436,14 @@ class SqliteSaver(BaseCheckpointSaver, AbstractContextManager): def get_next_version(self, current: Optional[str], channel: BaseChannel) -> str: if current is None: - current_v = 1 + current_v = 0 else: current_v = int(current.split(".")[0]) next_v = current_v + 1 - next_h = md5(self.serde.dumps(channel.checkpoint())).hexdigest() + try: + next_h = md5(self.serde.dumps(channel.checkpoint())).hexdigest() + except EmptyChannelError: + next_h = "" return f"{next_v:032}.{next_h}" diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 29c6ebc36..f87126b3d 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -853,6 +853,9 @@ class Pregel( config, -1, for_execution=True, + get_next_version=self.checkpointer.get_next_version + if self.checkpointer + else _increment, ) # apply input writes _apply_writes( @@ -898,6 +901,9 @@ class Pregel( step, for_execution=True, manager=run_manager, + get_next_version=self.checkpointer.get_next_version + if self.checkpointer + else _increment, ) # if no more tasks, we're done @@ -1197,6 +1203,9 @@ class Pregel( config, -1, for_execution=True, + get_next_version=self.checkpointer.get_next_version + if self.checkpointer + else _increment, ) # apply input writes _apply_writes( @@ -1239,6 +1248,9 @@ class Pregel( step, for_execution=True, manager=run_manager, + get_next_version=self.checkpointer.get_next_version + if self.checkpointer + else _increment, ) # if no more tasks, we're done @@ -1630,12 +1642,12 @@ def _apply_writes( for chan, vals in pending_writes_by_channel.items(): if chan in channels: try: - channels[chan].update(vals) + updated = channels[chan].update(vals) except InvalidUpdateError as e: raise InvalidUpdateError( f"Invalid update for channel {chan} with values {vals}" ) from e - if get_next_version is not None: + if updated and get_next_version is not None: checkpoint["channel_versions"][chan] = get_next_version( max_version, channels[chan] ) @@ -1643,7 +1655,10 @@ def _apply_writes( # Channels that weren't updated in this step are notified of a new step for chan in channels: if chan not in updated_channels: - channels[chan].update([]) + if channels[chan].update([]) and get_next_version is not None: + checkpoint["channel_versions"][chan] = get_next_version( + max_version, channels[chan] + ) @overload @@ -1655,6 +1670,7 @@ def _prepare_next_tasks( config: RunnableConfig, step: int, for_execution: Literal[False], + get_next_version: Literal[None] = None, manager: Literal[None] = None, ) -> tuple[Checkpoint, list[PregelTaskDescription]]: ... @@ -1669,7 +1685,8 @@ def _prepare_next_tasks( config: RunnableConfig, step: int, for_execution: Literal[True], - manager: Union[ParentRunManager, AsyncParentRunManager], + get_next_version: Callable[[int, BaseChannel], int], + manager: Union[None, ParentRunManager, AsyncParentRunManager], ) -> tuple[Checkpoint, list[PregelExecutableTask]]: ... @@ -1683,6 +1700,7 @@ def _prepare_next_tasks( step: int, *, for_execution: bool, + get_next_version: Union[None, Callable[[int, BaseChannel], int]] = None, manager: Union[None, ParentRunManager, AsyncParentRunManager] = None, ) -> tuple[Checkpoint, Union[list[PregelTaskDescription], list[PregelExecutableTask]]]: checkpoint = copy_checkpoint(checkpoint) @@ -1810,10 +1828,18 @@ def _prepare_next_tasks( ) else: tasks.append(PregelTaskDescription(name, val)) + # Find the highest version of all channels + if checkpoint["channel_versions"]: + max_version = max(checkpoint["channel_versions"].values()) + else: + max_version = None # Consume all channels that were read if for_execution: for chan in channels_to_consume: - channels[chan].consume() + if channels[chan].consume(): + checkpoint["channel_versions"][chan] = get_next_version( + max_version, channels[chan] + ) return checkpoint, tasks diff --git a/tests/test_channels.py b/tests/test_channels.py index 313b4373a..bfa48f55d 100644 --- a/tests/test_channels.py +++ b/tests/test_channels.py @@ -56,13 +56,15 @@ def test_topic() -> None: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] - channel.update(["a", "b"]) + assert channel.update(["a", "b"]) assert channel.get() == ["a", "b"] - channel.update([["c", "d"], "d"]) + assert channel.update([["c", "d"], "d"]) assert channel.get() == ["c", "d", "d"] - channel.update([]) - assert channel.get() == [] - channel.update(["e"]) + assert channel.update([]) + with pytest.raises(EmptyChannelError): + channel.get() + assert not channel.update([]), "channel already empty" + assert channel.update(["e"]) assert channel.get() == ["e"] checkpoint = channel.checkpoint() with Topic(str).from_checkpoint(checkpoint) as channel: @@ -78,13 +80,15 @@ async def test_topic_async() -> None: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] - channel.update(["a", "b"]) + assert channel.update(["a", "b"]) assert channel.get() == ["a", "b"] - channel.update(["b", ["c", "d"], "d"]) + assert channel.update(["b", ["c", "d"], "d"]) assert channel.get() == ["b", "c", "d", "d"] - channel.update([]) - assert channel.get() == [] - channel.update(["e"]) + assert channel.update([]) + with pytest.raises(EmptyChannelError): + channel.get() + assert not channel.update([]), "channel already empty" + assert channel.update(["e"]) assert channel.get() == ["e"] checkpoint = channel.checkpoint() async with Topic(str).afrom_checkpoint(checkpoint) as channel: @@ -96,18 +100,20 @@ def test_topic_unique() -> None: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] - channel.update(["a", "b"]) + assert channel.update(["a", "b"]) assert channel.get() == ["a", "b"] - channel.update(["b", ["c", "d"], "d"]) + assert channel.update(["b", ["c", "d"], "d"]) assert channel.get() == ["c", "d"], "de-dupes from current and previous steps" - channel.update([]) - assert channel.get() == [] - channel.update(["e"]) + assert channel.update([]) + with pytest.raises(EmptyChannelError): + channel.get() + assert not channel.update([]), "channel already empty" + assert channel.update(["e"]) assert channel.get() == ["e"] checkpoint = channel.checkpoint() with Topic(str, unique=True).from_checkpoint(checkpoint) as channel: assert channel.get() == ["e"] - channel.update(["d", "f"]) + assert channel.update(["d", "f"]) assert channel.get() == ["f"], "de-dupes from checkpoint" @@ -116,18 +122,20 @@ async def test_topic_unique_async() -> None: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] - channel.update(["a", "b"]) + assert channel.update(["a", "b"]) assert channel.get() == ["a", "b"] - channel.update(["b", ["c", "d"], "d"]) + assert channel.update(["b", ["c", "d"], "d"]) assert channel.get() == ["c", "d"], "de-dupes from current and previous steps" - channel.update([]) - assert channel.get() == [] - channel.update(["e"]) + assert channel.update([]) + with pytest.raises(EmptyChannelError): + channel.get() + assert not channel.update([]), "channel already empty" + assert channel.update(["e"]) assert channel.get() == ["e"] checkpoint = channel.checkpoint() async with Topic(str, unique=True).afrom_checkpoint(checkpoint) as channel: assert channel.get() == ["e"] - channel.update(["d", "f"]) + assert channel.update(["d", "f"]) assert channel.get() == ["f"], "de-dupes from checkpoint" @@ -136,16 +144,16 @@ def test_topic_accumulate() -> None: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] - channel.update(["a", "b"]) + assert channel.update(["a", "b"]) assert channel.get() == ["a", "b"] - channel.update(["b", ["c", "d"], "d"]) + assert channel.update(["b", ["c", "d"], "d"]) assert channel.get() == ["a", "b", "b", "c", "d", "d"] - channel.update([]) + assert not channel.update([]) assert channel.get() == ["a", "b", "b", "c", "d", "d"] checkpoint = channel.checkpoint() with Topic(str, accumulate=True).from_checkpoint(checkpoint) as channel: assert channel.get() == ["a", "b", "b", "c", "d", "d"] - channel.update(["e"]) + assert channel.update(["e"]) assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"] @@ -154,16 +162,16 @@ async def test_topic_accumulate_async() -> None: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] - channel.update(["a", "b"]) + assert channel.update(["a", "b"]) assert channel.get() == ["a", "b"] - channel.update(["b", ["c", "d"], "d"]) + assert channel.update(["b", ["c", "d"], "d"]) assert channel.get() == ["a", "b", "b", "c", "d", "d"] - channel.update([]) + assert not channel.update([]) assert channel.get() == ["a", "b", "b", "c", "d", "d"] checkpoint = channel.checkpoint() 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.update(["e"]) assert channel.get() == ["a", "b", "b", "c", "d", "d", "e"] @@ -172,18 +180,19 @@ def test_topic_unique_accumulate() -> None: assert channel.ValueType is Sequence[str] assert channel.UpdateType is Union[str, list[str]] - channel.update(["a", "b"]) + assert channel.update(["a", "b"]) assert channel.get() == ["a", "b"] - channel.update(["b", ["c", "d"], "d"]) + assert channel.update(["b", ["c", "d"], "d"]) assert channel.get() == ["a", "b", "c", "d"] - channel.update([]) + assert not channel.update(["c"]), "no new values" + assert not channel.update([]) assert channel.get() == ["a", "b", "c", "d"] checkpoint = channel.checkpoint() 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.update(["d", "e"]) assert channel.get() == ["a", "b", "c", "d", "e"] diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 810925bb7..b6f06eac1 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -619,7 +619,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: ] assert [*app.stream({"input": 2, "inbox": 12})] == [ {"inbox": [3], "output": 13}, - {"inbox": [], "output": 4}, + {"output": 4}, ] assert [*app.stream({"input": 2, "inbox": 12}, stream_mode="debug")] == [ { @@ -1216,7 +1216,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: if i == 0: assert chunk == {"inbox": [3]} elif i == 1: - assert chunk == {"inbox": [], "output": 4} + assert chunk == {"output": 4} else: assert False, "Expected only two chunks" assert cleanup.call_count == 1, "Expected cleanup to be called once" diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 3517c0c13..273f16d2c 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -684,7 +684,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: ] assert [c async for c in app.astream({"input": 2, "inbox": 12})] == [ {"inbox": [3], "output": 13}, - {"inbox": [], "output": 4}, + {"output": 4}, ] assert [ c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="debug") @@ -1301,7 +1301,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: if i == 0: assert chunk == {"inbox": [3]} elif i == 1: - assert chunk == {"inbox": [], "output": 4} + assert chunk == {"output": 4} else: assert False, "Expected only two chunks" assert setup_sync.call_count == 0 From 4ff56edc26adf3cb45fdeee03bde579178baae24 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Fri, 14 Jun 2024 11:15:08 -0700 Subject: [PATCH 2/2] Update docstrings --- langgraph/channels/base.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/langgraph/channels/base.py b/langgraph/channels/base.py index 8133edf70..bb88ba2c6 100644 --- a/langgraph/channels/base.py +++ b/langgraph/channels/base.py @@ -63,7 +63,8 @@ class BaseChannel(Generic[Value, Update, C], ABC): The order of the updates in the sequence is arbitrary. This method is called by Pregel for all channels at the end of each step. If there are no updates, it is called with an empty sequence. - Raises InvalidUpdateError if the sequence of updates is invalid.""" + Raises InvalidUpdateError if the sequence of updates is invalid. + Returns True if the channel was updated, False otherwise.""" @abstractmethod def get(self) -> Value: @@ -74,7 +75,7 @@ class BaseChannel(Generic[Value, Update, C], ABC): def consume(self) -> bool: """Mark the current value of the channel as consumed. By default, no-op. This is called by Pregel before the start of the next step, for all - channels that triggered a node. + channels that triggered a node. If the channel was updated, return True. """ return False