diff --git a/examples/draft-revise-loop.py b/examples/draft-revise-loop.py index 48f8573da..e00db152b 100644 --- a/examples/draft-revise-loop.py +++ b/examples/draft-revise-loop.py @@ -4,8 +4,10 @@ from langchain.chat_models.openai import ChatOpenAI from langchain.output_parsers.openai_functions import JsonOutputFunctionsParser from langchain.prompts import SystemMessagePromptTemplate from langchain.schema.output_parser import StrOutputParser +from langchain.schema.runnable import RunnablePassthrough from permchain import Pregel, channels +from permchain.pregel import PregelIO # prompts @@ -62,9 +64,12 @@ gpt4 = ChatOpenAI(model="gpt-4") # chains -drafter_chain = drafter_prompt | gpt3 | StrOutputParser() - -reviser_chain = reviser_prompt | gpt3 | StrOutputParser() +drafter_chain = ( + {"question": RunnablePassthrough(input_type=str)} + | drafter_prompt + | gpt3 + | StrOutputParser() +) editor_chain = ( editor_prompt @@ -72,61 +77,51 @@ editor_chain = ( | JsonOutputFunctionsParser(args_only=False) ) +reviser_chain = reviser_prompt | gpt3 | StrOutputParser() + # application -drafter = ( - # subscribe to question channel as a dict with a single key, "question" - Pregel.subscribe_to(["question"]) - | drafter_chain - | Pregel.send_to("draft") -) - -editor = ( - # subscribe to draft channel as a dict with a single key, "draft" - Pregel.subscribe_to(["draft"]) - | editor_chain - | Pregel.send_to( - # send to "notes" channel if the editor does not accept the draft - notes=lambda x: x["arguments"]["notes"] - if x["name"] == "revise" - else None - ) -) - -reviser = ( - # subscribe to new values of "notes" channel, - # and join them with the current values of "question" and "draft" - Pregel.subscribe_to(["notes"]).join(["question", "draft"]) - | reviser_chain - | Pregel.send_to("draft") -) - draft_revise_loop = Pregel( - [drafter, reviser, editor], + drafter_chain | Pregel.send_to("draft"), + chains={ + "editor": ( + # subscribe to draft channel as a dict with a single key, "draft" + Pregel.subscribe_to(["draft"]) + | editor_chain + | Pregel.send_to( + # send to "notes" channel if the editor does not accept the draft + notes=lambda x: x["arguments"]["notes"] + if x["name"] == "revise" + else None + ) + ), + "reviser": ( + # subscribe to new values of "notes" channel, + # and join them with the input value (question) and "draft" + Pregel.subscribe_to(["notes"]).join([PregelIO.IN, "draft"]) + | reviser_chain + | Pregel.send_to("draft") + ), + }, channels={ - "question": channels.LastValue(str), "draft": channels.LastValue(str), "notes": channels.LastValue(str), }, # output will be a dict with keys "draft" and "notes" output=["draft", "notes"], - # input will be a dict with a single key, "question" - input=["question"], # debug logging debug=True, ) # run -for draft in draft_revise_loop.stream({"question": "What food do turtles eat?"}): +for draft in draft_revise_loop.stream("What food do turtles eat?"): print(draft) print("---") async def main(): - async for draft in draft_revise_loop.astream( - {"question": "What food do turtles eat?"} - ): + async for draft in draft_revise_loop.astream("What food do turtles eat?"): print(draft) print("---") diff --git a/examples/readme.py b/examples/readme.py index 67160e445..c014537a4 100644 --- a/examples/readme.py +++ b/examples/readme.py @@ -7,7 +7,7 @@ grow_value = ( ) app = Pregel( - grow_value, + chains={"grow_value": grow_value}, channels={"value": channels.LastValue(str)}, input="value", output="value", diff --git a/examples/recursive-web-loader.py b/examples/recursive-web-loader.py index 857bb40c9..cb805f16c 100644 --- a/examples/recursive-web-loader.py +++ b/examples/recursive-web-loader.py @@ -66,7 +66,7 @@ def recursive_web_loader( extractor = extractor or (lambda x: x) metadata_extractor = metadata_extractor or _metadata_extractor # the main chain that gets executed recursively - chain = ( + 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 @@ -96,10 +96,12 @@ def recursive_web_loader( ) ) return Pregel( - # use the base_url as the first url to visit - Pregel.subscribe_to("base_url") | Pregel.send_to("next_urls"), - # add the main chain - chain, + chains={ + # use the base_url as the first url to visit + "input": Pregel.subscribe_to("base_url") | Pregel.send_to("next_urls"), + # add the main chain + "visitor": visitor, + }, # define the channels channels={ "base_url": channels.LastValue(str), diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index 35bd56d25..37153598d 100644 --- a/permchain/pregel/__init__.py +++ b/permchain/pregel/__init__.py @@ -2,6 +2,7 @@ from __future__ import annotations import asyncio import concurrent.futures +import enum import logging from collections import defaultdict, deque from typing import Any, AsyncIterator, Iterator, Mapping, Optional, Sequence, Type, cast @@ -23,7 +24,6 @@ from langchain.schema.runnable.config import ( get_executor_for_config, patch_config, ) -from langchain.utils.input import get_bolded_text, get_colored_text from permchain.channels.base import ( AsyncChannelsManager, @@ -31,7 +31,14 @@ from permchain.channels.base import ( ChannelsManager, EmptyChannelError, ) -from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_STEP +from permchain.channels.last_value import LastValue +from permchain.pregel.constants import ( + CHAINS_MAIN, + CONFIG_KEY_READ, + CONFIG_KEY_SEND, + CONFIG_KEY_STEP, +) +from permchain.pregel.debug import print_step_start from permchain.pregel.read import PregelBatch, PregelInvoke from permchain.pregel.validate import validate_chains_channels from permchain.pregel.write import PregelSink @@ -39,10 +46,24 @@ from permchain.pregel.write import PregelSink logger = logging.getLogger(__name__) +# Before Python 3.11 native StrEnum is not available +class StrEnum(str, enum.Enum): + """A string enum.""" + + pass + + +class PregelIO(StrEnum): + """Pregel IO channels.""" + + IN = "__pregel_input" + OUT = "__pregel_output" + + class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): channels: Mapping[str, Channel] - chains: Sequence[PregelInvoke | PregelBatch] + chains: Mapping[str, PregelInvoke | PregelBatch] output: str | Sequence[str] @@ -57,24 +78,40 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): def __init__( self, - *chains: Sequence[PregelInvoke | PregelBatch] | PregelInvoke | PregelBatch, + main: Optional[Runnable] = None, + *, + chains: Mapping[str, PregelInvoke | PregelBatch], channels: Mapping[str, Channel], - output: str | Sequence[str], - input: str | Sequence[str], + output: str | Sequence[str] = PregelIO.OUT, + input: str | Sequence[str] = PregelIO.IN, step_timeout: Optional[float] = None, debug: Optional[bool] = None, ) -> None: - chains_flat: list[PregelInvoke | PregelBatch] = [] - for chain in chains: - if isinstance(chain, (list, tuple)): - chains_flat.extend(chain) - else: - chains_flat.append(cast(PregelInvoke | PregelBatch, chain)) + chains = {**chains} + channels = {**channels} - validate_chains_channels(chains_flat, channels, input, output) + if main is not None: + chains[CHAINS_MAIN] = ( + Pregel.subscribe_to(PregelIO.IN) | main | Pregel.send_to(PregelIO.OUT) + ) + elif output is PregelIO.OUT: + raise ValueError( + f"When no main runnable is provided, output must be one or more of the channels in {channels.keys()}" + ) + + if input is PregelIO.IN: + channels[PregelIO.IN] = LastValue( + main.input_schema if main is not None else Any # type: ignore[arg-type] + ) + if output is PregelIO.OUT: + channels[PregelIO.OUT] = LastValue( + main.output_schema if main is not None else Any # type: ignore[arg-type] + ) + + validate_chains_channels(chains, channels, input, output) super().__init__( - chains=chains_flat, + chains=chains, channels=channels, output=output, input=input, @@ -157,7 +194,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): run_manager: CallbackManagerForChainRun, config: RunnableConfig, ) -> Iterator[dict[str, Any] | Any]: - processes = tuple(self.chains) + processes = {**self.chains} # TODO this is where we'd restore from checkpoint with ChannelsManager(self.channels) as channels, get_executor_for_config( config @@ -188,16 +225,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): # with channel updates applied only at the transition between steps for step in range(config["recursion_limit"]): if self.debug: - from pprint import pformat - - n_tasks = len(next_tasks) - print( - f"{get_colored_text('[pregel/step]', color='blue')} " - + get_bolded_text( - f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Current values:\n" - ) - + pformat({k: read(k) for k in channels}) - ) + print_step_start(step, next_tasks) # collect all writes to channels, without applying them yet pending_writes = deque[tuple[str, Any]]() @@ -220,7 +248,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): }, ), ) - for proc, input in next_tasks + for proc, input, _ in next_tasks ), return_when=concurrent.futures.FIRST_EXCEPTION, timeout=self.step_timeout, @@ -251,7 +279,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): # 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): + if any(chan == 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}: @@ -269,7 +297,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): run_manager: AsyncCallbackManagerForChainRun, config: RunnableConfig, ) -> AsyncIterator[dict[str, Any] | Any]: - processes = tuple(self.chains) + processes = {**self.chains} # TODO this is where we'd restore from checkpoint async with AsyncChannelsManager(self.channels) as channels: next_tasks = _apply_writes_and_prepare_next_tasks( @@ -300,16 +328,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): # channel updates being applied only at the transition between steps for step in range(config["recursion_limit"]): if self.debug: - from pprint import pformat - - n_tasks = len(next_tasks) - print( - f"{get_colored_text('[pregel/step]', color='blue')} " - + get_bolded_text( - f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Current values:\n" - ) - + pformat({k: read(k) for k in channels}) - ) + print_step_start(step, next_tasks) # collect all writes to channels, without applying them yet pending_writes = deque[tuple[str, Any]]() @@ -335,7 +354,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): ), ) ) - for proc, input in next_tasks + for proc, input, _ in next_tasks ], return_when=asyncio.FIRST_EXCEPTION, timeout=self.step_timeout, @@ -366,7 +385,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): # 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): + if any(chan == 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}: @@ -441,10 +460,10 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): def _apply_writes_and_prepare_next_tasks( - processes: Sequence[PregelInvoke | PregelBatch], + processes: Mapping[str, PregelInvoke | PregelBatch], channels: Mapping[str, Channel], pending_writes: Sequence[tuple[str, Any]], -) -> list[tuple[Runnable, Any]]: +) -> list[tuple[Runnable, Any, str]]: pending_writes_by_channel: dict[str, list[Any]] = defaultdict(list) # Group writes by channel for chan, val in pending_writes: @@ -456,13 +475,13 @@ def _apply_writes_and_prepare_next_tasks( if chan in channels: channels[chan].update(vals) updated_channels.add(chan) - else: + elif chan != PregelIO.OUT: logger.warning(f"Skipping write for channel {chan} which has no readers") - tasks: list[tuple[Runnable, Any]] = [] + tasks: list[tuple[Runnable, Any, str]] = [] # Check if any processes should be run in next step # If so, prepare the values to be passed to them - for proc in processes: + for name, proc in processes.items(): if isinstance(proc, PregelInvoke): # If any of the channels read by this process were updated if any(chan in updated_channels for chan in proc.channels.values()): @@ -475,9 +494,9 @@ def _apply_writes_and_prepare_next_tasks( # Processes that subscribe to a single keyless channel get # the value directly, instead of a dict if list(proc.channels.keys()) == [None]: - tasks.append((proc, val[None])) + tasks.append((proc, val[None], name)) else: - tasks.append((proc, val)) + tasks.append((proc, val, name)) elif isinstance(proc, PregelBatch): # If the channel read by this process was updated if proc.channel in updated_channels: @@ -487,6 +506,6 @@ def _apply_writes_and_prepare_next_tasks( if proc.key is not None: val = [{proc.key: v} for v in val] - tasks.append((proc, val)) + tasks.append((proc, val, name)) return tasks diff --git a/permchain/pregel/constants.py b/permchain/pregel/constants.py index eda362e78..5d572c3b6 100644 --- a/permchain/pregel/constants.py +++ b/permchain/pregel/constants.py @@ -1,3 +1,5 @@ +CHAINS_MAIN = "__pregel_main" + CONFIG_KEY_STEP = "__pregel_step" CONFIG_KEY_SEND = "__pregel_send" CONFIG_KEY_READ = "__pregel_read" diff --git a/permchain/pregel/debug.py b/permchain/pregel/debug.py new file mode 100644 index 000000000..737276fbf --- /dev/null +++ b/permchain/pregel/debug.py @@ -0,0 +1,28 @@ +from pprint import pformat +from typing import Any + +from langchain.schema.runnable import Runnable +from langchain.utils.input import get_bolded_text, get_colored_text + +from permchain.pregel.constants import CHAINS_MAIN + + +def print_step_start(step: int, next_tasks: list[tuple[Runnable, Any, str]]) -> None: + n_tasks = len(next_tasks) + print( + f"{get_colored_text('[pregel/step]', color='blue')} " + + get_bolded_text( + f"Starting step {step} with {n_tasks} task{'s' if n_tasks > 1 else ''}. Next tasks:\n" + ) + + pformat( + {name if name != CHAINS_MAIN else _main: val for _, val, name in next_tasks} + ) + ) + + +class MainChainStr: + def __repr__(cls) -> str: + return "main" + + +_main = MainChainStr() diff --git a/permchain/pregel/validate.py b/permchain/pregel/validate.py index dd5d06249..aa7287641 100644 --- a/permchain/pregel/validate.py +++ b/permchain/pregel/validate.py @@ -5,13 +5,13 @@ from permchain.pregel.read import PregelBatch, PregelInvoke def validate_chains_channels( - chains: Sequence[PregelInvoke | PregelBatch], + chains: Mapping[str, PregelInvoke | PregelBatch], channels: Mapping[str, Channel], input: str | Sequence[str], output: str | Sequence[str], ) -> None: subscribed_channels = set[str]() - for chain in chains: + for chain in chains.values(): if isinstance(chain, PregelInvoke): subscribed_channels.update(chain.channels.values()) elif isinstance(chain, PregelBatch): diff --git a/tests/test_pregel.py b/tests/test_pregel.py index f9ce13708..6c9bfede2 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -4,11 +4,10 @@ from contextlib import contextmanager from typing import Generator import pytest -from langchain.schema.runnable import RunnablePassthrough +from langchain.schema.runnable import RunnableLambda, 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) -> None: @@ -16,7 +15,9 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") app = Pregel( - (chain,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -30,12 +31,26 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: assert app.invoke(2) == 3 +def test_invoke_single_process_in_out_implicit(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = RunnableLambda(add_one) + + app = Pregel(chain, channels={}, chains={}) + + # TODO add back when langchain is updated with with_types() + # assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"} + # assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"} + assert app.invoke(2) == 3 + + 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,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -58,7 +73,9 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") app = Pregel( - (chain,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -86,7 +103,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") app = Pregel( - [chain_one, chain_two], + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -105,7 +122,7 @@ def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") app = Pregel( - [chain_one, chain_two], + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -131,8 +148,7 @@ def test_batch_two_processes_in_out() -> None: ) app = Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -154,17 +170,15 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: "output": channels.LastValue(int), "-1": channels.LastValue(int), } - chains: list[PregelInvoke] = [ - Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1") - ] + chains = {"-1": 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( + chains[str(i)] = ( 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")) + chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output") - app = Pregel(*chains, channels=chans, input="input", output="output") + app = Pregel(chains=chains, channels=chans, input="input", output="output") for _ in range(10): assert app.invoke(2, {"recursion_limit": test_size}) == 2 + test_size @@ -184,17 +198,15 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: "output": channels.LastValue(int), "-1": channels.LastValue(int), } - chains: list[PregelInvoke] = [ - Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1") - ] + chains = {"-1": 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( + chains[str(i)] = ( 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")) + chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output") - app = Pregel(*chains, channels=chans, input="input", output="output") + app = Pregel(chains=chains, channels=chans, input="input", output="output") for _ in range(10): assert app.batch([2, 1, 3, 4, 5], {"recursion_limit": test_size}) == [ @@ -222,8 +234,7 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") app = Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -244,8 +255,7 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") app = Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.Inbox(int), @@ -267,9 +277,11 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.send_to("output") app = Pregel( - chain_one, - chain_three, - chain_four, + chains={ + "chain_one": chain_one, + "chain_three": chain_three, + "chain_four": chain_four, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -294,7 +306,9 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: 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"), + chains={ + "one": Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -315,9 +329,11 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.send_to("output") app = Pregel( - chain_one, - chain_two, - chain_three, + chains={ + "chain_one": chain_one, + "chain_two": chain_two, + "chain_three": chain_three, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -346,8 +362,7 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: chain_two = Pregel.subscribe_to("between") | add_one | Pregel.send_to("output") app = Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -366,8 +381,7 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: chain_two = Pregel.subscribe_to("between") | add_one app = Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -390,8 +404,7 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None: with pytest.raises(ValueError): Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -419,7 +432,7 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") app = Pregel( - [chain_one, chain_two], + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 4883b1014..b29d4e41a 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -3,11 +3,10 @@ from contextlib import asynccontextmanager, contextmanager from typing import Any, AsyncGenerator, AsyncIterator, Generator import pytest -from langchain.schema.runnable import RunnablePassthrough +from langchain.schema.runnable import RunnableLambda, 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: @@ -15,7 +14,9 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") app = Pregel( - (chain,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -27,12 +28,26 @@ 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_implicit(mocker: MockerFixture) -> None: + add_one = mocker.Mock(side_effect=lambda x: x + 1) + chain = RunnableLambda(add_one) + + app = Pregel(chain, channels={}, chains={}) + + # TODO add back when langchain is updated with with_types() + # assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"} + # assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"} + 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,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -55,7 +70,9 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") app = Pregel( - (chain,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -83,7 +100,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") app = Pregel( - [chain_one, chain_two], + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -102,7 +119,7 @@ async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") pubsub = Pregel( - [chain_one, chain_two], + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -129,8 +146,7 @@ async def test_batch_two_processes_in_out() -> None: ) app = Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -152,17 +168,15 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: "output": channels.LastValue(int), "-1": channels.LastValue(int), } - chains: list[PregelInvoke] = [ - Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1") - ] + chains = {"-1": 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( + chains[str(i)] = ( 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")) + chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output") - app = Pregel(*chains, channels=chans, input="input", output="output") + app = Pregel(chains=chains, channels=chans, input="input", output="output") # No state is left over from previous invocations for _ in range(10): @@ -183,17 +197,15 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: "output": channels.LastValue(int), "-1": channels.LastValue(int), } - chains: list[PregelInvoke] = [ - Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1") - ] + chains = {"-1": 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( + chains[str(i)] = ( 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")) + chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output") - app = Pregel(*chains, channels=chans, input="input", output="output") + app = Pregel(chains=chains, channels=chans, input="input", output="output") # No state is left over from previous invocations for _ in range(10): @@ -227,8 +239,7 @@ async def test_invoke_two_processes_two_in_two_out_invalid( chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") app = Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -249,8 +260,7 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") app = Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.Inbox(int), @@ -272,9 +282,11 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) - chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.send_to("output") app = Pregel( - chain_one, - chain_three, - chain_four, + chains={ + "chain_one": chain_one, + "chain_three": chain_three, + "chain_four": chain_four, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -300,7 +312,9 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None 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"), + chains={ + "one": Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -321,9 +335,11 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.send_to("output") app = Pregel( - chain_one, - chain_two, - chain_three, + chains={ + "chain_one": chain_one, + "chain_two": chain_two, + "chain_three": chain_three, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -354,8 +370,7 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non chain_two = Pregel.subscribe_to("between") | add_one | Pregel.send_to("output") app = Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -375,8 +390,7 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None: chain_two = Pregel.subscribe_to("between") | add_one app = Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -419,7 +433,7 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") app = Pregel( - [chain_one, chain_two], + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int),