From 8adfa694053075c6bd995ffd465554aa5f5cba9e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Thu, 19 Oct 2023 12:57:26 +0100 Subject: [PATCH] Implement dict input, add more tests --- README.md | 4 +- permchain/channels/__init__.py | 4 +- permchain/pregel.py | 43 ++++++++--- pyproject.toml | 2 +- tests/test_pregel.py | 106 ++++++++++++++++++++++--- tests/test_pregel_async.py | 137 +++++++++++++++++++++++++++++++++ 6 files changed, 270 insertions(+), 26 deletions(-) diff --git a/README.md b/README.md index dc4f0a315..278e04042 100644 --- a/README.md +++ b/README.md @@ -36,8 +36,8 @@ Check `examples` for more examples. - [x] Finish updating tests to new API - [x] Implement input_schema and output_schema in Pregel - [ ] More tests - - [ ] Test different input and output types (str, str sequence, None) - - [ ] Add tests for Stream, UniqueInbox + - [x] Test different input and output types (str, str sequence) + - [x] Add tests for Stream, UniqueInbox - [ ] Add tests for subscribe_to_each().join() - [ ] Implement checkpointing - [ ] Save checkpoints at end of each step diff --git a/permchain/channels/__init__.py b/permchain/channels/__init__.py index 23bca3a44..9fe8df138 100644 --- a/permchain/channels/__init__.py +++ b/permchain/channels/__init__.py @@ -1,9 +1,9 @@ from permchain.channels.base import Channel, EmptyChannelError, InvalidUpdateError +from permchain.channels.binop import BinaryOperatorAggregate +from permchain.channels.context import ContextManager from permchain.channels.inbox import Inbox, UniqueInbox from permchain.channels.last_value import LastValue -from permchain.channels.binop import BinaryOperatorAggregate from permchain.channels.stream import Set, Stream -from permchain.channels.context import ContextManager __all__ = [ "Channel", diff --git a/permchain/pregel.py b/permchain/pregel.py index 4934ebf2b..e47695886 100644 --- a/permchain/pregel.py +++ b/permchain/pregel.py @@ -288,7 +288,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): output: str | Sequence[str] - input: str | None + input: str | Sequence[str] step_timeout: Optional[float] = None @@ -300,7 +300,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): *chains: Sequence[PregelInvoke | PregelBatch] | PregelInvoke | PregelBatch, channels: Mapping[str, Channel], output: str | Sequence[str], - input: str | None = None, + input: str | Sequence[str], step_timeout: Optional[float] = None, ) -> None: chains_flat: list[PregelInvoke | PregelBatch] = [] @@ -402,8 +402,13 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): 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()), + if isinstance(self.input, str) + else deque( + (k, v) + for chunk in input + for k, v in chunk.items() + if k in self.input + ), ) def read(chan: str) -> Any: @@ -495,8 +500,15 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): 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 isinstance(self.input, str) + else deque( + [ + (k, v) + async for chunk in input + for k, v in chunk.items() + if k in self.input + ] + ), ) def read(chan: str) -> Any: @@ -695,7 +707,7 @@ def _apply_writes_and_prepare_next_tasks( def validate_chains_channels( chains: Sequence[PregelInvoke | PregelBatch], channels: Mapping[str, Channel], - input: str | None, + input: str | Sequence[str], output: str | Sequence[str], ) -> None: subscribed_channels = set[str]() @@ -709,13 +721,24 @@ def validate_chains_channels( 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(input, str): + if input not in subscribed_channels: + raise ValueError(f"Input channel {input} is not subscribed to by any chain") + else: + for chan in input: + if chan not in subscribed_channels: + raise ValueError( + f"Input channel {chan} is not subscribed to by any chain" + ) + if isinstance(output, str): if output not in channels: raise ValueError(f"Output channel {output} is not initialized") + else: + for chan in output: + if chan not in channels: + raise ValueError(f"Output channel {chan} is not initialized") diff --git a/pyproject.toml b/pyproject.toml index 8685cfdf3..b02d0bc85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,6 +67,6 @@ asyncio_mode = "auto" # # https://github.com/tophat/syrupy # --snapshot-warn-unused Prints a warning on unused snapshots rather than fail the test suite. -addopts = "--full-trace --strict-markers --strict-config --durations=5 --snapshot-warn-unused" +addopts = "-x --full-trace --strict-markers --strict-config --durations=5 --snapshot-warn-unused" # Registering custom markers. # https://docs.pytest.org/en/7.1.x/example/markers.html#registering-markers diff --git a/tests/test_pregel.py b/tests/test_pregel.py index e8f7f727a..f9ce13708 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -1,5 +1,7 @@ import time from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from typing import Generator import pytest from langchain.schema.runnable import RunnablePassthrough @@ -51,12 +53,39 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: assert app.invoke(2) == {"output": 3} +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.send_to("output") + + app = Pregel( + (chain,), + channels={ + "input": channels.LastValue(int), + "output": channels.LastValue(int), + }, + input=["input"], + output=["output"], + ) + + assert app.input_schema.schema() == { + "title": "PregelInput", + "type": "object", + "properties": {"input": {"title": "Input", "type": "integer"}}, + } + assert app.output_schema.schema() == { + "title": "PregelOutput", + "type": "object", + "properties": {"output": {"title": "Output", "type": "integer"}}, + } + assert app.invoke({"input": 2}) == {"output": 3} + + 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") - pubsub = Pregel( + app = Pregel( [chain_one, chain_two], channels={ "input": channels.LastValue(int), @@ -67,8 +96,26 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: output="output", ) - # Then invoke pubsub - assert pubsub.invoke(2) == 4 + assert app.invoke(2) == 4 + + +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.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", "inbox"], + output="output", + ) + + assert [*app.stream({"input": 2, "inbox": 12})] == [13, 4] # [12 + 1, 2 + 1 + 1] def test_batch_two_processes_in_out() -> None: @@ -83,7 +130,7 @@ def test_batch_two_processes_in_out() -> None: Pregel.subscribe_to("one") | add_one_with_delay | Pregel.send_to("output") ) - pubsub = Pregel( + app = Pregel( chain_one, chain_two, channels={ @@ -95,8 +142,7 @@ def test_batch_two_processes_in_out() -> None: output="output", ) - # Then invoke pubsub - assert pubsub.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] + assert app.batch([3, 2, 1, 3, 5]) == [5, 4, 3, 5, 7] def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: @@ -151,7 +197,6 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: app = Pregel(*chains, channels=chans, input="input", output="output") for _ in range(10): - # Then invoke pubsub assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ 2 + test_size, 1 + test_size, @@ -244,7 +289,7 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None assert [*executor.map(app.invoke, [2] * 100)] == [[13, 13]] * 100 -def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None: +def test_invoke_join_then_call_other_app(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]) @@ -283,7 +328,6 @@ def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None: output="output", ) - # Then invoke pubsub for _ in range(10): assert app.invoke([2, 3]) == 27 @@ -313,7 +357,6 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: output="output", ) - # Then invoke pubsub assert [c for c in app.stream(2)] == [3, 4] @@ -334,7 +377,6 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: 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 app.invoke(2) is None @@ -358,3 +400,45 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None: input="input", output="output", ) + + +def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: + setup = mocker.Mock() + cleanup = mocker.Mock() + + @contextmanager + def an_int() -> Generator[int, None, None]: + setup() + try: + yield 5 + finally: + cleanup() + + 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), + "ctx": channels.ContextManager(an_int, typ=int), + }, + input="input", + output=["inbox", "output"], + ) + + assert setup.call_count == 0 + assert cleanup.call_count == 0 + for i, chunk in enumerate(app.stream(2)): + assert setup.call_count == 1, "Expected setup to be called once" + assert cleanup.call_count == 0, "Expected cleanup to not be called yet" + if i == 0: + assert chunk == {"inbox": (3,)} + elif i == 1: + assert chunk == {"output": 4} + else: + assert False, "Expected only two chunks" + assert cleanup.call_count == 1, "Expected cleanup to be called once" diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 4181473de..4883b1014 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -1,4 +1,6 @@ import asyncio +from contextlib import asynccontextmanager, contextmanager +from typing import Any, AsyncGenerator, AsyncIterator, Generator import pytest from langchain.schema.runnable import RunnablePassthrough @@ -25,6 +27,56 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: assert await app.ainvoke(2) == 3 +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.send_to("output") + + app = Pregel( + (chain,), + channels={ + "input": channels.LastValue(int), + "output": channels.LastValue(int), + }, + input="input", + output=["output"], + ) + + assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"} + assert app.output_schema.schema() == { + "title": "PregelOutput", + "type": "object", + "properties": {"output": {"title": "Output", "type": "integer"}}, + } + assert await app.ainvoke(2) == {"output": 3} + + +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.send_to("output") + + app = Pregel( + (chain,), + channels={ + "input": channels.LastValue(int), + "output": channels.LastValue(int), + }, + input=["input"], + output=["output"], + ) + + assert app.input_schema.schema() == { + "title": "PregelInput", + "type": "object", + "properties": {"input": {"title": "Input", "type": "integer"}}, + } + assert app.output_schema.schema() == { + "title": "PregelOutput", + "type": "object", + "properties": {"output": {"title": "Output", "type": "integer"}}, + } + assert await app.ainvoke({"input": 2}) == {"output": 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") @@ -44,6 +96,26 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert await app.ainvoke(2) == 4 +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.send_to("inbox") + chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") + + pubsub = Pregel( + [chain_one, chain_two], + channels={ + "input": channels.LastValue(int), + "output": channels.LastValue(int), + "inbox": channels.Inbox(int), + }, + input=["input", "inbox"], + output="output", + ) + + # [12 + 1, 2 + 1 + 1] + assert [c async for c in pubsub.astream({"input": 2, "inbox": 12})] == [13, 4] + + async def test_batch_two_processes_in_out() -> None: async def add_one_with_delay(inp: int) -> int: await asyncio.sleep(inp / 10) @@ -318,3 +390,68 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: # 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 + + +async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: + setup_sync = mocker.Mock() + cleanup_sync = mocker.Mock() + setup_async = mocker.Mock() + cleanup_async = mocker.Mock() + + @contextmanager + def an_int() -> Generator[int, None, None]: + setup_sync() + try: + yield 5 + finally: + cleanup_sync() + + @asynccontextmanager + async def an_int_async() -> AsyncGenerator[int, None]: + setup_async() + try: + yield 5 + finally: + cleanup_async() + + 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), + "ctx": channels.ContextManager(an_int, an_int_async, typ=int), + }, + input="input", + output=["inbox", "output"], + ) + + async def aenumerate(aiter: AsyncIterator[Any]) -> AsyncIterator[tuple[int, Any]]: + i = 0 + async for chunk in aiter: + yield i, chunk + i += 1 + + assert setup_sync.call_count == 0 + assert cleanup_sync.call_count == 0 + assert setup_async.call_count == 0 + assert cleanup_async.call_count == 0 + async for i, chunk in aenumerate(app.astream(2)): + assert setup_sync.call_count == 0, "Sync context manager should not be used" + assert cleanup_sync.call_count == 0, "Sync context manager should not be used" + assert setup_async.call_count == 1, "Expected setup to be called once" + assert cleanup_async.call_count == 0, "Expected cleanup to not be called yet" + if i == 0: + assert chunk == {"inbox": (3,)} + elif i == 1: + assert chunk == {"output": 4} + else: + assert False, "Expected only two chunks" + assert setup_sync.call_count == 0 + assert cleanup_sync.call_count == 0 + assert setup_async.call_count == 1, "Expected setup to be called once" + assert cleanup_async.call_count == 1, "Expected cleanup to be called once"