diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index c766b6873..5c19a09fb 100644 --- a/langgraph/pregel/__init__.py +++ b/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, ): diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 6be844479..392ce7c80 100644 --- a/tests/test_pregel.py +++ b/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") diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 9d8e0e971..2c84b34c2 100644 --- a/tests/test_pregel_async.py +++ b/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: