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

Make channels context managers, Add async tests, Add concurrency tests

Browse Source
This commit is contained in:
Nuno Campos
2023-10-18 12:23:30 +01:00
parent b146b0dd7e
commit 285577fc77
8 changed files with 1681 additions and 1101 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+2
View File
@@ -35,6 +35,8 @@ Check `examples` for more examples.
- [x] do we want api to send input to multiple channels in invoke()
- [x] Finish updating tests to new API
- [ ] Implement input_schema and output_schema in Pregel
- [ ] More tests
- [ ] Test different input and output types (str, str sequence, None)
- [ ] Implement checkpointing
- [ ] Save checkpoints at end of each step
- [ ] Load checkpoint at start of invocation
examples/runnable-pregel.py
+2 -2
View File
@@ -1,7 +1,5 @@
from __future__ import annotations
import asyncio
from langchain.chat_models.openai import ChatOpenAI
from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser
from langchain.prompts import SystemMessagePromptTemplate
@@ -131,4 +129,6 @@ async def main():
print("---")
# import asyncio
# asyncio.run(main())
permchain/channels.py
+79 -7
View File
@@ -1,8 +1,16 @@
import json
from abc import ABC, abstractmethod
from typing import Callable, FrozenSet, Generic, Optional, Sequence, TypeVar
from types import TracebackType
from typing import (
Callable,
FrozenSet,
Generic,
Optional,
Sequence,
TypeVar,
)
from typing_extensions import Self, get_args
from typing_extensions import Self # , get_args
Value = TypeVar("Value")
Update = TypeVar("Update")
@@ -33,9 +41,29 @@ class Channel(Generic[Value, Update], ABC):
# return type_args[1]
@abstractmethod
def _empty(self, checkpoint: Optional[str] = None) -> Self:
def __enter__(self, checkpoint: Optional[str] = None) -> Self:
...
@abstractmethod
def __exit__(
self,
__exc_type: type[BaseException] | None,
__exc_value: BaseException | None,
__traceback: TracebackType | None,
) -> bool | None:
...
async def __aenter__(self, checkpoint: Optional[str] = None) -> Self:
return self.__enter__(checkpoint)
async def __aexit__(
self,
__exc_type: type[BaseException] | None,
__exc_value: BaseException | None,
__traceback: TracebackType | None,
) -> None:
self.__exit__(__exc_type, __exc_value, __traceback)
@abstractmethod
def _update(self, values: Sequence[Update]) -> None:
...
@@ -63,12 +91,23 @@ class BinaryOperatorAggregate(Generic[Value], Channel[Value, Value]):
super().__init__()
self.operator = operator
def _empty(self, checkpoint: Optional[str] = None) -> Self:
def __enter__(self, checkpoint: Optional[str] = None) -> Self:
empty = self.__class__(self.operator)
if checkpoint is not None:
empty.value = json.loads(checkpoint)
return empty
def __exit__(
self,
__exc_type: type[BaseException] | None,
__exc_value: BaseException | None,
__traceback: TracebackType | None,
) -> None:
try:
del self.value
except AttributeError:
pass
def _update(self, values: Sequence[Value]) -> None:
if not hasattr(self, "value"):
self.value = values[0]
@@ -90,12 +129,23 @@ class BinaryOperatorAggregate(Generic[Value], Channel[Value, Value]):
class LastValue(Generic[Value], Channel[Value, Value]):
"""Stores the last value received."""
def _empty(self, checkpoint: Optional[str] = None) -> Self:
def __enter__(self, checkpoint: Optional[str] = None) -> Self:
empty = self.__class__()
if checkpoint is not None:
empty.value = json.loads(checkpoint)
return empty
def __exit__(
self,
__exc_type: type[BaseException] | None,
__exc_value: BaseException | None,
__traceback: TracebackType | None,
) -> None:
try:
del self.value
except AttributeError:
pass
def _update(self, values: Sequence[Value]) -> None:
if len(values) != 1:
raise InvalidUpdateError()
@@ -115,12 +165,23 @@ class LastValue(Generic[Value], Channel[Value, Value]):
class Inbox(Generic[Value], Channel[Sequence[Value], Value]):
"""Stores all values received, resets in each step."""
def _empty(self, checkpoint: Optional[str] = None) -> Self:
def __enter__(self, checkpoint: Optional[str] = None) -> Self:
empty = self.__class__()
if checkpoint is not None:
empty.queue = tuple(json.loads(checkpoint))
return empty
def __exit__(
self,
__exc_type: type[BaseException] | None,
__exc_value: BaseException | None,
__traceback: TracebackType | None,
) -> None:
try:
del self.queue
except AttributeError:
pass
def _update(self, values: Sequence[Value]) -> None:
self.queue = tuple(values)
@@ -139,12 +200,23 @@ class Set(Generic[Value], Channel[FrozenSet[Value], Value]):
set: set[Value]
def _empty(self, checkpoint: Optional[str] = None) -> Self:
def __enter__(self, checkpoint: Optional[str] = None) -> Self:
empty = self.__class__()
if checkpoint is not None:
empty.set = set(json.loads(checkpoint))
return empty
def __exit__(
self,
__exc_type: type[BaseException] | None,
__exc_value: BaseException | None,
__traceback: TracebackType | None,
) -> None:
try:
del self.set
except AttributeError:
pass
def _update(self, values: Sequence[Value]) -> None:
if not hasattr(self, "set"):
self.set = set()
permchain/pregel.py
+168 -118
View File
@@ -4,17 +4,18 @@ import asyncio
import concurrent.futures
import logging
from collections import defaultdict, deque
from contextlib import asynccontextmanager, contextmanager
from typing import (
Any,
AsyncGenerator,
AsyncIterator,
Callable,
Generic,
Generator,
Iterator,
Mapping,
Optional,
Sequence,
cast,
overload,
)
from langchain.callbacks.manager import (
@@ -40,9 +41,9 @@ from langchain.schema.runnable.config import (
get_executor_for_config,
patch_config,
)
from langchain.schema.runnable.utils import ConfigurableFieldSpec, Input, Output
from langchain.schema.runnable.utils import ConfigurableFieldSpec
from permchain.channels import Channel, EmptyChannelError, Inbox
from permchain.channels import Channel, EmptyChannelError
logger = logging.getLogger(__name__)
@@ -101,6 +102,23 @@ class PregelInvoke(RunnableBinding):
kwargs: Mapping[str, Any] = Field(default_factory=dict)
def __init__(
self,
channels: Mapping[None, str] | Mapping[str, str],
*,
bound: Optional[Runnable[Any, Any]] = None,
kwargs: Optional[Mapping[str, Any]] = None,
config: Optional[RunnableConfig] = None,
**other_kwargs: Any,
) -> None:
super().__init__(
channels=channels,
bound=bound or RunnablePassthrough(),
kwargs=kwargs or {},
config=config,
**other_kwargs,
)
def join(self, channels: Sequence[str]) -> PregelInvoke:
joiner = RunnablePassthrough.assign(
**{chan: PregelRead(chan) for chan in channels}
@@ -115,7 +133,7 @@ class PregelInvoke(RunnableBinding):
other: Runnable[Any, Other]
| Callable[[Any], Other]
| Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
) -> Runnable:
) -> PregelInvoke:
if isinstance(self.bound, RunnablePassthrough):
return PregelInvoke(channels=self.channels, bound=coerce_to_runnable(other))
else:
@@ -140,7 +158,7 @@ class PregelBatch(RunnableEach):
other: Runnable[Any, Other]
| Callable[[Any], Other]
| Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]],
) -> Runnable:
) -> PregelBatch:
if isinstance(self.bound, RunnablePassthrough):
return PregelBatch(channel=self.channel, bound=coerce_to_runnable(other))
else:
@@ -222,7 +240,31 @@ class PregelSink(RunnableLambda):
return input
class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]):
@contextmanager
def ChannelsManager(
channels: Mapping[str, Channel]
) -> Generator[Mapping[str, Channel], None, None]:
empty = {k: v.__enter__() for k, v in channels.items()}
try:
yield empty
finally:
for v in empty.values():
v.__exit__(None, None, None)
@asynccontextmanager
async def AsyncChannelsManager(
channels: Mapping[str, Channel]
) -> AsyncGenerator[Mapping[str, Channel], None]:
empty = {k: await v.__aenter__() for k, v in channels.items()}
try:
yield empty
finally:
for v in empty.values():
await v.__aexit__(None, None, None)
class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]):
channels: Mapping[str, Channel]
chains: Sequence[PregelInvoke | PregelBatch]
@@ -249,7 +291,7 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
if isinstance(chain, (list, tuple)):
chains_flat.extend(chain)
else:
chains_flat.append(chain)
chains_flat.append(cast(PregelInvoke | PregelBatch, chain))
validate_chains_channels(chains_flat, channels, input, output)
@@ -272,7 +314,7 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
)
@classmethod
def subscribe_to_each(cls, inbox: Inbox) -> PregelBatch:
def subscribe_to_each(cls, inbox: str) -> PregelBatch:
"""Runs process.batch() with the content of inbox each time it is updated."""
return PregelBatch(channel=inbox)
@@ -297,25 +339,26 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
input: Iterator[dict[str, Any] | Any],
run_manager: CallbackManagerForChainRun,
config: RunnableConfig,
) -> Iterator[Output]:
) -> Iterator[dict[str, Any] | Any]:
processes = tuple(self.chains)
# TODO this is where we'd restore from checkpoint
channels = {k: v._empty() for k, v in self.channels.items()}
next_tasks = _apply_writes_and_prepare_next_tasks(
processes,
channels,
deque((self.input, chunk) for chunk in input)
if self.input is not None
else deque((k, v) for chunk in input for k, v in chunk.items()),
)
with ChannelsManager(self.channels) as channels, get_executor_for_config(
config
) as executor:
next_tasks = _apply_writes_and_prepare_next_tasks(
processes,
channels,
deque((self.input, chunk) for chunk in input)
if self.input is not None
else deque((k, v) for chunk in input for k, v in chunk.items()),
)
def read(chan: Channel) -> Any:
try:
return channels[chan]._get()
except EmptyChannelError:
return None
def read(chan: str) -> Any:
try:
return channels[chan]._get()
except EmptyChannelError:
return None
with get_executor_for_config(config) as executor:
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
# channel updates from step N are only visible in step N+1
@@ -323,7 +366,7 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
# with channel updates applied only at the transition between steps
for step in range(config["recursion_limit"]):
# collect all writes to channels, without applying them yet
pending_writes = deque[tuple[Channel, Any]]()
pending_writes = deque[tuple[str, Any]]()
# execute tasks, and wait for one to fail or all to finish
# each task is independent from all other concurrent tasks
@@ -373,8 +416,12 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
)
# if any write to output channel in this step, yield current value
if any(chan is self.output for chan, _ in pending_writes):
yield channels[self.output]._get()
if isinstance(self.output, str):
if any(chan is self.output for chan, _ in pending_writes):
yield channels[self.output]._get()
else:
if updated := {c for c, _ in pending_writes if c in self.output}:
yield {chan: channels[chan]._get() for chan in updated}
# TODO this is where we'd save checkpoint
@@ -387,99 +434,102 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
input: AsyncIterator[dict[str, Any] | Any],
run_manager: AsyncCallbackManagerForChainRun,
config: RunnableConfig,
) -> AsyncIterator[Output]:
) -> AsyncIterator[dict[str, Any] | Any]:
processes = tuple(self.chains)
channels = {k: v._empty() for k, v in self.channels.items()}
next_tasks = _apply_writes_and_prepare_next_tasks(
processes,
channels,
deque((self.input, chunk) async for chunk in input)
if self.input is not None
else deque((k, v) async for chunk in input for k, v in chunk.items()),
)
def read(chan: Channel) -> Any:
try:
return channels[chan]._get()
except EmptyChannelError:
return None
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
# channel updates from step N are only visible in step N+1,
# channels are guaranteed to be immutable for the duration of the step,
# channel updates being applied only at the transition between steps
for step in range(config["recursion_limit"]):
# collect all writes to channels, without applying them yet
pending_writes = deque[tuple[Channel, Any]]()
# execute tasks, and wait for one to fail or all to finish
# each task is independent from all other concurrent tasks
done, inflight = await asyncio.wait(
[
asyncio.create_task(
proc.ainvoke(
input,
patch_config(
config,
callbacks=run_manager.get_child(f"pregel:step:{step}"),
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: pending_writes.extend,
CONFIG_KEY_READ: read,
CONFIG_KEY_STEP: step,
},
),
)
)
for proc, input in next_tasks
],
return_when=asyncio.FIRST_EXCEPTION,
timeout=self.step_timeout,
)
while done:
# if any task failed
if exc := done.pop().exception():
# cancel all pending tasks
while inflight:
inflight.pop().cancel()
# raise the exception
raise exc
# TODO this is where retry of an entire step would happen
if inflight:
# if we got here means we timed out
while inflight:
# cancel all pending tasks
inflight.pop().cancel()
# raise timeout error
raise TimeoutError(f"Timed out at step {step}")
# apply writes to channels, decide on next step
# TODO this is where we'd restore from checkpoint
async with AsyncChannelsManager(self.channels) as channels:
next_tasks = _apply_writes_and_prepare_next_tasks(
processes, channels, pending_writes
processes,
channels,
deque([(self.input, chunk) async for chunk in input])
if self.input is not None
else deque([(k, v) async for chunk in input for k, v in chunk.items()]),
)
# if any write to output channel in this step, yield current value
if isinstance(self.output, str):
if any(chan is self.output for chan, _ in pending_writes):
yield channels[self.output]._get()
else:
if updated := {c for c, _ in pending_writes if c in self.output}:
yield {chan: channels[chan]._get() for chan in updated}
def read(chan: str) -> Any:
try:
return channels[chan]._get()
except EmptyChannelError:
return None
# if no more tasks, we're done
if not next_tasks:
break
# Similarly to Bulk Synchronous Parallel / Pregel model
# computation proceeds in steps, while there are channel updates
# channel updates from step N are only visible in step N+1,
# channels are guaranteed to be immutable for the duration of the step,
# channel updates being applied only at the transition between steps
for step in range(config["recursion_limit"]):
# collect all writes to channels, without applying them yet
pending_writes = deque[tuple[str, Any]]()
# execute tasks, and wait for one to fail or all to finish
# each task is independent from all other concurrent tasks
done, inflight = await asyncio.wait(
[
asyncio.create_task(
proc.ainvoke(
input,
patch_config(
config,
callbacks=run_manager.get_child(
f"pregel:step:{step}"
),
configurable={
# deque.extend is thread-safe
CONFIG_KEY_SEND: pending_writes.extend,
CONFIG_KEY_READ: read,
CONFIG_KEY_STEP: step,
},
),
)
)
for proc, input in next_tasks
],
return_when=asyncio.FIRST_EXCEPTION,
timeout=self.step_timeout,
)
while done:
# if any task failed
if exc := done.pop().exception():
# cancel all pending tasks
while inflight:
inflight.pop().cancel()
# raise the exception
raise exc
# TODO this is where retry of an entire step would happen
if inflight:
# if we got here means we timed out
while inflight:
# cancel all pending tasks
inflight.pop().cancel()
# raise timeout error
raise TimeoutError(f"Timed out at step {step}")
# apply writes to channels, decide on next step
next_tasks = _apply_writes_and_prepare_next_tasks(
processes, channels, pending_writes
)
# if any write to output channel in this step, yield current value
if isinstance(self.output, str):
if any(chan is self.output for chan, _ in pending_writes):
yield channels[self.output]._get()
else:
if updated := {c for c, _ in pending_writes if c in self.output}:
yield {chan: channels[chan]._get() for chan in updated}
# if no more tasks, we're done
if not next_tasks:
break
def invoke(
self,
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
**kwargs: Any,
) -> Output:
latest: Output | None = None
) -> dict[str, Any] | Any:
latest: dict[str, Any] | Any = None
for chunk in self.stream(input, config, **kwargs):
latest = chunk
return latest
@@ -489,7 +539,7 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
**kwargs: Any,
) -> Iterator[Output]:
) -> Iterator[dict[str, Any] | Any]:
return self.transform(iter([input]), config, **kwargs)
def transform(
@@ -497,7 +547,7 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
input: Iterator[dict[str, Any] | Any],
config: RunnableConfig | None = None,
**kwargs: Any | None,
) -> Iterator[Output]:
) -> Iterator[dict[str, Any] | Any]:
return self._transform_stream_with_config(
input, self._transform, config, **kwargs
)
@@ -507,8 +557,8 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
**kwargs: Any,
) -> Output:
latest: Output | None = None
) -> dict[str, Any] | Any:
latest: dict[str, Any] | Any = None
async for chunk in self.astream(input, config, **kwargs):
latest = chunk
return latest
@@ -518,8 +568,8 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
input: dict[str, Any] | Any,
config: RunnableConfig | None = None,
**kwargs: Any,
) -> AsyncIterator[Output]:
async def input_stream() -> AsyncIterator[Input]:
) -> AsyncIterator[dict[str, Any] | Any]:
async def input_stream() -> AsyncIterator[dict[str, Any] | Any]:
yield input
async for chunk in self.atransform(input_stream(), config, **kwargs):
@@ -530,7 +580,7 @@ class Pregel(Generic[Output], RunnableSerializable[dict[str, Any] | Any, Output]
input: AsyncIterator[dict[str, Any] | Any],
config: RunnableConfig | None = None,
**kwargs: Any | None,
) -> AsyncIterator[Output]:
) -> AsyncIterator[dict[str, Any] | Any]:
async for chunk in self._atransform_stream_with_config(
input, self._atransform, config, **kwargs
):
@@ -547,7 +597,7 @@ def _apply_writes_and_prepare_next_tasks(
for chan, val in pending_writes:
pending_writes_by_channel[chan].append(val)
updated_channels: set[Channel] = set()
updated_channels: set[str] = set()
# Apply writes to channels
for chan, vals in pending_writes_by_channel.items():
if chan in channels:
@@ -595,7 +645,7 @@ def validate_chains_channels(
input: str | None,
output: str | Sequence[str],
) -> None:
subscribed_channels = set()
subscribed_channels = set[str]()
for chain in chains:
if isinstance(chain, PregelInvoke):
subscribed_channels.update(chain.channels.values())
poetry.lock
Generated
+1073 -959
View File
File diff suppressed because it is too large Load Diff
pyproject.toml
+2 -1
View File
@@ -9,7 +9,7 @@ repository = "https://www.github.com/langchain-ai/permchain"
[tool.poetry.dependencies]
python = ">=3.8.1,<4.0"
langchain = "^0.0.313"
langchain = ">=0.0.313"
[tool.poetry.group.test.dependencies]
@@ -57,6 +57,7 @@ requires = ["poetry-core>=1.0.0"]
build-backend = "poetry.core.masonry.api"
[tool.pytest.ini_options]
asyncio_mode = "auto"
# --strict-markers will raise errors on unknown marks.
# https://docs.pytest.org/en/7.1.x/how-to/mark.html#raising-errors-on-unknown-marks
#
tests/test_pregel.py
+35 -14
View File
@@ -1,14 +1,15 @@
import time
from concurrent.futures import ThreadPoolExecutor
import pytest
from pytest_mock import MockerFixture
from langchain.schema.runnable import RunnablePassthrough
from pytest_mock import MockerFixture
from permchain import Pregel, channels
from permchain.pregel import PregelInvoke
def test_invoke_single_process_in_out(mocker: MockerFixture):
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.send_to("output")
@@ -26,7 +27,7 @@ def test_invoke_single_process_in_out(mocker: MockerFixture):
assert app.invoke(2) == 3
def test_invoke_two_processes_in_out(mocker: MockerFixture):
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.send_to("inbox")
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
@@ -46,7 +47,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture):
assert pubsub.invoke(2) == 4
def test_batch_two_processes_in_out():
def test_batch_two_processes_in_out() -> None:
def add_one_with_delay(inp: int) -> int:
time.sleep(inp / 10)
return inp + 1
@@ -74,7 +75,7 @@ def test_batch_two_processes_in_out():
assert pubsub.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
def test_invoke_many_processes_in_out(mocker: MockerFixture):
def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
test_size = 100
add_one = mocker.Mock(side_effect=lambda x: x + 1)
@@ -98,8 +99,13 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture):
for _ in range(10):
assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size
with ThreadPoolExecutor() as executor:
assert [
*executor.map(app.invoke, [2] * 10, [{"recursion_limit": test_size}] * 10)
] == [2 + test_size] * 10
def test_batch_many_processes_in_out(mocker: MockerFixture):
def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
test_size = 100
add_one = mocker.Mock(side_effect=lambda x: x + 1)
@@ -130,8 +136,17 @@ def test_batch_many_processes_in_out(mocker: MockerFixture):
5 + test_size,
]
with ThreadPoolExecutor() as executor:
assert [
*executor.map(
app.batch, [[2, 1, 3, 4, 5]] * 10, [{"recursion_limit": test_size}] * 10
)
] == [
[2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size]
] * 10
def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture):
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.send_to("output")
@@ -153,7 +168,7 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture):
app.invoke(2)
def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture):
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.send_to("output")
@@ -174,7 +189,7 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture):
assert app.invoke(2) == (3, 3)
def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture):
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))
@@ -201,8 +216,11 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture):
for _ in range(100):
assert app.invoke(2) == [13, 13]
with ThreadPoolExecutor() as executor:
assert [*executor.map(app.invoke, [2] * 100)] == [[13, 13]] * 100
def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture):
def test_invoke_join_then_call_other_pubsub(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])
@@ -245,8 +263,11 @@ def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture):
for _ in range(10):
assert app.invoke([2, 3]) == 27
with ThreadPoolExecutor() as executor:
assert [*executor.map(app.invoke, [[2, 3]] * 10)] == [27] * 10
def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture):
def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
add_one = mocker.Mock(side_effect=lambda x: x + 1)
chain_one = (
@@ -272,7 +293,7 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture):
assert [c for c in app.stream(2)] == [3, 4]
def test_invoke_two_processes_no_out(mocker: MockerFixture):
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.send_to("between")
chain_two = Pregel.subscribe_to("between") | add_one
@@ -295,14 +316,14 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture):
assert app.invoke(2) is None
def test_invoke_two_processes_no_in(mocker: MockerFixture):
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.send_to("output")
chain_two = Pregel.subscribe_to("between") | add_one
with pytest.raises(ValueError):
app = Pregel(
Pregel(
chain_one,
chain_two,
channels={
tests/test_pregel_async.py
+320
View File
@@ -0,0 +1,320 @@
import asyncio
import pytest
from langchain.schema.runnable import RunnablePassthrough
from pytest_mock import MockerFixture
from permchain import Pregel, channels
from permchain.pregel import PregelInvoke
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.send_to("output")
app = Pregel(
(chain,),
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
},
input="input",
output="output",
)
assert await app.ainvoke(2) == 3
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.send_to("inbox")
chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output")
app = Pregel(
[chain_one, chain_two],
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
"inbox": channels.Inbox[int](),
},
input="input",
output="output",
)
assert await app.ainvoke(2) == 4
async def test_batch_two_processes_in_out() -> None:
async def add_one_with_delay(inp: int) -> int:
await asyncio.sleep(inp / 10)
return inp + 1
chain_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")
)
app = Pregel(
chain_one,
chain_two,
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
"one": channels.LastValue[int](),
},
input="input",
output="output",
)
assert await app.abatch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7]
async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
test_size = 100
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") | add_one | Pregel.send_to("-1")
]
for i in range(test_size - 2):
chans[str(i)] = channels.LastValue[int]()
chains.append(
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i))
)
chains.append(Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output"))
app = Pregel(*chains, channels=chans, input="input", output="output")
# No state is left over from previous invocations
for _ in range(10):
assert await app.ainvoke(2, {"recursion_limit": test_size}) == 2 + test_size
# Concurrent invocations do not interfere with each other
assert await asyncio.gather(
*(app.ainvoke(2, {"recursion_limit": test_size}) for _ in range(10))
) == [2 + test_size for _ in range(10)]
async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
test_size = 100
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") | add_one | Pregel.send_to("-1")
]
for i in range(test_size - 2):
chans[str(i)] = channels.LastValue[int]()
chains.append(
Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i))
)
chains.append(Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output"))
app = Pregel(*chains, channels=chans, input="input", output="output")
# No state is left over from previous invocations
for _ in range(10):
# Then invoke pubsub
assert await app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [
2 + test_size,
1 + test_size,
3 + test_size,
4 + test_size,
5 + test_size,
]
# Concurrent invocations do not interfere with each other
assert await asyncio.gather(
*(
app.abatch([2, 1, 3, 4, 5], {"recursion_limit": test_size})
for _ in range(10)
)
) == [
[2 + test_size, 1 + test_size, 3 + test_size, 4 + test_size, 5 + test_size]
for _ in range(10)
]
async 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.send_to("output")
chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("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
await app.ainvoke(2)
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.send_to("output")
chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("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 await app.ainvoke(2) == (3, 3)
async 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.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")
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",
)
# 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 await app.ainvoke(2) == [13, 13]
assert await asyncio.gather(*(app.ainvoke(2) for _ in range(100))) == [
[13, 13] for _ in range(100)
]
async def test_invoke_join_then_call_other_pubsub(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(
Pregel.subscribe_to("input") | add_one | Pregel.send_to("output"),
channels={
"input": channels.LastValue[int](),
"output": channels.LastValue[int](),
},
input="input",
output="output",
)
chain_one = (
Pregel.subscribe_to("input") | add_10_each | Pregel.send_to("inbox_one").map()
)
chain_two = (
Pregel.subscribe_to("inbox_one")
| inner_app.map()
| sorted
| Pregel.send_to("outbox_one")
)
chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.send_to("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 await app.ainvoke([2, 3]) == 27
assert await asyncio.gather(*(app.ainvoke([2, 3]) for _ in range(10))) == [
27 for _ in range(10)
]
async 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")
| add_one
| Pregel.send_to(output=RunnablePassthrough(), between=RunnablePassthrough())
)
chain_two = Pregel.subscribe_to("between") | add_one | Pregel.send_to("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 async for c in app.astream(2)] == [3, 4]
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.send_to("between")
chain_two = Pregel.subscribe_to("between") | add_one
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 await app.ainvoke(2) is None