From 46e67cff87d2502c3da7383e0f7fac07df8346ed Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 11 Apr 2024 08:44:41 -0700 Subject: [PATCH] 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(