diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 166dd5274..1d4031df3 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -971,9 +971,8 @@ class Pregel( **kwargs: Any, ) -> Union[dict[str, Any], Any]: output_keys = output_keys if output_keys is not None else self.output_channels - output_is_dict = not isinstance(output_keys, str) if stream_mode == "values": - latest: Union[dict[str, Any], Any] = {} if output_is_dict else None + latest: Union[dict[str, Any], Any] = None else: chunks = [] for chunk in self.stream( @@ -988,7 +987,7 @@ class Pregel( **kwargs, ): if stream_mode == "values": - latest = {**latest, **chunk} if output_is_dict else chunk + latest = chunk else: chunks.append(chunk) if stream_mode == "values": @@ -1010,9 +1009,8 @@ class Pregel( **kwargs: Any, ) -> Union[dict[str, Any], Any]: output_keys = output_keys if output_keys is not None else self.output_channels - output_is_dict = not isinstance(output_keys, str) if stream_mode == "values": - latest: Union[dict[str, Any], Any] = {} if output_is_dict else None + latest: Union[dict[str, Any], Any] = None else: chunks = [] async for chunk in self.astream( @@ -1027,7 +1025,7 @@ class Pregel( **kwargs, ): if stream_mode == "values": - latest = {**latest, **chunk} if output_is_dict else chunk + latest = chunk else: chunks.append(chunk) if stream_mode == "values": diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index cdce903c7..4687a5af6 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -71,8 +71,8 @@ def map_output_values( if any(chan == output_channels for chan, _ in pending_writes): yield read_channel(channels, output_channels) else: - if updated := {c for c, _ in pending_writes if c in output_channels}: - yield read_channels(channels, updated) + if {c for c, _ in pending_writes if c in output_channels}: + yield read_channels(channels, output_channels) def map_output_updates( diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 8fbd10c19..4f901d25c 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -272,6 +272,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: nodes={"one": one, "two": two}, channels={"inbox": Topic(int)}, input_channels=["input", "inbox"], + stream_channels=["output", "inbox"], ) # [12 + 1, 2 + 1 + 1] @@ -294,7 +295,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: ] assert [*app.stream({"input": 2, "inbox": 12})] == [ {"inbox": [3], "output": 13}, - {"output": 4}, + {"inbox": [], "output": 4}, ] @@ -637,13 +638,16 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: ) two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") - app = Pregel(nodes={"one": one, "two": two}) + app = Pregel(nodes={"one": one, "two": two}, stream_channels=["output", "between"]) assert [c for c in app.stream(2, stream_mode="updates")] == [ {"one": {"between": 3, "output": 3}}, {"two": {"output": 4}}, ] - assert [c for c in app.stream(2)] == [{"between": 3, "output": 3}, {"output": 4}] + assert [c for c in app.stream(2)] == [ + {"between": 3, "output": 3}, + {"between": 3, "output": 4}, + ] def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: @@ -695,6 +699,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: "ctx": Context(an_int, typ=int), }, output_channels=["inbox", "output"], + stream_channels=["inbox", "output"], ) assert setup.call_count == 0 @@ -705,7 +710,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: if i == 0: assert chunk == {"inbox": [3]} elif i == 1: - assert chunk == {"output": 4} + assert chunk == {"inbox": [], "output": 4} else: assert False, "Expected only two chunks" assert cleanup.call_count == 1, "Expected cleanup to be called once" @@ -1825,6 +1830,7 @@ def test_conditional_entrypoint_graph_state(snapshot: SnapshotAssertion) -> None assert app.invoke({"input": "what is weather in sf"}) == { "input": "what is weather in sf", "output": "what is weather in sf->right", + "steps": [], } assert [*app.stream({"input": "what is weather in sf"})] == [ @@ -2887,10 +2893,17 @@ def test_in_one_fan_out_out_one_graph_state() -> None: ] assert [*app.stream({"query": "what is weather in sf"}, stream_mode="values")] == [ - {"query": "what is weather in sf"}, - {"query": "query: what is weather in sf"}, - {"docs": ["doc1", "doc2", "doc3", "doc4"]}, - {"answer": "doc1,doc2,doc3,doc4"}, + {"query": "what is weather in sf", "docs": []}, + {"query": "query: what is weather in sf", "docs": []}, + { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + }, ] diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 8e8d012d9..eee44f555 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -164,7 +164,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - app = Pregel(nodes={"one": one, "two": two}) + app = Pregel(nodes={"one": one, "two": two}, stream_channels=["inbox", "output"]) assert await app.ainvoke(2) == 4 @@ -182,6 +182,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: } elif step == 2: assert values == { + "inbox": 3, "output": 4, } assert step == 2 @@ -280,6 +281,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: nodes={"one": one, "two": two}, channels={"inbox": Topic(int)}, input_channels=["input", "inbox"], + stream_channels=["inbox", "output"], ) # [12 + 1, 2 + 1 + 1] @@ -304,7 +306,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: ] assert [c async for c in app.astream({"input": 2, "inbox": 12})] == [ {"inbox": [3], "output": 13}, - {"output": 4}, + {"inbox": [], "output": 4}, ] @@ -669,12 +671,12 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non ) two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") - app = Pregel(nodes={"one": one, "two": two}) + app = Pregel(nodes={"one": one, "two": two}, stream_channels=["output", "between"]) # Then invoke pubsub assert [c async for c in app.astream(2)] == [ {"between": 3, "output": 3}, - {"output": 4}, + {"between": 3, "output": 4}, ] @@ -727,6 +729,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: "ctx": Context(an_int, an_int_async, typ=int), }, output_channels=["inbox", "output"], + stream_channels=["inbox", "output"], ) async def aenumerate(aiter: AsyncIterator[Any]) -> AsyncIterator[tuple[int, Any]]: @@ -747,7 +750,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: if i == 0: assert chunk == {"inbox": [3]} elif i == 1: - assert chunk == {"output": 4} + assert chunk == {"inbox": [], "output": 4} else: assert False, "Expected only two chunks" assert setup_sync.call_count == 0 @@ -1956,6 +1959,7 @@ async def test_conditional_entrypoint_graph_state() -> None: assert await app.ainvoke({"input": "what is weather in sf"}) == { "input": "what is weather in sf", "output": "what is weather in sf->right", + "steps": [], } assert [c async for c in app.astream({"input": "what is weather in sf"})] == [ @@ -2653,6 +2657,25 @@ async def test_in_one_fan_out_out_one_graph_state() -> None: {"qa": {"answer": "doc1,doc2,doc3,doc4"}}, ] + assert [ + c + async for c in app.astream( + {"query": "what is weather in sf"}, stream_mode="values" + ) + ] == [ + {"query": "what is weather in sf", "docs": []}, + {"query": "query: what is weather in sf", "docs": []}, + { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + }, + { + "query": "query: what is weather in sf", + "docs": ["doc1", "doc2", "doc3", "doc4"], + "answer": "doc1,doc2,doc3,doc4", + }, + ] + @pytest.mark.parametrize( "checkpoint_at", [CheckpointAt.END_OF_RUN, CheckpointAt.END_OF_STEP]