Merge pull request #1499 from langchain-ai/nc/27aug/print-output-skipped-tasks

Emit output for skipped tasks, with `cached` mark

libs/langgraph/langgraph/pregel/io.py

@@ -96,6 +96,7 @@ class AddableUpdatesDict(AddableDict):
 def map_output_updates(
     output_channels: Union[str, Sequence[str]],
     tasks: list[tuple[PregelExecutableTask, Sequence[tuple[str, Any]]]],
+    cached: bool = False,
 ) -> Iterator[dict[str, Union[Any, dict[str, Any]]]]:
     """Map pending writes (a sequence of tuples (channel, value)) to output chunk."""
     output_tasks = [
@@ -130,6 +131,8 @@ def map_output_updates(
             grouped[node] = None
         if len(value) == 1:
             grouped[node] = value[0]
+    if cached:
+        grouped["__metadata__"] = {"cached": cached}
     yield AddableUpdatesDict(grouped)
 
 

libs/langgraph/langgraph/pregel/loop.py

@@ -205,17 +205,23 @@ class PregelLoop:
                 writes,
                 task_id,
             )
+        self._output_writes(task_id, writes)
+
+    def _output_writes(
+        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.stream.extend(
                 (self.config["configurable"].get("checkpoint_ns", ""), "updates", v)
-                for v in map_output_updates(self.output_keys, [(task, writes)])
+                for v in map_output_updates(self.output_keys, [(task, writes)], cached)
             )
-            self.stream.extend(
-                (self.config["configurable"].get("checkpoint_ns", ""), "debug", v)
-                for v in map_debug_task_results(
-                    self.step, [(task, writes)], self.stream_keys
+            if not cached:
+                self.stream.extend(
+                    (self.config["configurable"].get("checkpoint_ns", ""), "debug", v)
+                    for v in map_debug_task_results(
+                        self.step, [(task, writes)], self.stream_keys
+                    )
                 )
-            )
 
     def tick(
         self,
@@ -320,6 +326,10 @@ class PregelLoop:
                     continue
                 if task := next((t for t in self.tasks if t.id == tid), None):
                     task.writes.append((k, v))
+            # print output for any tasks we applied previous writes to
+            for task in self.tasks:
+                if task.writes:
+                    self._output_writes(task.id, task.writes, cached=True)
 
         # if all tasks have finished, re-tick
         if all(task.writes for task in self.tasks):

libs/langgraph/tests/test_pregel.py

@@ -798,8 +798,12 @@ def test_invoke_two_processes_in_out_interrupt(
 
     # re-running from any previous checkpoint w/out forking should do nothing
     assert [c for c in app.stream(None, history[0].config, stream_mode="updates")] == []
-    assert [c for c in app.stream(None, history[1].config, stream_mode="updates")] == []
-    assert [c for c in app.stream(None, history[2].config, stream_mode="updates")] == []
+    assert [c for c in app.stream(None, history[1].config, stream_mode="updates")] == [
+        {"two": {"output": 5}, "__metadata__": {"cached": True}},
+    ]
+    assert [c for c in app.stream(None, history[2].config, stream_mode="updates")] == [
+        {"one": {"inbox": 4}, "__metadata__": {"cached": True}},
+    ]
 
     # forking and re-running from any prev checkpoint should re-run nodes
     fork_config = app.update_state(history[0].config, None)
@@ -987,7 +991,7 @@ def test_fork_always_re_runs_nodes(
     ] == []
     assert [
         c for c in graph.stream(None, history[1].config, stream_mode="updates")
-    ] == []
+    ] == [{"add_one": 1, "__metadata__": {"cached": True}}]
 
     # forking and re-running from any prev checkpoint should re-run nodes
     fork_config = graph.update_state(history[0].config, None)

libs/langgraph/tests/test_pregel_async.py

@@ -1025,10 +1025,14 @@ async def test_invoke_two_processes_in_out_interrupt(
     ] == []
     assert [
         c async for c in app.astream(None, history[1].config, stream_mode="updates")
-    ] == []
+    ] == [
+        {"two": {"output": 5}, "__metadata__": {"cached": True}},
+    ]
     assert [
         c async for c in app.astream(None, history[2].config, stream_mode="updates")
-    ] == []
+    ] == [
+        {"one": {"inbox": 4}, "__metadata__": {"cached": True}},
+    ]
 
     # forking and re-running from any prev checkpoint should re-run nodes
     fork_config = await app.aupdate_state(history[0].config, None)
@@ -1220,7 +1224,7 @@ async def test_fork_always_re_runs_nodes(
     ] == []
     assert [
         c async for c in graph.astream(None, history[1].config, stream_mode="updates")
-    ] == []
+    ] == [{"add_one": 1, "__metadata__": {"cached": True}}]
 
     # forking and re-running from any prev checkpoint should re-run nodes
     fork_config = await graph.aupdate_state(history[0].config, None)