Merge pull request #298 from hmasdev/feature-handle-falsy-output

Pregel: Add support for falsy outputs in Pregel class

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 (

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

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")

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: