diff --git a/libs/langgraph/langgraph/channels/any_value.py b/libs/langgraph/langgraph/channels/any_value.py index 58b923a03..51493ab19 100644 --- a/libs/langgraph/langgraph/channels/any_value.py +++ b/libs/langgraph/langgraph/channels/any_value.py @@ -30,9 +30,14 @@ class AnyValue(Generic[Value], BaseChannel[Value, Value, Value]): """The type of the update received by the channel.""" return self.typ + def copy(self) -> Self: + """Return a copy of the channel.""" + empty = self.__class__(self.typ, self.key) + empty.value = self.value + return empty + def from_checkpoint(self, checkpoint: Value) -> Self: - empty = self.__class__(self.typ) - empty.key = self.key + empty = self.__class__(self.typ, self.key) if checkpoint is not MISSING: empty.value = checkpoint return empty diff --git a/libs/langgraph/langgraph/channels/base.py b/libs/langgraph/langgraph/channels/base.py index 33143fffb..dc8888acd 100644 --- a/libs/langgraph/langgraph/channels/base.py +++ b/libs/langgraph/langgraph/channels/base.py @@ -30,6 +30,12 @@ class BaseChannel(Generic[Value, Update, C], ABC): # serialize/deserialize methods + def copy(self) -> Self: + """Return a copy of the channel. + By default, delegates to checkpoint() and from_checkpoint(). + Subclasses can override this method with a more efficient implementation.""" + return self.from_checkpoint(self.checkpoint()) + def checkpoint(self) -> C: """Return a serializable representation of the channel's current state. Raises EmptyChannelError if the channel is empty (never updated yet), diff --git a/libs/langgraph/langgraph/channels/binop.py b/libs/langgraph/langgraph/channels/binop.py index 59cc11550..1f95d5562 100644 --- a/libs/langgraph/langgraph/channels/binop.py +++ b/libs/langgraph/langgraph/channels/binop.py @@ -66,6 +66,13 @@ class BinaryOperatorAggregate(Generic[Value], BaseChannel[Value, Value, Value]): """The type of the update received by the channel.""" return self.typ + def copy(self) -> Self: + """Return a copy of the channel.""" + empty = self.__class__(self.typ, self.operator) + empty.key = self.key + empty.value = self.value + return empty + def from_checkpoint(self, checkpoint: Value) -> Self: empty = self.__class__(self.typ, self.operator) empty.key = self.key diff --git a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py index 511c311f3..4f75f2a8c 100644 --- a/libs/langgraph/langgraph/channels/dynamic_barrier_value.py +++ b/libs/langgraph/langgraph/channels/dynamic_barrier_value.py @@ -46,6 +46,14 @@ class DynamicBarrierValue( """The type of the update received by the channel.""" return self.typ + def copy(self) -> Self: + """Return a copy of the channel.""" + empty = self.__class__(self.typ) + empty.key = self.key + empty.names = self.names + empty.seen = self.seen.copy() + return empty + def checkpoint(self) -> tuple[Optional[set[Value]], set[Value]]: return (self.names, self.seen) diff --git a/libs/langgraph/langgraph/channels/ephemeral_value.py b/libs/langgraph/langgraph/channels/ephemeral_value.py index 29471afb4..e2beaf05d 100644 --- a/libs/langgraph/langgraph/channels/ephemeral_value.py +++ b/libs/langgraph/langgraph/channels/ephemeral_value.py @@ -30,6 +30,13 @@ class EphemeralValue(Generic[Value], BaseChannel[Value, Value, Value]): """The type of the update received by the channel.""" return self.typ + def copy(self) -> Self: + """Return a copy of the channel.""" + empty = self.__class__(self.typ, self.guard) + empty.key = self.key + empty.value = self.value + return empty + def from_checkpoint(self, checkpoint: Value) -> Self: empty = self.__class__(self.typ, self.guard) empty.key = self.key diff --git a/libs/langgraph/langgraph/channels/last_value.py b/libs/langgraph/langgraph/channels/last_value.py index 90707575c..7232b8f65 100644 --- a/libs/langgraph/langgraph/channels/last_value.py +++ b/libs/langgraph/langgraph/channels/last_value.py @@ -34,9 +34,14 @@ class LastValue(Generic[Value], BaseChannel[Value, Value, Value]): """The type of the update received by the channel.""" return self.typ + def copy(self) -> Self: + """Return a copy of the channel.""" + empty = self.__class__(self.typ, self.key) + empty.value = self.value + return empty + def from_checkpoint(self, checkpoint: Value) -> Self: - empty = self.__class__(self.typ) - empty.key = self.key + empty = self.__class__(self.typ, self.key) if checkpoint is not MISSING: empty.value = checkpoint return empty diff --git a/libs/langgraph/langgraph/channels/named_barrier_value.py b/libs/langgraph/langgraph/channels/named_barrier_value.py index 2145d0f73..1d4b32097 100644 --- a/libs/langgraph/langgraph/channels/named_barrier_value.py +++ b/libs/langgraph/langgraph/channels/named_barrier_value.py @@ -33,6 +33,13 @@ class NamedBarrierValue(Generic[Value], BaseChannel[Value, Value, set[Value]]): """The type of the update received by the channel.""" return self.typ + def copy(self) -> Self: + """Return a copy of the channel.""" + empty = self.__class__(self.typ, self.names) + empty.key = self.key + empty.seen = self.seen.copy() + return empty + def checkpoint(self) -> set[Value]: return self.seen diff --git a/libs/langgraph/langgraph/channels/topic.py b/libs/langgraph/langgraph/channels/topic.py index 8fc998353..42c665bb5 100644 --- a/libs/langgraph/langgraph/channels/topic.py +++ b/libs/langgraph/langgraph/channels/topic.py @@ -48,6 +48,13 @@ class Topic( """The type of the update received by the channel.""" return Union[self.typ, list[self.typ]] # type: ignore[name-defined] + def copy(self) -> Self: + """Return a copy of the channel.""" + empty = self.__class__(self.typ, self.accumulate) + empty.key = self.key + empty.values = self.values.copy() + return empty + def checkpoint(self) -> list[Value]: return self.values diff --git a/libs/langgraph/langgraph/channels/untracked_value.py b/libs/langgraph/langgraph/channels/untracked_value.py index 24ffbd425..035beac02 100644 --- a/libs/langgraph/langgraph/channels/untracked_value.py +++ b/libs/langgraph/langgraph/channels/untracked_value.py @@ -30,6 +30,13 @@ class UntrackedValue(Generic[Value], BaseChannel[Value, Value, Value]): """The type of the update received by the channel.""" return self.typ + def copy(self) -> Self: + """Return a copy of the channel.""" + empty = self.__class__(self.typ, self.guard) + empty.key = self.key + empty.value = self.value + return empty + def checkpoint(self) -> Value: return MISSING diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index a550de23d..8b342ed3d 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -1535,12 +1535,9 @@ class Pregel(PregelProtocol): ), CONFIG_KEY_READ: partial( local_read, - step + 1, - checkpoint, channels, managed, task, - config, ), }, ), @@ -1944,12 +1941,9 @@ class Pregel(PregelProtocol): ), CONFIG_KEY_READ: partial( local_read, - step + 1, - checkpoint, channels, managed, task, - config, ), }, ), diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index ffcf9c815..e43890498 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -33,7 +33,6 @@ from langgraph.checkpoint.base import ( Checkpoint, PendingWrite, V, - copy_checkpoint, ) from langgraph.constants import ( CONF, @@ -69,12 +68,10 @@ from langgraph.managed.base import ManagedValueMapping from langgraph.pregel.call import get_runnable_for_task from langgraph.pregel.io import read_channel, read_channels from langgraph.pregel.log import logger -from langgraph.pregel.manager import ChannelsManager from langgraph.pregel.read import PregelNode from langgraph.store.base import BaseStore from langgraph.types import ( All, - LoopProtocol, PregelExecutableTask, PregelScratchpad, PregelTask, @@ -169,39 +166,39 @@ def should_interrupt( def local_read( - step: int, - checkpoint: Checkpoint, channels: Mapping[str, BaseChannel], managed: ManagedValueMapping, task: WritesProtocol, - config: RunnableConfig, select: Union[list[str], str], fresh: bool = False, ) -> Union[dict[str, Any], Any]: """Function injected under CONFIG_KEY_READ in task config, to read current state. Used by conditional edges to read a copy of the state with reflecting the writes from that node only.""" + updated: dict[str, list[Any]] = defaultdict(list) if isinstance(select, str): managed_keys = [] - for c, _ in task.writes: + for c, v in task.writes: if c == select: - updated = {c} - break - else: - updated = set() + updated[c].append(v) else: managed_keys = [k for k in select if k in managed] select = [k for k in select if k not in managed] - updated = set(select).intersection(c for c, _ in task.writes) + for c, v in task.writes: + if c in select: + updated[c].append(v) if fresh and updated: - with ChannelsManager( - {k: v for k, v in channels.items() if k in updated}, - checkpoint, - LoopProtocol(config=config, step=step, stop=step + 1), - skip_context=True, - ) as (local_channels, _): - apply_writes(copy_checkpoint(checkpoint), local_channels, [task], None) - values = read_channels({**channels, **local_channels}, select) + # apply writes + local_channels: dict[str, BaseChannel] = {} + for k in channels: + if k in updated: + cc = channels[k].copy() + cc.update(updated[k]) + else: + cc = channels[k] + local_channels[k] = cc + # read fresh values + values = read_channels(local_channels, select) else: values = read_channels(channels, select) if managed_keys: @@ -335,6 +332,17 @@ def apply_writes( return pending_writes_by_managed, updated_channels +def has_next_tasks( + trigger_to_nodes: Mapping[str, Sequence[str]], + updated_channels: set[str], + checkpoint: Checkpoint, +) -> bool: + """Check if there are any tasks that should be run in the next step.""" + return bool(checkpoint["pending_sends"]) or not updated_channels.isdisjoint( + trigger_to_nodes + ) + + @overload def prepare_next_tasks( checkpoint: Checkpoint, @@ -562,12 +570,9 @@ def prepare_single_task( ), CONFIG_KEY_READ: partial( local_read, - step, - checkpoint, channels, managed, PregelTaskWrites(task_path[:3], name, writes, triggers), - config, ), CONFIG_KEY_STORE: (store or configurable.get(CONFIG_KEY_STORE)), CONFIG_KEY_CHECKPOINTER: ( @@ -667,14 +672,11 @@ def prepare_single_task( ), CONFIG_KEY_READ: partial( local_read, - step, - checkpoint, channels, managed, PregelTaskWrites( task_path[:3], packet.node, writes, triggers ), - config, ), CONFIG_KEY_STORE: ( store or configurable.get(CONFIG_KEY_STORE) @@ -789,8 +791,6 @@ def prepare_single_task( ), CONFIG_KEY_READ: partial( local_read, - step, - checkpoint, channels, managed, PregelTaskWrites( @@ -799,7 +799,6 @@ def prepare_single_task( writes, triggers, ), - config, ), CONFIG_KEY_STORE: ( store or configurable.get(CONFIG_KEY_STORE) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 031d4d2a2..41f959723 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -155,6 +155,8 @@ class PregelLoop(LoopProtocol): manager: Union[None, AsyncParentRunManager, ParentRunManager] interrupt_after: Union[All, Sequence[str]] interrupt_before: Union[All, Sequence[str]] + checkpoint_every_step: bool + debug: bool checkpointer_get_next_version: GetNextVersion checkpointer_put_writes: Optional[ @@ -211,6 +213,7 @@ class PregelLoop(LoopProtocol): input_model: Optional[Type[BaseModel]] = None, debug: bool = False, trigger_to_nodes: Optional[Mapping[str, Sequence[str]]] = None, + checkpoint_every_step: bool = True, ) -> None: super().__init__( step=0, @@ -235,6 +238,7 @@ class PregelLoop(LoopProtocol): or CONFIG_KEY_DEDUPE_TASKS in config[CONF] ) self.trigger_to_nodes = trigger_to_nodes + self.checkpoint_every_step = checkpoint_every_step self.debug = debug if self.stream is not None and CONFIG_KEY_STREAM in config[CONF]: self.stream = DuplexStream(self.stream, config[CONF][CONFIG_KEY_STREAM]) @@ -703,8 +707,6 @@ class PregelLoop(LoopProtocol): return updated_channels def _put_checkpoint(self, metadata: CheckpointMetadata) -> None: - for k, v in self.config["metadata"].items(): - metadata.setdefault(k, v) # type: ignore # assign step and parents metadata["step"] = self.step metadata["parents"] = self.config[CONF].get(CONFIG_KEY_CHECKPOINT_MAP, {}) @@ -719,10 +721,15 @@ class PregelLoop(LoopProtocol): else self.stream_keys ), ) - # create new checkpoint - self.checkpoint = create_checkpoint(self.checkpoint, self.channels, self.step) # bail if no checkpointer if self._checkpointer_put_after_previous is not None: + for k, v in self.config["metadata"].items(): + metadata.setdefault(k, v) # type: ignore + + # create new checkpoint + self.checkpoint = create_checkpoint( + self.checkpoint, self.channels, self.step + ) self.checkpoint_metadata = metadata self.prev_checkpoint_config = (