diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 6c9205248..acc6574e8 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/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) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 87551b83f..be92c9018 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/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): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 269e2bbcc..4aefed25a 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/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) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index b24848e6d..b12209977 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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)