Add support for allow_falsy_output parameter in Pregel class

langgraph/pregel/__init__.py

@@ -557,6 +557,7 @@ class Pregel(
         input_keys: Optional[Union[str, Sequence[str]]] = None,
         interrupt_before_nodes: Optional[Sequence[str]] = None,
         interrupt_after_nodes: Optional[Sequence[str]] = None,
+        allow_falsy_output: bool = False,
         debug: Optional[bool] = None,
     ) -> Iterator[Union[dict[str, Any], Any]]:
         config = ensure_config(config)
@@ -698,12 +699,16 @@ class Pregel(
 
                     # yield current value or updates
                     if stream_mode == "values":
-                        if step_output := map_output_values(
-                            output_keys, pending_writes, channels
-                        ):
+                        if (
+                            step_output := map_output_values(
+                                output_keys, pending_writes, channels
+                            )
+                        ) or allow_falsy_output:
                             yield step_output
                     else:
-                        if step_output := map_output_updates(output_keys, next_tasks):
+                        if (
+                            step_output := map_output_updates(output_keys, next_tasks)
+                        ) or allow_falsy_output:
                             yield step_output
 
                     # save end of step checkpoint
@@ -755,6 +760,7 @@ class Pregel(
         input_keys: Optional[Union[str, Sequence[str]]] = None,
         interrupt_before_nodes: Optional[Sequence[str]] = None,
         interrupt_after_nodes: Optional[Sequence[str]] = None,
+        allow_falsy_output: bool = False,
         debug: Optional[bool] = None,
     ) -> AsyncIterator[Union[dict[str, Any], Any]]:
         config = ensure_config(config)
@@ -912,12 +918,16 @@ class Pregel(
 
                     # yield current value or updates
                     if stream_mode == "values":
-                        if step_output := map_output_values(
-                            output_keys, pending_writes, channels
-                        ):
+                        if (
+                            step_output := map_output_values(
+                                output_keys, pending_writes, channels
+                            )
+                        ) or allow_falsy_output:
                             yield step_output
                     else:
-                        if step_output := map_output_updates(output_keys, next_tasks):
+                        if (
+                            step_output := map_output_updates(output_keys, next_tasks)
+                        ) or allow_falsy_output:
                             yield step_output
 
                     # save end of step checkpoint
@@ -969,6 +979,7 @@ class Pregel(
         input_keys: Optional[Union[str, Sequence[str]]] = None,
         interrupt_before_nodes: Optional[Sequence[str]] = None,
         interrupt_after_nodes: Optional[Sequence[str]] = None,
+        allow_falsy_output: bool = False,
         debug: Optional[bool] = None,
         **kwargs: Any,
     ) -> Union[dict[str, Any], Any]:
@@ -986,6 +997,7 @@ class Pregel(
             input_keys=input_keys,
             interrupt_before_nodes=interrupt_before_nodes,
             interrupt_after_nodes=interrupt_after_nodes,
+            allow_falsy_output=allow_falsy_output,
             debug=debug,
             **kwargs,
         ):
@@ -1008,6 +1020,7 @@ class Pregel(
         input_keys: Optional[Union[str, Sequence[str]]] = None,
         interrupt_before_nodes: Optional[Sequence[str]] = None,
         interrupt_after_nodes: Optional[Sequence[str]] = None,
+        allow_falsy_output: bool = False,
         debug: Optional[bool] = None,
         **kwargs: Any,
     ) -> Union[dict[str, Any], Any]:
@@ -1025,6 +1038,7 @@ class Pregel(
             input_keys=input_keys,
             interrupt_before_nodes=interrupt_before_nodes,
             interrupt_after_nodes=interrupt_after_nodes,
+            allow_falsy_output=allow_falsy_output,
             debug=debug,
             **kwargs,
         ):

tests/test_pregel.py

@@ -68,6 +68,20 @@ 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("add_one", lambda *args, **kwargs: falsy_value)
+    graph.set_entry_point("add_one")
+    graph.set_finish_point("add_one")
+    gapp = graph.compile()
+    assert gapp.invoke(1, allow_falsy_output=True) == falsy_value
+    assert gapp.invoke(1) is None
+
+
 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,20 @@ 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("add_one", lambda *args, **kwargs: falsy_value)
+    graph.set_entry_point("add_one")
+    graph.set_finish_point("add_one")
+    gapp = graph.compile()
+    assert (await gapp.ainvoke(1, allow_falsy_output=True)) == falsy_value
+    assert (await gapp.ainvoke(1)) is None
+
+
 async def test_invoke_single_process_in_out_implicit_channels(
     mocker: MockerFixture,
 ) -> None: