Speed improvements to local_read, emit output

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:

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(

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)

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
                     )
                 )