.stream() defaults to yielding output from all channels

langgraph/pregel/__init__.py

@@ -222,7 +222,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         if config["recursion_limit"] < 1:
             raise ValueError("recursion_limit must be at least 1")
         # assign defaults
-        output = output if output is not None else self.output
+        output = output if output is not None else [chan for chan in self.channels]
         # copy nodes to ignore mutations during execution
         processes = {**self.nodes}
         # get checkpoint from saver, or create an empty one
@@ -339,7 +339,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
             None,
         )
         # assign defaults
-        output = output if output is not None else self.output
+        output = output if output is not None else [chan for chan in self.channels]
         # copy nodes to ignore mutations during execution
         processes = {**self.nodes}
         # get checkpoint from saver, or create an empty one
@@ -448,7 +448,12 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         **kwargs: Any,
     ) -> dict[str, Any] | Any:
         latest: dict[str, Any] | Any = None
-        for chunk in self.stream(input, config, output=output, **kwargs):
+        for chunk in self.stream(
+            input,
+            config,
+            output=output if output is not None else self.output,
+            **kwargs,
+        ):
             latest = chunk
         return latest
 
@@ -498,7 +503,12 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
         **kwargs: Any,
     ) -> dict[str, Any] | Any:
         latest: dict[str, Any] | Any = None
-        async for chunk in self.astream(input, config, output=output, **kwargs):
+        async for chunk in self.astream(
+            input,
+            config,
+            output=output if output is not None else self.output,
+            **kwargs,
+        ):
             latest = chunk
         return latest
 

tests/test_pregel.py

@@ -162,14 +162,14 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "inbox": 3,
                 "input": 2,
             }
-            assert output is None
+            assert output == {"inbox": 3}
         elif view.step == 2:
             assert view.values == {
                 "output": 4,
                 "inbox": 3,
                 "input": 2,
             }
-            assert output == 4
+            assert output == {"output": 4}
 
     for output, view in app.step(2):
         if view.step == 1:
@@ -177,7 +177,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "inbox": 3,
                 "input": 2,
             }
-            assert output is None
+            assert output == {"inbox": 3}
             # modify inbox value
             view.values["inbox"] = 5
         elif view.step == 2:
@@ -187,7 +187,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "input": 2,
             }
             # output is different now
-            assert output == 6
+            assert output == {"output": 6}
 
     graph = Graph()
     graph.add_node("add_one", add_one)
@@ -205,14 +205,14 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "add_one": 2,
                 "add_one_more": 3,
             }
-            assert output is None
+            assert output == {"add_one_more": 3}
         elif view.step == 2:
             assert view.values == {
                 "add_one": 2,
                 "add_one_more": 3,
                 "__end__": 4,
             }
-            assert output == 4
+            assert output == {"__end__": 4}
 
     for output, view in gapp.step(2):
         if view.step == 1:
@@ -220,7 +220,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "add_one": 2,
                 "add_one_more": 3,
             }
-            assert output is None
+            assert output == {"add_one_more": 3}
             # modify inbox value
             view.values["add_one_more"] = 5
         elif view.step == 2:
@@ -230,7 +230,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "__end__": 6,
             }
             # output is different now
-            assert output == 6
+            assert output == {"__end__": 6}
 
 
 def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
@@ -244,9 +244,12 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
         input=["input", "inbox"],
     )
 
-    assert [*app.stream({"input": 2, "inbox": 12})] == [13, 4]  # [12 + 1, 2 + 1 + 1]
-    assert [*app.stream({"input": 2, "inbox": 12}, output=["output"])] == [
-        {"output": 13},
+    assert [*app.stream({"input": 2, "inbox": 12}, output="output")] == [
+        13,
+        4,
+    ]  # [12 + 1, 2 + 1 + 1]
+    assert [*app.stream({"input": 2, "inbox": 12})] == [
+        {"inbox": [3], "output": 13},
         {"output": 4},
     ]
 
@@ -493,7 +496,7 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
 
     app = Pregel(nodes={"one": one, "two": two})
 
-    assert [c for c in app.stream(2)] == [3, 4]
+    assert [c for c in app.stream(2)] == [{"between": 3, "output": 3}, {"output": 4}]
 
 
 def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:

tests/test_pregel_async.py

@@ -163,14 +163,14 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "inbox": 3,
                 "input": 2,
             }
-            assert output is None
+            assert output == {"inbox": 3}
         elif view.step == 2:
             assert view.values == {
                 "output": 4,
                 "inbox": 3,
                 "input": 2,
             }
-            assert output == 4
+            assert output == {"output": 4}
 
     async for output, view in app.astep(2):
         if view.step == 1:
@@ -178,7 +178,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "inbox": 3,
                 "input": 2,
             }
-            assert output is None
+            assert output == {"inbox": 3}
             # modify inbox value
             view.values["inbox"] = 5
         elif view.step == 2:
@@ -188,7 +188,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "input": 2,
             }
             # output is different now
-            assert output == 6
+            assert output == {"output": 6}
 
     graph = Graph()
     graph.add_node("add_one", add_one)
@@ -206,14 +206,14 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "add_one": 2,
                 "add_one_more": 3,
             }
-            assert output is None
+            assert output == {"add_one_more": 3}
         elif view.step == 2:
             assert view.values == {
                 "add_one": 2,
                 "add_one_more": 3,
                 "__end__": 4,
             }
-            assert output == 4
+            assert output == {"__end__": 4}
 
     async for output, view in gapp.astep(2):
         if view.step == 1:
@@ -221,7 +221,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "add_one": 2,
                 "add_one_more": 3,
             }
-            assert output is None
+            assert output == {"add_one_more": 3}
             # modify inbox value
             view.values["add_one_more"] = 5
         elif view.step == 2:
@@ -231,7 +231,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
                 "__end__": 6,
             }
             # output is different now
-            assert output == 6
+            assert output == {"__end__": 6}
 
 
 async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
@@ -246,10 +246,13 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
     )
 
     # [12 + 1, 2 + 1 + 1]
-    assert [c async for c in pubsub.astream({"input": 2, "inbox": 12})] == [13, 4]
     assert [
-        c async for c in pubsub.astream({"input": 2, "inbox": 12}, output=["output"])
-    ] == [{"output": 13}, {"output": 4}]
+        c async for c in pubsub.astream({"input": 2, "inbox": 12}, output="output")
+    ] == [13, 4]
+    assert [c async for c in pubsub.astream({"input": 2, "inbox": 12})] == [
+        {"inbox": [3], "output": 13},
+        {"output": 4},
+    ]
 
 
 async def test_batch_two_processes_in_out() -> None:
@@ -508,7 +511,10 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
     app = Pregel(nodes={"one": one, "two": two})
 
     # Then invoke pubsub
-    assert [c async for c in app.astream(2)] == [3, 4]
+    assert [c async for c in app.astream(2)] == [
+        {"between": 3, "output": 3},
+        {"output": 4},
+    ]
 
 
 async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: