diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index c3dfd1b80..044ffbbf3 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -102,17 +102,25 @@ def local_read( ) -> Union[dict[str, Any], Any]: if isinstance(select, str): managed_keys = [] + for c, _ in task.writes: + if c == select: + updated = {c} + break + else: + updated = set() else: managed_keys = [k for k in select if k in managed] select = [k for k in select if k not in managed] - if fresh: - new_checkpoint = create_checkpoint(copy_checkpoint(checkpoint), channels, -1) - with ChannelsManager(channels, new_checkpoint, config, skip_context=True) as ( - channels, - _, - ): - apply_writes(new_checkpoint, channels, [task], None) - values = read_channels(channels, select) + updated = set(select).intersection(c for c, _ in task.writes) + if fresh and updated: + with ChannelsManager( + {k: v for k, v in channels.items() if k in updated}, + checkpoint, + config, + skip_context=True, + ) as (local_channels, _): + apply_writes(copy_checkpoint(checkpoint), local_channels, [task], None) + values = read_channels({**channels, **local_channels}, select) else: values = read_channels(channels, select) if managed_keys: @@ -171,7 +179,10 @@ def apply_writes( # Consume all channels that were read for chan in { - chan for task in tasks for chan in task.triggers if chan not in RESERVED + chan + for task in tasks + for chan in task.triggers + if chan not in RESERVED and chan in channels }: if channels[chan].consume(): if get_next_version is not None: diff --git a/libs/langgraph/langgraph/pregel/debug.py b/libs/langgraph/langgraph/pregel/debug.py index f24fad6e3..cb6e45e0d 100644 --- a/libs/langgraph/langgraph/pregel/debug.py +++ b/libs/langgraph/langgraph/pregel/debug.py @@ -96,33 +96,25 @@ def map_debug_tasks( def map_debug_task_results( step: int, - tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]], + task_tup: tuple[PregelExecutableTask, Sequence[tuple[str, Any]]], stream_keys: Union[str, Sequence[str]], ) -> Iterator[DebugOutputTaskResult]: stream_channels_list = ( [stream_keys] if isinstance(stream_keys, str) else stream_keys ) - ts = datetime.now(timezone.utc) - for task, writes in tasks: - if task.config is not None and TAG_HIDDEN in task.config.get("tags", []): - continue - - metadata = task.config["metadata"].copy() - metadata.pop("checkpoint_id", None) - # TODO: make task IDs deterministic in tests and reuse task IDs for payload ID - - yield { - "type": "task_result", - "timestamp": ts.isoformat(), - "step": step, - "payload": { - "id": task.id, - "name": task.name, - "error": next((w[1] for w in writes if w[0] == ERROR), None), - "result": [w for w in writes if w[0] in stream_channels_list], - "interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT], - }, - } + task, writes = task_tup + yield { + "type": "task_result", + "timestamp": datetime.now(timezone.utc).isoformat(), + "step": step, + "payload": { + "id": task.id, + "name": task.name, + "error": next((w[1] for w in writes if w[0] == ERROR), None), + "result": [w for w in writes if w[0] in stream_channels_list], + "interrupts": [asdict(w[1]) for w in writes if w[0] == INTERRUPT], + }, + } def map_debug_checkpoint( diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 8568e8781..286282f1d 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -107,26 +107,27 @@ def map_output_updates( (t, ww) for t, ww in tasks if (not t.config or TAG_HIDDEN not in t.config.get("tags")) - and all(k not in (ERROR, INTERRUPT) for k, _ in ww) + and ww[0][0] != ERROR + and ww[0][0] != INTERRUPT ] if not output_tasks: return if isinstance(output_channels, str): - updated = [ + updated = ( (task.name, value) for task, writes in output_tasks for chan, value in writes if chan == output_channels - ] + ) else: - updated = [ + updated = ( ( task.name, {chan: value for chan, value in task.writes if chan in output_channels}, ) for task, writes in output_tasks if any(chan in output_channels for chan, _ in writes) - ] + ) grouped = {t.name: [] for t, _ in output_tasks} for node, value in updated: grouped[node].append(value) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index d2052b17a..703882ab2 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -45,6 +45,7 @@ from langgraph.constants import ( ERROR, INPUT, INTERRUPT, + TAG_HIDDEN, ) from langgraph.errors import EmptyInputError, GraphInterrupt from langgraph.managed.base import ( @@ -437,14 +438,11 @@ class PregelLoop: if isinstance(self.stream_keys, str) 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: - # create new checkpoint self.checkpoint_metadata = metadata - self.checkpoint = create_checkpoint( - self.checkpoint, self.channels, self.step - ) - self.checkpoint_config = { **self.checkpoint_config, "configurable": { @@ -459,7 +457,6 @@ class PregelLoop: new_versions = get_new_channel_versions( self.checkpoint_previous_versions, channel_versions ) - self.checkpoint_previous_versions = channel_versions # save it, without blocking @@ -510,15 +507,20 @@ class PregelLoop: self, task_id: str, writes: Sequence[tuple[str, Any]], *, cached: bool = False ) -> None: if task := next((t for t in self.tasks if t.id == task_id), None): - self._emit( - (self.config["configurable"].get("checkpoint_ns", ""), "updates", v) - for v in map_output_updates(self.output_keys, [(task, writes)], cached) - ) + if task.config is not None and TAG_HIDDEN in task.config.get("tags"): + return + if writes[0][0] != ERROR and writes[0][0] != INTERRUPT: + self._emit( + (self.config["configurable"].get("checkpoint_ns", ""), "updates", v) + for v in map_output_updates( + self.output_keys, [(task, writes)], cached + ) + ) if not cached: self._emit( (self.config["configurable"].get("checkpoint_ns", ""), "debug", v) for v in map_debug_task_results( - self.step, [(task, writes)], self.stream_keys + self.step, (task, writes), self.stream_keys ) )