From 42de3304dcb619f2b20def34373f57a57dede182 Mon Sep 17 00:00:00 2001 From: hmasdev Date: Thu, 11 Apr 2024 19:44:21 +0900 Subject: [PATCH 1/3] Add support for allow_falsy_output parameter in Pregel class --- langgraph/pregel/__init__.py | 30 ++++++++++++++++++++++-------- tests/test_pregel.py | 14 ++++++++++++++ tests/test_pregel_async.py | 14 ++++++++++++++ 3 files changed, 50 insertions(+), 8 deletions(-) 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: From 236ecb1982084af73e46faec851b347930b3c8c1 Mon Sep 17 00:00:00 2001 From: hmasdev Date: Thu, 11 Apr 2024 20:34:45 +0900 Subject: [PATCH 2/3] rename node name in tests --- tests/test_pregel.py | 6 +++--- tests/test_pregel_async.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 392ce7c80..ad5c32a09 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -74,9 +74,9 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: ) 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") + 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, allow_falsy_output=True) == falsy_value assert gapp.invoke(1) is None diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 2c84b34c2..df1c1c024 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -71,9 +71,9 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: ) 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") + 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 (await gapp.ainvoke(1, allow_falsy_output=True)) == falsy_value assert (await gapp.ainvoke(1)) is None From 46e67cff87d2502c3da7383e0f7fac07df8346ed Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 11 Apr 2024 08:44:41 -0700 Subject: [PATCH 3/3] Remove flag --- langgraph/pregel/__init__.py | 36 ++++++++++-------------------------- langgraph/pregel/io.py | 14 ++++++-------- tests/test_pregel.py | 3 +-- tests/test_pregel_async.py | 3 +-- 4 files changed, 18 insertions(+), 38 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 5c19a09fb..166dd5274 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -557,7 +557,6 @@ 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) @@ -699,17 +698,11 @@ class Pregel( # yield current value or updates if stream_mode == "values": - if ( - step_output := map_output_values( - output_keys, pending_writes, channels - ) - ) or allow_falsy_output: - yield step_output + yield from map_output_values( + output_keys, pending_writes, channels + ) else: - if ( - step_output := map_output_updates(output_keys, next_tasks) - ) or allow_falsy_output: - yield step_output + yield from map_output_updates(output_keys, next_tasks) # save end of step checkpoint if self.checkpointer is not None and ( @@ -760,7 +753,6 @@ 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) @@ -918,17 +910,13 @@ class Pregel( # yield current value or updates if stream_mode == "values": - if ( - step_output := map_output_values( - output_keys, pending_writes, channels - ) - ) or allow_falsy_output: - yield step_output + for chunk in map_output_values( + output_keys, pending_writes, channels + ): + yield chunk else: - if ( - step_output := map_output_updates(output_keys, next_tasks) - ) or allow_falsy_output: - 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 ( @@ -979,7 +967,6 @@ 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]: @@ -997,7 +984,6 @@ 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, ): @@ -1020,7 +1006,6 @@ 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]: @@ -1038,7 +1023,6 @@ 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/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 ad5c32a09..8654a6720 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -78,8 +78,7 @@ def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None: graph.set_entry_point("return_falsy_const") graph.set_finish_point("return_falsy_const") gapp = graph.compile() - assert gapp.invoke(1, allow_falsy_output=True) == falsy_value - assert gapp.invoke(1) is None + assert gapp.invoke(1) == falsy_value def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) -> None: diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index df1c1c024..936ef0e55 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -75,8 +75,7 @@ async def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> No graph.set_entry_point("return_falsy_const") graph.set_finish_point("return_falsy_const") gapp = graph.compile() - assert (await gapp.ainvoke(1, allow_falsy_output=True)) == falsy_value - assert (await gapp.ainvoke(1)) is None + assert falsy_value == await gapp.ainvoke(1) async def test_invoke_single_process_in_out_implicit_channels(