WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-31 04:09:49 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #366 from langchain-ai/nc/30apr/pregel-no-implicit-channels

Browse Source 
Pregel: Remove implicit creation of channels referenced by nodes
This commit is contained in:
Nuno Campos
2024-04-30 10:06:14 -07:00
committed by GitHub
parent f6d332beb9 9a79eb54f9
commit c293d3d71d
4 changed files with 326 additions and 93 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/pregel/__init__.py
+2 -7
View File
@@ -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:
langgraph/pregel/validate.py
+5 -6
View File
@@ -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:
tests/test_pregel.py
+150 -40
View File
@@ -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
@@ -676,14 +740,20 @@ 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])
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):
tests/test_pregel_async.py
+169 -40
View File
@@ -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