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

Add validation, finish porting tests to new api

Browse Source
This commit is contained in:
Nuno Campos
2023-10-16 13:52:35 +01:00
parent 07c9199d18
commit 837c9c9af2
2 changed files with 207 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files
permchain/pregel.py
+32
View File
@@ -250,6 +250,9 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
chains_flat.extend(chain)
else:
chains_flat.append(chain)
validate_chains_channels(chains_flat, channels, input, output)
super().__init__(
chains=chains_flat,
channels=channels,
@@ -584,3 +587,32 @@ def _apply_writes_and_prepare_next_tasks(
tasks.append((proc, val))
return tasks
def validate_chains_channels(
chains: Sequence[PregelInvoke | PregelBatch],
channels: Mapping[str, Channel],
input: str | None,
output: str | Sequence[str],
) -> None:
subscribed_channels = set()
for chain in chains:
if isinstance(chain, PregelInvoke):
subscribed_channels.update(chain.channels.values())
elif isinstance(chain, PregelBatch):
subscribed_channels.add(chain.channel)
else:
raise TypeError(
f"Invalid chain type {type(chain)}, expected Pregel.subscribe_to() or Pregel.subscribe_to_each()"
)
if input is not None and input not in subscribed_channels:
raise ValueError(f"Input channel {input} is not subscribed to by any chain")
for chan in subscribed_channels:
if chan not in channels:
raise ValueError(f"Channel {chan} is subscribed to, but not initialized")
if isinstance(output, str):
if output not in channels:
raise ValueError(f"Output channel {output} is not initialized")
tests/test_pregel.py
+175 -112
View File
@@ -9,48 +9,66 @@ from permchain.pregel import PregelInvoke
def test_invoke_single_process_in_out(mocker: MockerFixture):
input_chan = channels.LastValue[int]("input")
output_chan = channels.LastValue[int]("output")
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output_chan)
chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
pubsub = Pregel(chain, input=input_chan, output=output_chan)
app = Pregel(
(chain,),
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
},
input="input",
output="output",
)
# Then invoke pubsub
assert pubsub.invoke(2) == 3
assert app.invoke(2) == 3
def test_invoke_two_processes_in_out(mocker: MockerFixture):
input_chan = channels.LastValue[int]("input")
output = channels.LastValue[int]("output")
inbox = channels.Inbox[int]("inbox")
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(inbox)
chain_two = Pregel.subscribe_to_each(inbox) | add_one | Pregel.send_to(output)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
pubsub = Pregel(
[chain_one, chain_two],
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
"inbox": channels.Inbox[int](),
},
input="input",
output="output",
)
# Then invoke pubsub
assert pubsub.invoke(2) == 4
def test_batch_two_processes_in_out(mocker: MockerFixture):
def test_batch_two_processes_in_out():
def add_one_with_delay(inp: int) -> int:
time.sleep(inp / 10)
return inp + 1
input_chan = channels.LastValue[int]("input")
output = channels.LastValue[int]("output")
one = channels.LastValue[int]("one")
chain_one = (
Pregel.subscribe_to(input_chan) | add_one_with_delay | Pregel.send_to(one)
Pregel.subscribe_to("input") | add_one_with_delay | Pregel.send_to("one")
)
chain_two = (
Pregel.subscribe_to("one") | add_one_with_delay | Pregel.send_to("output")
)
chain_two = Pregel.subscribe_to(one) | add_one_with_delay | Pregel.send_to(output)
pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
pubsub = Pregel(
chain_one,
chain_two,
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
"one": channels.LastValue[int](),
},
input="input",
output="output",
)
# Then invoke pubsub
assert pubsub.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
@@ -58,52 +76,53 @@ def test_batch_two_processes_in_out(mocker: MockerFixture):
def test_invoke_many_processes_in_out(mocker: MockerFixture):
test_size = 100
input_chan = channels.LastValue[int]("input")
output = channels.LastValue[int]("output")
topics: list[channels.Channel] = [channels.LastValue[int]("zero")]
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chans = {
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
"-1": channels.LastValue[int](),
}
chains: list[PregelInvoke] = [
Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(topics[0])
Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1")
]
for i in range(test_size - 2):
topics.append(channels.LastValue[int](str(i)))
chans[str(i)] = channels.LastValue[int]()
chains.append(
Pregel.subscribe_to(topics[-2]) | add_one | Pregel.send_to(topics[-1])
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i))
)
chains.append(Pregel.subscribe_to(topics[-1]) | add_one | Pregel.send_to(output))
chains.append(Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output"))
pubsub = Pregel(*chains, input=input_chan, output=output)
app = Pregel(*chains, channels=chans, input="input", output="output")
for _ in range(10):
# Then invoke pubsub
assert pubsub.invoke(2, {"recursion_limit": test_size}) == 2 + test_size
assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size
def test_batch_many_processes_in_out(mocker: MockerFixture):
test_size = 100
input_chan = channels.LastValue[int]("input")
output = channels.LastValue[int]("output")
topics: list[channels.Channel] = [channels.LastValue[int]("zero")]
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chans = {
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
"-1": channels.LastValue[int](),
}
chains: list[PregelInvoke] = [
Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(topics[0])
Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1")
]
for i in range(test_size - 2):
topics.append(channels.LastValue[int](str(i)))
chans[str(i)] = channels.LastValue[int]()
chains.append(
Pregel.subscribe_to(topics[-2]) | add_one | Pregel.send_to(topics[-1])
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i))
)
chains.append(Pregel.subscribe_to(topics[-1]) | add_one | Pregel.send_to(output))
chains.append(Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output"))
pubsub = Pregel(*chains, input=input_chan, output=output)
app = Pregel(*chains, channels=chans, input="input", output="output")
for _ in range(10):
# Then invoke pubsub
assert pubsub.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
2 + test_size,
1 + test_size,
3 + test_size,
@@ -115,138 +134,182 @@ def test_batch_many_processes_in_out(mocker: MockerFixture):
def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture):
add_one = mocker.Mock(side_effect=lambda x: x + 1)
input_chan = channels.LastValue[int]("input")
output = channels.LastValue[int]("output")
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output)
chain_two = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output)
pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
app = Pregel(
chain_one,
chain_two,
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
},
input="input",
output="output",
)
with pytest.raises(channels.InvalidUpdateError):
# LastValue channels can only be updated once per iteration
pubsub.invoke(2)
app.invoke(2)
def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture):
add_one = mocker.Mock(side_effect=lambda x: x + 1)
input_chan = channels.LastValue[int]("input")
output = channels.Inbox[int]("output")
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output")
chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output)
chain_two = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output)
pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
app = Pregel(
chain_one,
chain_two,
channels={
"input": channels.LastValue[int](),
"output": channels.Inbox[int](),
},
input="input",
output="output",
)
# An Inbox channel accumulates updates into a sequence
assert pubsub.invoke(2) == (3, 3)
assert app.invoke(2) == (3, 3)
def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture):
add_one = mocker.Mock(side_effect=lambda x: x + 1)
add_10_each = mocker.Mock(side_effect=lambda x: sorted(y + 10 for y in x))
input_chan = channels.LastValue[int]("input")
output = channels.LastValue[int]("output")
inbox = channels.Inbox[int]("inbox")
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
chain_three = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox")
chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.send_to("output")
chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(inbox)
chain_three = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(inbox)
chain_four = Pregel.subscribe_to(inbox) | add_10_each | Pregel.send_to(output)
app = Pregel(
chain_one,
chain_three,
chain_four,
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
"inbox": channels.Inbox[int](),
},
input="input",
output="output",
)
pubsub = Pregel(chain_one, chain_three, chain_four, input=input_chan, output=output)
# Then invoke pubsub
# Then invoke app
# We get a single array result as chain_four waits for all publishers to finish
# before operating on all elements published to topic_two as an array
for _ in range(100):
assert pubsub.invoke(2) == [13, 13]
assert app.invoke(2) == [13, 13]
def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture):
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])
input_chan = channels.LastValue[int]("input")
output = channels.LastValue[int]("output")
inner_pubsub = Pregel(
Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(output),
input=input_chan,
output=output,
inner_app = Pregel(
Pregel.subscribe_to("input") | add_one | Pregel.send_to("output"),
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
},
input="input",
output="output",
)
inbox_one = channels.Inbox[int]("inbox_one")
outbox_one = channels.LastValue[int]("outbox_one")
chain_one = (
Pregel.subscribe_to(input_chan) | add_10_each | Pregel.send_to(inbox_one).map()
Pregel.subscribe_to("input") | add_10_each | Pregel.send_to("inbox_one").map()
)
chain_two = (
Pregel.subscribe_to(inbox_one)
| inner_pubsub.map()
Pregel.subscribe_to("inbox_one")
| inner_app.map()
| sorted
| Pregel.send_to(outbox_one)
| Pregel.send_to("outbox_one")
)
chain_three = Pregel.subscribe_to(outbox_one) | sum | Pregel.send_to(output)
chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.send_to("output")
pubsub = Pregel(chain_one, chain_two, chain_three, input=input_chan, output=output)
app = Pregel(
chain_one,
chain_two,
chain_three,
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
"inbox_one": channels.Inbox[int](),
"outbox_one": channels.LastValue[int](),
},
input="input",
output="output",
)
# Then invoke pubsub
for _ in range(10):
assert pubsub.invoke([2, 3]) == 27
assert app.invoke([2, 3]) == 27
def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture):
add_one = mocker.Mock(side_effect=lambda x: x + 1)
input_chan = channels.LastValue[int]("input")
output = channels.LastValue[int]("output")
between = channels.LastValue[int]("between")
chain_one = (
Pregel.subscribe_to(input_chan)
Pregel.subscribe_to("input")
| add_one
| Pregel.send_to(
{output: RunnablePassthrough(), between: RunnablePassthrough()}
)
| Pregel.send_to(output=RunnablePassthrough(), between=RunnablePassthrough())
)
chain_two = Pregel.subscribe_to(between) | add_one | Pregel.send_to(output)
chain_two = Pregel.subscribe_to("between") | add_one | Pregel.send_to("output")
pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
app = Pregel(
chain_one,
chain_two,
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
"between": channels.LastValue[int](),
},
input="input",
output="output",
)
# Then invoke pubsub
assert [c for c in pubsub.stream(2)] == [3, 4]
assert [c for c in app.stream(2)] == [3, 4]
def test_invoke_two_processes_no_out(mocker: MockerFixture):
input_chan = channels.LastValue[int]("input")
output = channels.LastValue[int]("output")
between = channels.LastValue[int]("between")
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to(input_chan) | add_one | Pregel.send_to(between)
chain_two = Pregel.subscribe_to(between) | add_one
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.send_to("between")
chain_two = Pregel.subscribe_to("between") | add_one
pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
app = Pregel(
chain_one,
chain_two,
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
"between": channels.LastValue[int](),
},
input="input",
output="output",
)
# Then invoke pubsub
# It finishes executing (once no more messages being published)
# but returns nothing, as nothing was published to OUT topic
assert pubsub.invoke(2) is None
assert app.invoke(2) is None
def test_invoke_two_processes_no_in(mocker: MockerFixture):
input_chan = channels.LastValue[int]("input")
output = channels.LastValue[int]("output")
between = channels.LastValue[int]("between")
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to(between) | add_one | Pregel.send_to(output)
chain_two = Pregel.subscribe_to(between) | add_one
pubsub = Pregel(chain_one, chain_two, input=input_chan, output=output)
chain_one = Pregel.subscribe_to("between") | add_one | Pregel.send_to("output")
chain_two = Pregel.subscribe_to("between") | add_one
with pytest.raises(ValueError):
assert pubsub.invoke(2) is None
app = Pregel(
chain_one,
chain_two,
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
"between": channels.LastValue[int](),
},
input="input",
output="output",
)