diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index c766b6873..166dd5274 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -698,13 +698,11 @@ class Pregel( # yield current value or updates if stream_mode == "values": - if step_output := map_output_values( + yield from map_output_values( output_keys, pending_writes, channels - ): - yield step_output + ) else: - if step_output := map_output_updates(output_keys, next_tasks): - yield step_output + yield from map_output_updates(output_keys, next_tasks) # save end of step checkpoint if self.checkpointer is not None and ( @@ -912,13 +910,13 @@ class Pregel( # yield current value or updates if stream_mode == "values": - if step_output := map_output_values( + for chunk in map_output_values( output_keys, pending_writes, channels ): - yield step_output + yield chunk else: - if step_output := map_output_updates(output_keys, next_tasks): - yield step_output + for chunk in map_output_updates(output_keys, next_tasks): + yield chunk # save end of step checkpoint if self.checkpointer is not None and ( diff --git a/langgraph/pregel/io.py b/langgraph/pregel/io.py index a59089384..cdce903c7 100644 --- a/langgraph/pregel/io.py +++ b/langgraph/pregel/io.py @@ -65,21 +65,20 @@ def map_output_values( output_channels: Union[str, Sequence[str]], pending_writes: Sequence[tuple[str, Any]], channels: Mapping[str, BaseChannel], -) -> Optional[Union[dict[str, Any], Any]]: +) -> Iterator[Union[dict[str, Any], Any]]: """Map pending writes (a sequence of tuples (channel, value)) to output chunk.""" if isinstance(output_channels, str): if any(chan == output_channels for chan, _ in pending_writes): - return read_channel(channels, output_channels) + yield read_channel(channels, output_channels) else: if updated := {c for c, _ in pending_writes if c in output_channels}: - return read_channels(channels, updated) - return None + yield read_channels(channels, updated) def map_output_updates( output_channels: Union[str, Sequence[str]], tasks: list[PregelExecutableTask], -) -> Optional[dict[str, Union[Any, dict[str, Any]]]]: +) -> Iterator[dict[str, Union[Any, dict[str, Any]]]]: """Map pending writes (a sequence of tuples (channel, value)) to output chunk.""" output_tasks = [ t for t in tasks if not t.config or TAG_HIDDEN not in t.config.get("tags") @@ -91,12 +90,11 @@ def map_output_updates( for chan, value in writes if chan == output_channels }: - return updated + yield updated else: if updated := { node: {chan: value for chan, value in writes if chan in output_channels} for node, _, _, writes, _ in output_tasks if any(chan in output_channels for chan, _ in writes) }: - return updated - return None + yield updated diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 6be844479..8654a6720 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -68,6 +68,19 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: assert gapp.invoke(2, debug=True) == 3 +@pytest.mark.parametrize( + "falsy_value", + [None, False, 0, "", [], {}, set(), frozenset(), 0.0, 0j], +) +def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None: + graph = Graph() + graph.add_node("return_falsy_const", lambda *args, **kwargs: falsy_value) + graph.set_entry_point("return_falsy_const") + graph.set_finish_point("return_falsy_const") + gapp = graph.compile() + assert gapp.invoke(1) == falsy_value + + def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 9d8e0e971..936ef0e55 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -65,6 +65,19 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: assert await gapp.ainvoke(2) == 3 +@pytest.mark.parametrize( + "falsy_value", + [None, False, 0, "", [], {}, set(), frozenset(), 0.0, 0j], +) +async def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None: + graph = Graph() + graph.add_node("return_falsy_const", lambda *args, **kwargs: falsy_value) + graph.set_entry_point("return_falsy_const") + graph.set_finish_point("return_falsy_const") + gapp = graph.compile() + assert falsy_value == await gapp.ainvoke(1) + + async def test_invoke_single_process_in_out_implicit_channels( mocker: MockerFixture, ) -> None: