Remove flag

This commit is contained in:
Nuno Campos
2024-04-11 08:44:41 -07:00
parent 236ecb1982
commit 46e67cff87
4 changed files with 18 additions and 38 deletions
+10 -26
View File
@@ -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,
):
+6 -8
View File
@@ -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
+1 -2
View File
@@ -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:
+1 -2
View File
@@ -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(