diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 76df9f591..aef1a69d0 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -65,7 +65,7 @@ from langgraph.pregel.debug import print_checkpoint, print_step_start from langgraph.pregel.io import map_input, map_output from langgraph.pregel.log import logger from langgraph.pregel.read import ChannelBatch, ChannelInvoke -from langgraph.pregel.reserved import ReservedChannels +from langgraph.pregel.reserved import AllReservedChannels, ReservedChannels from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry @@ -274,7 +274,8 @@ class Pregel( return ( [self.snapshot_channels] if isinstance(self.snapshot_channels, str) - else self.snapshot_channels or list(self.channels.keys()) + else self.snapshot_channels + or [k for k in self.channels if k not in AllReservedChannels] ) def get_state(self, config: RunnableConfig) -> StateSnapshot: @@ -291,7 +292,6 @@ class Pregel( k: _read_channel(channels, k) for k in channels if k in self.snapshot_channels_list - and k not in [k.value for k in ReservedChannels] } return StateSnapshot( values[self.snapshot_channels] @@ -314,7 +314,6 @@ class Pregel( k: _read_channel(channels, k) for k in channels if k in self.snapshot_channels_list - and k not in [k.value for k in ReservedChannels] } return StateSnapshot( values[self.snapshot_channels] @@ -947,7 +946,7 @@ def _apply_writes( pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) # Group writes by channel for chan, val in pending_writes: - if chan in [c.value for c in ReservedChannels]: + if chan in AllReservedChannels: raise ValueError(f"Can't write to reserved channel {chan}") pending_writes_by_channel[chan].append(val) diff --git a/langgraph/pregel/reserved.py b/langgraph/pregel/reserved.py index b6e66b945..2fad3b348 100644 --- a/langgraph/pregel/reserved.py +++ b/langgraph/pregel/reserved.py @@ -6,3 +6,6 @@ class ReservedChannels(StrEnum): is_last_step = "is_last_step" """A channel that is True if the current step is the last step, False otherwise.""" + + +AllReservedChannels = {channel.value for channel in ReservedChannels}