From 96870c09351c2988ac1f6e899bfa975e708f3c2d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 30 Apr 2024 09:04:19 -0700 Subject: [PATCH 1/2] Pregel: Remove implicit creation of channels referenced by nodes - channels must be declared in the constructor args --- langgraph/pregel/__init__.py | 9 +- langgraph/pregel/validate.py | 11 +- tests/test_pregel.py | 188 +++++++++++++++++++++++++++-------- 3 files changed, 156 insertions(+), 52 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 5d5f90633..552021e4d 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -54,7 +54,6 @@ from langgraph.channels.base import ( InvalidUpdateError, create_checkpoint, ) -from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import ( BaseCheckpointSaver, Checkpoint, @@ -185,13 +184,11 @@ class Pregel( channels: Mapping[str, BaseChannel] = Field(default_factory=dict) - default_channel_cls: Type[BaseChannel] = Field(default=LastValue) - auto_validate: bool = True stream_mode: StreamMode = "values" - output_channels: Union[str, Sequence[str]] = "output" + output_channels: Union[str, Sequence[str]] """Channels to output, defaults to channel named 'output'.""" stream_channels: Optional[Union[str, Sequence[str]]] = None @@ -201,7 +198,7 @@ class Pregel( interrupt_before_nodes: Sequence[str] = Field(default_factory=list) - input_channels: Union[str, Sequence[str]] = "input" + input_channels: Union[str, Sequence[str]] step_timeout: Optional[float] = None @@ -226,7 +223,6 @@ class Pregel( values["stream_channels"], values["interrupt_after_nodes"], values["interrupt_before_nodes"], - values["default_channel_cls"], ) if values["interrupt_after_nodes"] or values["interrupt_before_nodes"]: if not values["checkpointer"]: @@ -242,7 +238,6 @@ class Pregel( self.stream_channels, self.interrupt_after_nodes, self.interrupt_before_nodes, - self.default_channel_cls, ) if self.interrupt_after_nodes or self.interrupt_before_nodes: if not self.checkpointer: diff --git a/langgraph/pregel/validate.py b/langgraph/pregel/validate.py index 69f29cfd4..0b1bfd16b 100644 --- a/langgraph/pregel/validate.py +++ b/langgraph/pregel/validate.py @@ -1,4 +1,4 @@ -from typing import Any, Mapping, Optional, Sequence, Type, Union +from typing import Mapping, Optional, Sequence, Union from langgraph.channels.base import BaseChannel from langgraph.constants import INTERRUPT @@ -13,7 +13,6 @@ def validate_graph( stream_channels: Optional[Union[str, Sequence[str]]], interrupt_after_nodes: Sequence[str], interrupt_before_nodes: Sequence[str], - default_channel_cls: Type[BaseChannel], ) -> None: subscribed_channels = set[str]() for name, node in nodes.items(): @@ -28,11 +27,11 @@ def validate_graph( for chan in subscribed_channels: if chan not in channels: - channels[chan] = default_channel_cls(Any) # type: ignore[arg-type] + raise ValueError(f"Subscribed channel '{chan}' not in 'channels'") if isinstance(input_channels, str): if input_channels not in channels: - channels[input_channels] = default_channel_cls(Any) # type: ignore[arg-type] + raise ValueError(f"Input channel '{input_channels}' not in 'channels'") if input_channels not in subscribed_channels: raise ValueError( f"Input channel {input_channels} is not subscribed to by any node" @@ -40,7 +39,7 @@ def validate_graph( else: for chan in input_channels: if chan not in channels: - channels[chan] = default_channel_cls(Any) # type: ignore[arg-type] + raise ValueError(f"Input channel '{chan}' not in 'channels'") if all(chan not in subscribed_channels for chan in input_channels): raise ValueError( f"None of the input channels {input_channels} are subscribed to by any node" @@ -58,7 +57,7 @@ def validate_graph( for chan in all_output_channels: if chan not in channels: - channels[chan] = default_channel_cls(Any) # type: ignore[arg-type] + raise ValueError(f"Output channel '{chan}' not in 'channels'") for node in interrupt_after_nodes: if node not in nodes: diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 909be8cce..2a8025f9e 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -168,17 +168,6 @@ def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None: assert gapp.invoke(1) == falsy_value -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") - - app = Pregel(nodes={"one": chain}) - - assert app.input_schema.schema() == {"title": "LangGraphInput"} - assert app.output_schema.schema() == {"title": "LangGraphOutput"} - assert app.invoke(2) == 3 - - def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = ( @@ -188,17 +177,25 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: ) app = Pregel( - nodes={"one": chain}, output_channels=["output", "fixed", "output_plus_one"] + nodes={"one": chain}, + channels={ + "input": LastValue(int), + "output": LastValue(int), + "fixed": LastValue(int), + "output_plus_one": LastValue(int), + }, + output_channels=["output", "fixed", "output_plus_one"], + input_channels="input", ) - assert app.input_schema.schema() == {"title": "LangGraphInput"} + assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"} assert app.output_schema.schema() == { "title": "LangGraphOutput", "type": "object", "properties": { - "output": {"title": "Output"}, - "fixed": {"title": "Fixed"}, - "output_plus_one": {"title": "Output Plus One"}, + "output": {"title": "Output", "type": "integer"}, + "fixed": {"title": "Fixed", "type": "integer"}, + "output_plus_one": {"title": "Output Plus One", "type": "integer"}, }, } assert app.invoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} @@ -209,17 +206,17 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( - nodes={ - "one": chain, - }, + nodes={"one": chain}, + channels={"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", output_channels=["output"], ) - assert app.input_schema.schema() == {"title": "LangGraphInput"} + assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"} assert app.output_schema.schema() == { "title": "LangGraphOutput", "type": "object", - "properties": {"output": {"title": "Output"}}, + "properties": {"output": {"title": "Output", "type": "integer"}}, } assert app.invoke(2) == {"output": 3} @@ -229,9 +226,8 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( - nodes={ - "one": chain, - }, + nodes={"one": chain}, + channels={"input": LastValue(int), "output": LastValue(int)}, input_channels=["input"], output_channels=["output"], ) @@ -239,12 +235,12 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: assert app.input_schema.schema() == { "title": "LangGraphInput", "type": "object", - "properties": {"input": {"title": "Input"}}, + "properties": {"input": {"title": "Input", "type": "integer"}}, } assert app.output_schema.schema() == { "title": "LangGraphOutput", "type": "object", - "properties": {"output": {"title": "Output"}}, + "properties": {"output": {"title": "Output", "type": "integer"}}, } assert app.invoke({"input": 2}) == {"output": 3} @@ -256,6 +252,13 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: app = Pregel( nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", ) assert app.invoke(2) == 4 @@ -302,6 +305,13 @@ def test_invoke_two_processes_in_out_interrupt( memory = MemorySaverAssertImmutable(at=checkpoint_at) app = Pregel( nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", checkpointer=memory, interrupt_after_nodes=["one"], ) @@ -357,9 +367,14 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: app = Pregel( nodes={"one": one, "two": two}, - channels={"inbox": Topic(int)}, + channels={ + "inbox": Topic(int), + "output": LastValue(int), + "input": LastValue(int), + }, input_channels=["input", "inbox"], stream_channels=["output", "inbox"], + output_channels=["output"], ) # [12 + 1, 2 + 1 + 1] @@ -394,7 +409,16 @@ def test_batch_two_processes_in_out() -> None: one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output") - app = Pregel(nodes={"one": one, "two": two}) + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "one": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] assert app.batch([3, 2, 1, 3, 5], output_keys=["output"]) == [ @@ -427,7 +451,13 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: ) nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(nodes=nodes) + app = Pregel( + nodes=nodes, + channels={str(i): LastValue(int) for i in range(-1, test_size - 2)} + | {"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", + output_channels="output", + ) for _ in range(10): assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size @@ -449,7 +479,13 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: ) nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(nodes=nodes) + app = Pregel( + nodes=nodes, + channels={str(i): LastValue(int) for i in range(-1, test_size - 2)} + | {"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", + output_channels="output", + ) for _ in range(3): assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ @@ -476,7 +512,12 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - app = Pregel(nodes={"one": one, "two": two}) + app = Pregel( + nodes={"one": one, "two": two}, + channels={"output": LastValue(int), "input": LastValue(int)}, + input_channels="input", + output_channels="output", + ) with pytest.raises(InvalidUpdateError): # LastValue channels can only be updated once per iteration @@ -491,7 +532,12 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non app = Pregel( nodes={"one": one, "two": two}, - channels={"output": Topic(int)}, + channels={ + "input": LastValue(int), + "output": Topic(int), + }, + input_channels="input", + output_channels="output", ) # An Inbox channel accumulates updates into a sequence @@ -520,7 +566,13 @@ def test_invoke_checkpoint(mocker: MockerFixture, checkpoint_at: CheckpointAt) - app = Pregel( nodes={"one": one}, - channels={"total": BinaryOperatorAggregate(int, operator.add)}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", checkpointer=memory, ) @@ -575,7 +627,13 @@ def test_invoke_checkpoint_sqlite( memory.at = checkpoint_at app = Pregel( nodes={"one": one}, - channels={"total": BinaryOperatorAggregate(int, operator.add)}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", checkpointer=memory, ) @@ -663,7 +721,13 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None "chain_three": chain_three, "chain_four": chain_four, }, - channels={"inbox": Topic(int)}, + channels={ + "inbox": Topic(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", ) # Then invoke app @@ -683,7 +747,13 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: inner_app = Pregel( nodes={ "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") - } + }, + channels={ + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", ) one = ( @@ -705,7 +775,14 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: "two": two, "chain_three": chain_three, }, - channels={"inbox_one": Topic(int)}, + channels={ + "inbox_one": Topic(int), + "outbox_one": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", ) for _ in range(10): @@ -725,7 +802,17 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: ) two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") - app = Pregel(nodes={"one": one, "two": two}, stream_channels=["output", "between"]) + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "input": LastValue(int), + "between": LastValue(int), + "output": LastValue(int), + }, + stream_channels=["output", "between"], + input_channels="input", + output_channels="output", + ) assert [c for c in app.stream(2, stream_mode="updates")] == [ {"one": {"between": 3, "output": 3}}, @@ -742,7 +829,16 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") two = Channel.subscribe_to("between") | add_one - app = Pregel(nodes={"one": one, "two": two}) + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "input": LastValue(int), + "between": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) # It finishes executing (once no more messages being published) # but returns nothing, as nothing was published to OUT topic @@ -784,7 +880,10 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: channels={ "inbox": Topic(int), "ctx": Context(an_int, typ=int), + "output": LastValue(int), + "input": LastValue(int), }, + input_channels="input", output_channels=["inbox", "output"], stream_channels=["inbox", "output"], ) @@ -3928,6 +4027,17 @@ def test_nested_graph(snapshot: SnapshotAssertion) -> None: }, ] + chain = app | RunnablePassthrough() + + assert chain.invoke({"my_key": "my value", "never_called": never_called}) == { + "my_key": "my value there and back again", + "never_called": never_called, + } + assert [*chain.stream({"my_key": "my value", "never_called": never_called})] == [ + {"inner": {"my_key": "my value there"}}, + {"side": {"my_key": "my value there and back again"}}, + ] + def test_repeat_condition(snapshot: SnapshotAssertion) -> None: class AgentState(TypedDict): From 9a79eb54f984cbab248a5114f85a6dd020d731d8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 30 Apr 2024 09:23:41 -0700 Subject: [PATCH 2/2] Finish fixing tests --- tests/test_pregel.py | 2 +- tests/test_pregel_async.py | 209 ++++++++++++++++++++++++++++++------- 2 files changed, 170 insertions(+), 41 deletions(-) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 2a8025f9e..d12ce4898 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -740,7 +740,7 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None assert [*executor.map(app.invoke, [2] * 100)] == [[13, 13]] * 100 -def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: +def test_invoke_join_then_call_other_pregel(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 1196de6d6..238dfb5e9 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -80,19 +80,6 @@ async def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> No assert falsy_value == await gapp.ainvoke(1) -async 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") - - app = Pregel(nodes={"one": chain}) - - assert app.input_schema.schema() == {"title": "LangGraphInput"} - assert app.output_schema.schema() == {"title": "LangGraphOutput"} - assert await app.ainvoke(2) == 3 - - async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = ( @@ -102,17 +89,25 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N ) app = Pregel( - nodes={"one": chain}, output_channels=["output", "fixed", "output_plus_one"] + nodes={"one": chain}, + channels={ + "input": LastValue(int), + "output": LastValue(int), + "fixed": LastValue(int), + "output_plus_one": LastValue(int), + }, + output_channels=["output", "fixed", "output_plus_one"], + input_channels="input", ) - assert app.input_schema.schema() == {"title": "LangGraphInput"} + assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"} assert app.output_schema.schema() == { "title": "LangGraphOutput", "type": "object", "properties": { - "output": {"title": "Output"}, - "fixed": {"title": "Fixed"}, - "output_plus_one": {"title": "Output Plus One"}, + "output": {"title": "Output", "type": "integer"}, + "fixed": {"title": "Fixed", "type": "integer"}, + "output_plus_one": {"title": "Output Plus One", "type": "integer"}, }, } assert await app.ainvoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} @@ -124,14 +119,16 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: app = Pregel( nodes={"one": chain}, + channels={"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", output_channels=["output"], ) - assert app.input_schema.schema() == {"title": "LangGraphInput"} + assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"} assert app.output_schema.schema() == { "title": "LangGraphOutput", "type": "object", - "properties": {"output": {"title": "Output"}}, + "properties": {"output": {"title": "Output", "type": "integer"}}, } assert await app.ainvoke(2) == {"output": 3} @@ -141,9 +138,8 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( - nodes={ - "one": chain, - }, + nodes={"one": chain}, + channels={"input": LastValue(int), "output": LastValue(int)}, input_channels=["input"], output_channels=["output"], ) @@ -151,12 +147,12 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> assert app.input_schema.schema() == { "title": "LangGraphInput", "type": "object", - "properties": {"input": {"title": "Input"}}, + "properties": {"input": {"title": "Input", "type": "integer"}}, } assert app.output_schema.schema() == { "title": "LangGraphOutput", "type": "object", - "properties": {"output": {"title": "Output"}}, + "properties": {"output": {"title": "Output", "type": "integer"}}, } assert await app.ainvoke({"input": 2}) == {"output": 3} @@ -166,7 +162,17 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") - app = Pregel(nodes={"one": one, "two": two}, stream_channels=["inbox", "output"]) + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", + stream_channels=["inbox", "output"], + ) assert await app.ainvoke(2) == 4 @@ -226,6 +232,13 @@ async def test_invoke_two_processes_in_out_interrupt( memory = MemorySaverAssertImmutable(at=checkpoint_at) app = Pregel( nodes={"one": one, "two": two}, + channels={ + "inbox": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", checkpointer=memory, interrupt_after_nodes=["one"], ) @@ -281,9 +294,14 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: app = Pregel( nodes={"one": one, "two": two}, - channels={"inbox": Topic(int)}, + channels={ + "inbox": Topic(int), + "output": LastValue(int), + "input": LastValue(int), + }, input_channels=["input", "inbox"], - stream_channels=["inbox", "output"], + stream_channels=["output", "inbox"], + output_channels=["output"], ) # [12 + 1, 2 + 1 + 1] @@ -322,7 +340,13 @@ async def test_batch_two_processes_in_out() -> None: app = Pregel( nodes={"one": one, "two": two}, - channels={"one": LastValue(int)}, + channels={ + "one": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", ) assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] @@ -356,7 +380,13 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: ) nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(nodes=nodes) + app = Pregel( + nodes=nodes, + channels={str(i): LastValue(int) for i in range(-1, test_size - 2)} + | {"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", + output_channels="output", + ) # No state is left over from previous invocations for _ in range(10): @@ -379,7 +409,13 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: ) nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") - app = Pregel(nodes=nodes) + app = Pregel( + nodes=nodes, + channels={str(i): LastValue(int) for i in range(-1, test_size - 2)} + | {"input": LastValue(int), "output": LastValue(int)}, + input_channels="input", + output_channels="output", + ) # No state is left over from previous invocations for _ in range(3): @@ -409,7 +445,12 @@ async def test_invoke_two_processes_two_in_two_out_invalid( one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - app = Pregel(nodes={"one": one, "two": two}) + app = Pregel( + nodes={"one": one, "two": two}, + channels={"output": LastValue(int), "input": LastValue(int)}, + input_channels="input", + output_channels="output", + ) with pytest.raises(InvalidUpdateError): # LastValue channels can only be updated once per iteration @@ -424,7 +465,12 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) app = Pregel( nodes={"one": one, "two": two}, - channels={"output": Topic(int)}, + channels={ + "input": LastValue(int), + "output": Topic(int), + }, + input_channels="input", + output_channels="output", ) # An Topic channel accumulates updates into a sequence @@ -455,7 +501,13 @@ async def test_invoke_checkpoint( app = Pregel( nodes={"one": one}, - channels={"total": BinaryOperatorAggregate(int, operator.add)}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", checkpointer=memory, ) @@ -510,7 +562,13 @@ async def test_invoke_checkpoint_aiosqlite( memory.at = checkpoint_at app = Pregel( nodes={"one": one}, - channels={"total": BinaryOperatorAggregate(int, operator.add)}, + channels={ + "total": BinaryOperatorAggregate(int, operator.add), + "input": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", checkpointer=memory, debug=True, ) @@ -605,7 +663,13 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) - "chain_three": chain_three, "chain_four": chain_four, }, - channels={"inbox": Topic(int)}, + channels={ + "inbox": Topic(int), + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", ) # Then invoke app @@ -619,14 +683,20 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) - ] -async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None: +async def test_invoke_join_then_call_other_pregel(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) inner_app = Pregel( nodes={ "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") - } + }, + channels={ + "output": LastValue(int), + "input": LastValue(int), + }, + input_channels="input", + output_channels="output", ) one = ( @@ -651,7 +721,11 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None channels={ "inbox_one": Topic(int), "outbox_one": LastValue(int), + "output": LastValue(int), + "input": LastValue(int), }, + input_channels="input", + output_channels="output", ) # Then invoke pubsub @@ -673,7 +747,17 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non ) two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") - app = Pregel(nodes={"one": one, "two": two}, stream_channels=["output", "between"]) + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "input": LastValue(int), + "between": LastValue(int), + "output": LastValue(int), + }, + stream_channels=["output", "between"], + input_channels="input", + output_channels="output", + ) # Then invoke pubsub assert [c async for c in app.astream(2)] == [ @@ -687,7 +771,16 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") two = Channel.subscribe_to("between") | add_one - app = Pregel(nodes={"one": one, "two": two}) + app = Pregel( + nodes={"one": one, "two": two}, + channels={ + "input": LastValue(int), + "between": LastValue(int), + "output": LastValue(int), + }, + input_channels="input", + output_channels="output", + ) # It finishes executing (once no more messages being published) # but returns nothing, as nothing was published to "output" topic @@ -727,9 +820,12 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: app = Pregel( nodes={"one": one, "two": two}, channels={ + "input": LastValue(int), + "output": LastValue(int), "inbox": Topic(int), "ctx": Context(an_int, an_int_async, typ=int), }, + input_channels="input", output_channels=["inbox", "output"], stream_channels=["inbox", "output"], ) @@ -3563,3 +3659,36 @@ async def test_nested_graph(snapshot: SnapshotAssertion) -> None: ] } assert times_called == 1 + + chain = app | RunnablePassthrough() + + assert await chain.ainvoke( + {"my_key": "my value", "never_called": never_called} + ) == { + "my_key": "my value there and back again", + "never_called": never_called, + } + assert [ + chunk + async for chunk in chain.astream( + {"my_key": "my value", "never_called": never_called} + ) + ] == [ + {"inner": {"my_key": "my value there"}}, + {"side": {"my_key": "my value there and back again"}}, + ] + times_called = 0 + async for event in chain.astream_events( + {"my_key": "my value", "never_called": never_called}, + version="v1", + config={"run_id": UUID(int=0)}, + ): + if event["event"] == "on_chain_end" and event["run_id"] == str(UUID(int=0)): + times_called += 1 + assert event["data"] == { + "output": [ + {"inner": {"my_key": "my value there"}}, + {"side": {"my_key": "my value there and back again"}}, + ] + } + assert times_called == 1