From 5799b80261c83db3a74c2a152fed314b24a9dff0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sun, 7 Jan 2024 19:53:28 -0800 Subject: [PATCH] .stream() defaults to yielding output from all channels --- langgraph/pregel/__init__.py | 18 ++++++++++++++---- tests/test_pregel.py | 27 +++++++++++++++------------ tests/test_pregel_async.py | 30 ++++++++++++++++++------------ 3 files changed, 47 insertions(+), 28 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 20b8cc604..33a0e49e3 100644 --- a/langgraph/pregel/__init__.py +++ b/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 diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 30b9bd11e..47c4ba321 100644 --- a/tests/test_pregel.py +++ b/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: diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index e65835931..57b538398 100644 --- a/tests/test_pregel_async.py +++ b/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: