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

Move subscribe/write cls methods to Channel

Browse Source
This commit is contained in:
Nuno Campos
2023-10-23 15:42:51 +01:00
parent f363310d80
commit dbe8c1362f
12 changed files with 307 additions and 271 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+5 -4
View File
@@ -47,17 +47,18 @@ Repeat until no chains are planned for execution, or a maximum number of steps i
## Example
```python
from permchain import Channels, Pregel
from permchain import Channel, Pregel
from permchain.channels import LastValue
grow_value = (
Pregel.subscribe_to("value")
Channel.subscribe_to("value")
| (lambda x: x + x)
| Pregel.write_to(value=lambda x: x if len(x) < 10 else None)
| Channel.write_to(value=lambda x: x if len(x) < 10 else None)
)
app = Pregel(
chains={"grow_value": grow_value},
channels={"value": Channels.LastValue(str)},
channels={"value": LastValue(str)},
input="value",
output="value",
)
examples/combine_docs.ipynb
+13 -12
View File
@@ -25,7 +25,8 @@
"from langchain.schema.document import Document\n",
"from langchain.schema import format_document\n",
"\n",
"from permchain import Channels, Pregel, PregelRead"
"from permchain import Channel, Pregel, PregelRead\n",
"from permchain.channels import LastValue, Inbox"
]
},
{
@@ -208,12 +209,12 @@
"source": [
"channels = {\n",
" # input\n",
" \"question\": Channels.LastValue(str),\n",
" \"docs\": Channels.Inbox(Document),\n",
" \"question\": LastValue(str),\n",
" \"docs\": Inbox(Document),\n",
" # intermediate\n",
" \"docs_to_finalize\": Channels.Inbox(Document),\n",
" \"docs_to_finalize\": Inbox(Document),\n",
" # output\n",
" \"answer\": Channels.LastValue(str),\n",
" \"answer\": LastValue(str),\n",
"}"
]
},
@@ -227,14 +228,14 @@
"def decide(docs: list[Document]) -> Runnable:\n",
" if len(_split_list_of_docs(docs)) > 1:\n",
" # send back to the beginning if we still need to collapse more\n",
" return Pregel.write_to(\"docs\")\n",
" return Channel.write_to(\"docs\")\n",
" else:\n",
" # send to the finalizer if we're ready to produce final answer\n",
" return Pregel.write_to(\"docs_to_finalize\")\n",
" return Channel.write_to(\"docs_to_finalize\")\n",
"\n",
"\n",
"collapse = (\n",
" Pregel.subscribe_to(\"docs\")\n",
" Channel.subscribe_to(\"docs\")\n",
" | _split_list_of_docs\n",
" | {\"docs_list\": RunnablePassthrough(), \"question\": PregelRead(\"question\")}\n",
" # {docs: list[list[Doc]], question: str} -> list[{docs: list[Doc], question: str}]\n",
@@ -246,9 +247,9 @@
"\n",
"# Convert final set of docs to an answer\n",
"finalize = (\n",
" Pregel.subscribe_to(\"docs_to_finalize\", key=\"docs\").join([\"question\"])\n",
" Channel.subscribe_to(\"docs_to_finalize\", key=\"docs\").join([\"question\"])\n",
" | stuff_chain\n",
" | Pregel.write_to(\"answer\")\n",
" | Channel.write_to(\"answer\")\n",
")"
]
},
@@ -314,7 +315,7 @@
"\u001b[0m- finalize({'docs': (Document(page_content='Harrison used to work at Kensho.'),\n",
" Document(page_content='Harrison used to work at Kensho.'))})\n",
"\u001b[36;1m\u001b[1;3m[pregel/checkpoint]\u001b[0m \u001b[1mFinishing step 3. Channel values:\n",
"\u001b[0m{'answer': 'Harrison used to work at Kensho.',\n",
"\u001b[0m{'answer': 'Harrison worked at Kensho.',\n",
" 'docs': (...),\n",
" 'docs_to_finalize': (...),\n",
" 'question': 'where did harrison work'}\n"
@@ -323,7 +324,7 @@
{
"data": {
"text/plain": [
"'Harrison used to work at Kensho.'"
"'Harrison worked at Kensho.'"
]
},
"execution_count": 13,
examples/draft-revise-loop.py
+11 -10
View File
@@ -5,7 +5,8 @@ from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
from langchain.prompts import SystemMessagePromptTemplate
from langchain.schema.output_parser import StrOutputParser
from permchain import Channels, Pregel
from permchain import Channel, Pregel
from permchain.channels import LastValue
# prompts
@@ -75,23 +76,23 @@ reviser_chain = reviser_prompt | gpt3 | StrOutputParser()
# application
channels = {
"question": Channels.LastValue(str),
"draft": Channels.LastValue(str),
"notes": Channels.LastValue(str),
"question": LastValue(str),
"draft": LastValue(str),
"notes": LastValue(str),
}
drafter = (
# subscribe to question channel as a dict with a single key, "question"
Pregel.subscribe_to(["question"])
Channel.subscribe_to(["question"])
| drafter_chain
| Pregel.write_to("draft")
| Channel.write_to("draft")
)
editor = (
# subscribe to draft channel as a dict with a single key, "draft"
Pregel.subscribe_to(["draft"])
Channel.subscribe_to(["draft"])
| editor_chain
| Pregel.write_to(
| Channel.write_to(
# send to "notes" channel if the editor does not accept the draft
notes=lambda x: x["arguments"]["notes"]
if x["name"] == "revise"
@@ -102,9 +103,9 @@ editor = (
reviser = (
# subscribe to new values of "notes" channel,
# and join them with the input value (question) and "draft"
Pregel.subscribe_to(["notes"]).join(["question", "draft"])
Channel.subscribe_to(["notes"]).join(["question", "draft"])
| reviser_chain
| Pregel.write_to("draft")
| Channel.write_to("draft")
)
draft_revise_loop = Pregel(
examples/readme.py
+5 -4
View File
@@ -1,14 +1,15 @@
from permchain import Channels, Pregel
from permchain import Channel, Pregel
from permchain.channels import LastValue
grow_value = (
Pregel.subscribe_to("value")
Channel.subscribe_to("value")
| (lambda x: x + x)
| Pregel.write_to(value=lambda x: x if len(x) < 10 else None)
| Channel.write_to(value=lambda x: x if len(x) < 10 else None)
)
app = Pregel(
chains={"grow_value": grow_value},
channels={"value": Channels.LastValue(str)},
channels={"value": LastValue(str)},
input="value",
output="value",
)
examples/recursive-web-loader.py
+10 -9
View File
@@ -6,7 +6,8 @@ from langchain.schema import Document
from langchain.schema.runnable import RunnableLambda, RunnablePassthrough
from langchain.utils.html import extract_sub_links
from permchain import Channels, Pregel
from permchain import Channel, Pregel
from permchain.channels import Archive, Context, LastValue, UniqueArchive, UniqueInbox
# Load url with sync httpx client
@@ -84,23 +85,23 @@ def recursive_web_loader(
metadata_extractor = metadata_extractor or _metadata_extractor
# define the channels
channels = {
"base_url": Channels.LastValue(str),
"next_urls": Channels.UniqueInbox(str),
"documents": Channels.Archive(Document),
"visited": Channels.UniqueArchive(str),
"client": Channels.Context(httpx_client, httpx_aclient),
"base_url": LastValue(str),
"next_urls": UniqueInbox(str),
"documents": Archive(Document),
"visited": UniqueArchive(str),
"client": Context(httpx_client, httpx_aclient),
}
# the main chain that gets executed recursively
visitor = (
# while there are urls in next_urls
# run the chain below for each url in next_urls
# adding the current values of visited set, base_url and httpx client
Pregel.subscribe_to_each("next_urls", key="url").join(
Channel.subscribe_to_each("next_urls", key="url").join(
["visited", "client", "base_url"]
)
# load the url (with sync and async implementations)
| RunnablePassthrough.assign(body=RunnableLambda(load_url, load_url_async))
| Pregel.write_to(
| Channel.write_to(
# send this url to the visited set
visited=lambda x: x["url"],
# send a new document to the documents stream
@@ -123,7 +124,7 @@ def recursive_web_loader(
channels=channels,
chains={
# use the base_url as the first url to visit
"input": Pregel.subscribe_to("base_url") | Pregel.write_to("next_urls"),
"input": Channel.subscribe_to("base_url") | Channel.write_to("next_urls"),
# add the main chain
"visitor": visitor,
},
permchain/__init__.py
+2 -3
View File
@@ -1,5 +1,4 @@
import permchain.channels as Channels
from permchain.pregel import Pregel
from permchain.pregel import Channel, Pregel
from permchain.pregel.read import PregelRead
__all__ = ["Channels", "Pregel", "PregelRead"]
__all__ = ["Channel", "Pregel", "PregelRead"]
permchain/pregel/__init__.py
+50 -48
View File
@@ -48,6 +48,56 @@ from permchain.pregel.validate import validate_chains_channels
from permchain.pregel.write import PregelSink
class Channel:
@overload
@classmethod
def subscribe_to(cls, channels: str, key: Optional[str] = None) -> PregelInvoke:
...
@overload
@classmethod
def subscribe_to(cls, channels: Sequence[str], key: None = None) -> PregelInvoke:
...
@classmethod
def subscribe_to(
cls, channels: str | Sequence[str], key: Optional[str] = None
) -> PregelInvoke:
"""Runs process.invoke() each time channels are updated,
with a dict of the channel values as input."""
if not isinstance(channels, str) and key is not None:
raise ValueError(
"Can't specify a key when subscribing to multiple channels"
)
return PregelInvoke(
channels=cast(
Mapping[None, str] | Mapping[str, str],
{key: channels}
if isinstance(channels, str)
else {chan: chan for chan in channels},
)
)
@classmethod
def subscribe_to_each(cls, inbox: str, key: Optional[str] = None) -> PregelBatch:
"""Runs process.batch() with the content of inbox each time it is updated."""
return PregelBatch(channel=inbox, key=key)
@classmethod
def write_to(
cls,
*channels: str,
**kwargs: RunnableLike,
) -> PregelSink:
"""Writes to channels the result of the lambda, or None to skip writing."""
return PregelSink(
channels=(
[(c, RunnablePassthrough()) for c in channels]
+ [(k, coerce_to_runnable(v)) for k, v in kwargs.items()]
)
)
class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
channels: Mapping[str, BaseChannel]
@@ -106,54 +156,6 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
**{k: (self.channels[k].ValueType, None) for k in self.output},
)
@overload
@classmethod
def subscribe_to(cls, channels: str, key: Optional[str] = None) -> PregelInvoke:
...
@overload
@classmethod
def subscribe_to(cls, channels: Sequence[str], key: None = None) -> PregelInvoke:
...
@classmethod
def subscribe_to(
cls, channels: str | Sequence[str], key: Optional[str] = None
) -> PregelInvoke:
"""Runs process.invoke() each time channels are updated,
with a dict of the channel values as input."""
if not isinstance(channels, str) and key is not None:
raise ValueError(
"Can't specify a key when subscribing to multiple channels"
)
return PregelInvoke(
channels=cast(
Mapping[None, str] | Mapping[str, str],
{key: channels}
if isinstance(channels, str)
else {chan: chan for chan in channels},
)
)
@classmethod
def subscribe_to_each(cls, inbox: str, key: Optional[str] = None) -> PregelBatch:
"""Runs process.batch() with the content of inbox each time it is updated."""
return PregelBatch(channel=inbox, key=key)
@classmethod
def write_to(
cls,
*channels: str,
**kwargs: RunnableLike,
) -> PregelSink:
"""Writes to channels the result of the lambda, or None to skip writing."""
return PregelSink(
channels=(
[(c, RunnablePassthrough()) for c in channels]
+ [(k, coerce_to_runnable(v)) for k, v in kwargs.items()]
)
)
def _transform(
self,
input: Iterator[dict[str, Any] | Any],
permchain/pregel/read.py
+1 -1
View File
@@ -123,7 +123,7 @@ class PregelBatch(RunnableEach):
if self.key is None:
raise ValueError(
"Cannot join() additional channels without a key."
" Pass a key arg to Pregel.subscribe_to_each()."
" Pass a key arg to Channel.subscribe_to_each()."
)
joiner = RunnablePassthrough.assign(
permchain/pregel/validate.py
+1 -1
View File
@@ -18,7 +18,7 @@ def validate_chains_channels(
subscribed_channels.add(chain.channel)
else:
raise TypeError(
f"Invalid chain type {type(chain)}, expected Pregel.subscribe_to() or Pregel.subscribe_to_each()"
f"Invalid chain type {type(chain)}, expected Channel.subscribe_to() or Channel.subscribe_to_each()"
)
for chan in subscribed_channels:
tests/test_channels.py
+16 -12
View File
@@ -6,12 +6,16 @@ import httpx
import pytest
from pytest_mock import MockerFixture
import permchain.channels as Channels
from permchain.channels.archive import UniqueArchive
from permchain.channels.base import EmptyChannelError, InvalidUpdateError
from permchain.channels.binop import BinaryOperatorAggregate
from permchain.channels.context import Context
from permchain.channels.inbox import Inbox
from permchain.channels.last_value import LastValue
def test_last_value() -> None:
with Channels.LastValue(int).empty() as channel:
with LastValue(int).empty() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -27,7 +31,7 @@ def test_last_value() -> None:
async def test_last_value_async() -> None:
async with Channels.LastValue(int).aempty() as channel:
async with LastValue(int).aempty() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -43,7 +47,7 @@ async def test_last_value_async() -> None:
def test_inbox() -> None:
with Channels.Inbox(str).empty() as channel:
with Inbox(str).empty() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, Sequence[str]]
@@ -57,7 +61,7 @@ def test_inbox() -> None:
async def test_inbox_async() -> None:
async with Channels.Inbox(str).aempty() as channel:
async with Inbox(str).aempty() as channel:
assert channel.ValueType is Sequence[str]
assert channel.UpdateType is Union[str, Sequence[str]]
@@ -72,7 +76,7 @@ async def test_inbox_async() -> None:
def test_set() -> None:
with Channels.UniqueArchive(str).empty() as channel:
with UniqueArchive(str).empty() as channel:
assert channel.ValueType is FrozenSet[str]
assert channel.UpdateType is str
@@ -84,7 +88,7 @@ def test_set() -> None:
async def test_set_async() -> None:
async with Channels.UniqueArchive(str).aempty() as channel:
async with UniqueArchive(str).aempty() as channel:
assert channel.ValueType is FrozenSet[str]
assert channel.UpdateType is str
@@ -96,7 +100,7 @@ async def test_set_async() -> None:
def test_binop() -> None:
with Channels.BinaryOperatorAggregate(int, operator.add).empty() as channel:
with BinaryOperatorAggregate(int, operator.add).empty() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -110,7 +114,7 @@ def test_binop() -> None:
async def test_binop_async() -> None:
async with Channels.BinaryOperatorAggregate(int, operator.add).aempty() as channel:
async with BinaryOperatorAggregate(int, operator.add).aempty() as channel:
assert channel.ValueType is int
assert channel.UpdateType is int
@@ -135,7 +139,7 @@ def test_ctx_manager(mocker: MockerFixture) -> None:
finally:
cleanup()
with Channels.Context(an_int, None, int).empty() as channel:
with Context(an_int, None, int).empty() as channel:
assert setup.call_count == 1
assert cleanup.call_count == 0
@@ -153,7 +157,7 @@ def test_ctx_manager(mocker: MockerFixture) -> None:
def test_ctx_manager_ctx(mocker: MockerFixture) -> None:
with Channels.Context(httpx.Client).empty() as channel:
with Context(httpx.Client).empty() as channel:
assert channel.ValueType is httpx.Client
with pytest.raises(InvalidUpdateError):
assert channel.UpdateType is None
@@ -183,7 +187,7 @@ async def test_ctx_manager_async(mocker: MockerFixture) -> None:
finally:
cleanup()
async with Channels.Context(an_int_sync, an_int, int).aempty() as channel:
async with Context(an_int_sync, an_int, int).aempty() as channel:
assert setup.call_count == 1
assert cleanup.call_count == 0
tests/test_pregel.py
+99 -86
View File
@@ -7,21 +7,24 @@ import pytest
from langchain.schema.runnable import RunnablePassthrough
from pytest_mock import MockerFixture
from permchain import Channels, Pregel
from permchain import Channel, Pregel
from permchain.channels.base import InvalidUpdateError
from permchain.channels.context import Context
from permchain.channels.inbox import Inbox
from permchain.channels.last_value import LastValue
def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
"one": chain,
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
},
input="input",
output="output",
@@ -34,15 +37,15 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
"one": chain,
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
},
input="input",
output=["output"],
@@ -59,15 +62,15 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
"one": chain,
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
},
input=["input"],
output=["output"],
@@ -88,15 +91,17 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = (
Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
)
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
"input": LastValue(int),
"output": LastValue(int),
"inbox": Inbox(int),
},
input="input",
output="output",
@@ -107,15 +112,17 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = (
Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
)
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
"input": LastValue(int),
"output": LastValue(int),
"inbox": Inbox(int),
},
input=["input", "inbox"],
output="output",
@@ -130,18 +137,18 @@ def test_batch_two_processes_in_out() -> None:
return inp + 1
chain_one = (
Pregel.subscribe_to("input") | add_one_with_delay | Pregel.write_to("one")
Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one")
)
chain_two = (
Pregel.subscribe_to("one") | add_one_with_delay | Pregel.write_to("output")
Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output")
)
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"one": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"one": LastValue(int),
},
input="input",
output="output",
@@ -155,17 +162,17 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chans = {
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"-1": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"-1": LastValue(int),
}
chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
for i in range(test_size - 2):
chans[str(i)] = Channels.LastValue(int)
chans[str(i)] = LastValue(int)
chains[str(i)] = (
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
)
chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output")
chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(chains=chains, channels=chans, input="input", output="output")
@@ -183,17 +190,17 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chans = {
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"-1": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"-1": LastValue(int),
}
chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
for i in range(test_size - 2):
chans[str(i)] = Channels.LastValue(int)
chans[str(i)] = LastValue(int)
chains[str(i)] = (
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
)
chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output")
chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(chains=chains, channels=chans, input="input", output="output")
@@ -219,14 +226,14 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
},
input="input",
output="output",
@@ -240,14 +247,14 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N
def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.Inbox(int),
"input": LastValue(int),
"output": Inbox(int),
},
input="input",
output="output",
@@ -261,9 +268,11 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None
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))
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
chain_three = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_four = (
Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output")
)
app = Pregel(
chains={
@@ -272,9 +281,9 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None
"chain_four": chain_four,
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
"input": LastValue(int),
"output": LastValue(int),
"inbox": Inbox(int),
},
input="input",
output="output",
@@ -296,26 +305,28 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
inner_app = Pregel(
chains={
"one": Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
"one": Channel.subscribe_to("input") | add_one | Channel.write_to("output")
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
},
input="input",
output="output",
)
chain_one = (
Pregel.subscribe_to("input") | add_10_each | Pregel.write_to("inbox_one").map()
Channel.subscribe_to("input")
| add_10_each
| Channel.write_to("inbox_one").map()
)
chain_two = (
Pregel.subscribe_to("inbox_one")
Channel.subscribe_to("inbox_one")
| inner_app.map()
| sorted
| Pregel.write_to("outbox_one")
| Channel.write_to("outbox_one")
)
chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.write_to("output")
chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output")
app = Pregel(
chains={
@@ -324,10 +335,10 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None:
"chain_three": chain_three,
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox_one": Channels.Inbox(int),
"outbox_one": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"inbox_one": Inbox(int),
"outbox_one": LastValue(int),
},
input="input",
output="output",
@@ -344,18 +355,18 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = (
Pregel.subscribe_to("input")
Channel.subscribe_to("input")
| add_one
| Pregel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
| Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
)
chain_two = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output")
chain_two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"between": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"between": LastValue(int),
},
input="input",
output="output",
@@ -366,15 +377,15 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("between")
chain_two = Pregel.subscribe_to("between") | add_one
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("between")
chain_two = Channel.subscribe_to("between") | add_one
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"between": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"between": LastValue(int),
},
input="input",
output="output",
@@ -388,16 +399,16 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output")
chain_two = Pregel.subscribe_to("between") | add_one
chain_one = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
chain_two = Channel.subscribe_to("between") | add_one
with pytest.raises(ValueError):
Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"between": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"between": LastValue(int),
},
input="input",
output="output",
@@ -417,16 +428,18 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
cleanup()
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = (
Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
)
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
"ctx": Channels.Context(an_int, typ=int),
"input": LastValue(int),
"output": LastValue(int),
"inbox": Inbox(int),
"ctx": Context(an_int, typ=int),
},
input="input",
output=["inbox", "output"],
tests/test_pregel_async.py
+94 -81
View File
@@ -6,21 +6,24 @@ import pytest
from langchain.schema.runnable import RunnablePassthrough
from pytest_mock import MockerFixture
from permchain import Channels, Pregel
from permchain import Channel, Pregel
from permchain.channels.base import InvalidUpdateError
from permchain.channels.context import Context
from permchain.channels.inbox import Inbox
from permchain.channels.last_value import LastValue
async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
"one": chain,
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
},
input="input",
output="output",
@@ -31,15 +34,15 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
"one": chain,
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
},
input="input",
output=["output"],
@@ -56,15 +59,15 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={
"one": chain,
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
},
input=["input"],
output=["output"],
@@ -85,15 +88,17 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = (
Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
)
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
"input": LastValue(int),
"output": LastValue(int),
"inbox": Inbox(int),
},
input="input",
output="output",
@@ -104,15 +109,17 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = (
Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
)
pubsub = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
"input": LastValue(int),
"output": LastValue(int),
"inbox": Inbox(int),
},
input=["input", "inbox"],
output="output",
@@ -128,18 +135,18 @@ async def test_batch_two_processes_in_out() -> None:
return inp + 1
chain_one = (
Pregel.subscribe_to("input") | add_one_with_delay | Pregel.write_to("one")
Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one")
)
chain_two = (
Pregel.subscribe_to("one") | add_one_with_delay | Pregel.write_to("output")
Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output")
)
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"one": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"one": LastValue(int),
},
input="input",
output="output",
@@ -153,17 +160,17 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chans = {
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"-1": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"-1": LastValue(int),
}
chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
for i in range(test_size - 2):
chans[str(i)] = Channels.LastValue(int)
chans[str(i)] = LastValue(int)
chains[str(i)] = (
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
)
chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output")
chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(chains=chains, channels=chans, input="input", output="output")
@@ -182,17 +189,17 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chans = {
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"-1": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"-1": LastValue(int),
}
chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")}
chains = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")}
for i in range(test_size - 2):
chans[str(i)] = Channels.LastValue(int)
chans[str(i)] = LastValue(int)
chains[str(i)] = (
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i))
Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i))
)
chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output")
chains["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output")
app = Pregel(chains=chains, channels=chans, input="input", output="output")
@@ -224,14 +231,14 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
},
input="input",
output="output",
@@ -245,14 +252,14 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
chain_two = Channel.subscribe_to("input") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.Inbox(int),
"input": LastValue(int),
"output": Inbox(int),
},
input="input",
output="output",
@@ -266,9 +273,11 @@ async 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))
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
chain_three = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_four = (
Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output")
)
app = Pregel(
chains={
@@ -277,9 +286,9 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -
"chain_four": chain_four,
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
"input": LastValue(int),
"output": LastValue(int),
"inbox": Inbox(int),
},
input="input",
output="output",
@@ -302,26 +311,28 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
inner_app = Pregel(
chains={
"one": Pregel.subscribe_to("input") | add_one | Pregel.write_to("output")
"one": Channel.subscribe_to("input") | add_one | Channel.write_to("output")
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
},
input="input",
output="output",
)
chain_one = (
Pregel.subscribe_to("input") | add_10_each | Pregel.write_to("inbox_one").map()
Channel.subscribe_to("input")
| add_10_each
| Channel.write_to("inbox_one").map()
)
chain_two = (
Pregel.subscribe_to("inbox_one")
Channel.subscribe_to("inbox_one")
| inner_app.map()
| sorted
| Pregel.write_to("outbox_one")
| Channel.write_to("outbox_one")
)
chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.write_to("output")
chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output")
app = Pregel(
chains={
@@ -330,10 +341,10 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None
"chain_three": chain_three,
},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox_one": Channels.Inbox(int),
"outbox_one": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"inbox_one": Inbox(int),
"outbox_one": LastValue(int),
},
input="input",
output="output",
@@ -352,18 +363,18 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = (
Pregel.subscribe_to("input")
Channel.subscribe_to("input")
| add_one
| Pregel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
| Channel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough())
)
chain_two = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output")
chain_two = Channel.subscribe_to("between") | add_one | Channel.write_to("output")
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"between": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"between": LastValue(int),
},
input="input",
output="output",
@@ -375,15 +386,15 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("between")
chain_two = Pregel.subscribe_to("between") | add_one
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("between")
chain_two = Channel.subscribe_to("between") | add_one
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"between": Channels.LastValue(int),
"input": LastValue(int),
"output": LastValue(int),
"between": LastValue(int),
},
input="input",
output="output",
@@ -418,16 +429,18 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None:
cleanup_async()
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox")
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output")
chain_one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox")
chain_two = (
Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output")
)
app = Pregel(
chains={"chain_one": chain_one, "chain_two": chain_two},
channels={
"input": Channels.LastValue(int),
"output": Channels.LastValue(int),
"inbox": Channels.Inbox(int),
"ctx": Channels.Context(an_int, an_int_async, typ=int),
"input": LastValue(int),
"output": LastValue(int),
"inbox": Inbox(int),
"ctx": Context(an_int, an_int_async, typ=int),
},
input="input",
output=["inbox", "output"],