From 9dfffe5fc0963588b25d0f78159c13aafc6625a2 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sun, 22 Oct 2023 21:12:30 +0100 Subject: [PATCH 1/4] Add an optional "main" chain which has implicit input and output channels --- examples/draft-revise-loop.py | 71 +++++++++---------- examples/readme.py | 2 +- examples/recursive-web-loader.py | 12 ++-- permchain/pregel/__init__.py | 115 ++++++++++++++++++------------- permchain/pregel/constants.py | 2 + permchain/pregel/debug.py | 28 ++++++++ permchain/pregel/validate.py | 4 +- tests/test_pregel.py | 91 +++++++++++++----------- tests/test_pregel_async.py | 88 +++++++++++++---------- 9 files changed, 243 insertions(+), 170 deletions(-) create mode 100644 permchain/pregel/debug.py 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), From 274600436b61f51bf86e9653bd3bc204e4453f93 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 23 Oct 2023 10:26:10 +0100 Subject: [PATCH 2/4] Remove "main" chain, Improve debug logging, Remove max_steps from send_to, Add retries to recursive web loader --- README.md | 2 +- examples/draft-revise-loop.py | 63 ++++++---- examples/recursive-web-loader.py | 28 ++++- permchain/pregel/__init__.py | 198 +++++++++---------------------- permchain/pregel/constants.py | 3 - permchain/pregel/debug.py | 24 ++-- permchain/pregel/io.py | 34 ++++++ permchain/pregel/log.py | 3 + permchain/pregel/write.py | 25 +--- tests/test_pregel.py | 14 +-- tests/test_pregel_async.py | 14 +-- 11 files changed, 176 insertions(+), 232 deletions(-) create mode 100644 permchain/pregel/io.py create mode 100644 permchain/pregel/log.py diff --git a/README.md b/README.md index cc1c3bd09..2d2f3beef 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,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/draft-revise-loop.py b/examples/draft-revise-loop.py index e00db152b..e9d6b534d 100644 --- a/examples/draft-revise-loop.py +++ b/examples/draft-revise-loop.py @@ -7,7 +7,6 @@ from langchain.schema.output_parser import StrOutputParser from langchain.schema.runnable import RunnablePassthrough from permchain import Pregel, channels -from permchain.pregel import PregelIO # prompts @@ -81,32 +80,46 @@ 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 input value (question) and "draft" + Pregel.subscribe_to(["notes"]).join(["question", "draft"]) + | reviser_chain + | Pregel.send_to("draft") +) + draft_revise_loop = Pregel( - 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), }, + chains={ + "drafter": drafter, + "editor": editor, + "reviser": reviser, + }, + # input will be a dict with a single key, "question" + input=["question"], # output will be a dict with keys "draft" and "notes" output=["draft", "notes"], # debug logging @@ -115,13 +128,15 @@ draft_revise_loop = Pregel( # run -for draft in draft_revise_loop.stream("What food do turtles eat?"): +for draft in draft_revise_loop.stream({"question": "What food do turtles eat?"}): print(draft) print("---") async def main(): - async for draft in draft_revise_loop.astream("What food do turtles eat?"): + async for draft in draft_revise_loop.astream( + {"question": "What food do turtles eat?"} + ): print(draft) print("---") diff --git a/examples/recursive-web-loader.py b/examples/recursive-web-loader.py index cb805f16c..125933cd1 100644 --- a/examples/recursive-web-loader.py +++ b/examples/recursive-web-loader.py @@ -1,4 +1,5 @@ -from typing import Callable, FrozenSet, Optional, TypedDict +from contextlib import asynccontextmanager, contextmanager +from typing import AsyncGenerator, Callable, FrozenSet, Generator, Optional, TypedDict import httpx from langchain.schema import Document @@ -10,6 +11,14 @@ from permchain import Pregel, channels # Load url with sync httpx client +@contextmanager +def httpx_client() -> Generator[httpx.Client, None, None]: + with httpx.HTTPTransport(retries=3) as transport, httpx.Client( + transport=transport + ) as client: + yield client + + class LoadUrlInput(TypedDict): url: str visited: FrozenSet[str] @@ -24,6 +33,14 @@ def load_url(input: LoadUrlInput) -> str: # Same as above but with async httpx client +@asynccontextmanager +async def httpx_aclient() -> AsyncGenerator[httpx.AsyncClient, None]: + async with httpx.AsyncHTTPTransport(retries=3) as transport, httpx.AsyncClient( + transport=transport + ) as client: + yield client + + class LoadUrlInputAsync(TypedDict): url: str visited: FrozenSet[str] @@ -92,7 +109,6 @@ def recursive_web_loader( ) if url not in x["visited"] and url != x["url"] ], - _max_steps=max_depth, ) ) return Pregel( @@ -108,17 +124,19 @@ def recursive_web_loader( "next_urls": channels.UniqueInbox(str), "documents": channels.Stream(Document), "visited": channels.Set(str), - "client": channels.ContextManager(httpx.Client, httpx.AsyncClient), + "client": channels.ContextManager(httpx_client, httpx_aclient), }, # this will accept a string as input input="base_url", # and return a dict with documents and visited set output=["documents", "visited"], - ) + # debug logging + debug=True, + ).with_config({"recursion_limit": max_depth + 1}) loader = recursive_web_loader(max_depth=3) documents = loader.invoke("https://docs.python.org/3.9/") -print(documents) +print(len(documents["documents"])) diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index 37153598d..1474b3ce0 100644 --- a/permchain/pregel/__init__.py +++ b/permchain/pregel/__init__.py @@ -2,8 +2,6 @@ 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 @@ -12,7 +10,7 @@ from langchain.callbacks.manager import ( CallbackManagerForChainRun, ) from langchain.globals import get_debug -from langchain.pydantic_v1 import BaseModel, create_model +from langchain.pydantic_v1 import BaseModel, Field, create_model, root_validator from langchain.schema.runnable import ( Runnable, RunnablePassthrough, @@ -31,34 +29,14 @@ from permchain.channels.base import ( ChannelsManager, EmptyChannelError, ) -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.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND +from permchain.pregel.debug import print_checkpoint, print_step_start +from permchain.pregel.io import map_input, map_output +from permchain.pregel.log import logger from permchain.pregel.read import PregelBatch, PregelInvoke from permchain.pregel.validate import validate_chains_channels 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] @@ -71,53 +49,17 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): step_timeout: Optional[float] = None - debug: bool + debug: bool = Field(default_factory=get_debug) class Config: arbitrary_types_allowed = True - def __init__( - self, - main: Optional[Runnable] = None, - *, - chains: Mapping[str, PregelInvoke | PregelBatch], - channels: Mapping[str, Channel], - output: str | Sequence[str] = PregelIO.OUT, - input: str | Sequence[str] = PregelIO.IN, - step_timeout: Optional[float] = None, - debug: Optional[bool] = None, - ) -> None: - chains = {**chains} - channels = {**channels} - - 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, - channels=channels, - output=output, - input=input, - step_timeout=step_timeout, - debug=debug if debug is not None else get_debug(), + @root_validator(skip_on_failure=True) + def validate_pregel(cls, values: dict[str, Any]) -> dict[str, Any]: + validate_chains_channels( + values["chains"], values["channels"], values["input"], values["output"] ) + return values @property def InputType(self) -> Any: @@ -176,7 +118,6 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): def send_to( cls, *channels: str, - _max_steps: Optional[int] = None, **kwargs: RunnableLike, ) -> PregelSink: """Writes to channels the result of the lambda, or None to skip writing.""" @@ -184,8 +125,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): channels=( [(c, RunnablePassthrough()) for c in channels] + [(k, coerce_to_runnable(v)) for k, v in kwargs.items()] - ), - max_steps=_max_steps, + ) ) def _transform( @@ -202,14 +142,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): next_tasks = _apply_writes_and_prepare_next_tasks( processes, channels, - deque((self.input, chunk) for chunk in input) - if isinstance(self.input, str) - else deque( - (k, v) - for chunk in input - for k, v in chunk.items() - if k in self.input - ), + deque(w for c in input for w in map_input(self.input, c)), ) def read(chan: str) -> Any: @@ -244,7 +177,6 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): # deque.extend is thread-safe CONFIG_KEY_SEND: pending_writes.extend, CONFIG_KEY_READ: read, - CONFIG_KEY_STEP: step, }, ), ) @@ -254,36 +186,20 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): 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}") + # interrupt on failure or timeout + _interrupt_or_proceed(done, inflight, 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 == 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 self.debug: + print_checkpoint(step, channels) + + # if any write to output channels in this step, yield current value + for output in map_output(self.output, pending_writes, channels): + yield output # TODO this is where we'd save checkpoint @@ -303,16 +219,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): next_tasks = _apply_writes_and_prepare_next_tasks( processes, channels, - deque([(self.input, chunk) async for chunk in input]) - 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 - ] - ), + deque([w async for c in input for w in map_input(self.input, c)]), ) def read(chan: str) -> Any: @@ -349,7 +256,6 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): # deque.extend is thread-safe CONFIG_KEY_SEND: pending_writes.extend, CONFIG_KEY_READ: read, - CONFIG_KEY_STEP: step, }, ), ) @@ -360,36 +266,22 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): 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}") + # interrupt on failure or timeout + _interrupt_or_proceed(done, inflight, 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 == 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 self.debug: + print_checkpoint(step, channels) + + # if any write to output channels in this step, yield current value + for output in map_output(self.output, pending_writes, channels): + yield output + + # TODO this is where we'd save checkpoint # if no more tasks, we're done if not next_tasks: @@ -459,6 +351,30 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): yield chunk +def _interrupt_or_proceed( + done: set[concurrent.futures.Future[Any]] | set[asyncio.Task[Any]], + inflight: set[concurrent.futures.Future[Any]] | set[asyncio.Task[Any]], + step: int, +) -> None: + 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}") + + def _apply_writes_and_prepare_next_tasks( processes: Mapping[str, PregelInvoke | PregelBatch], channels: Mapping[str, Channel], @@ -475,7 +391,7 @@ def _apply_writes_and_prepare_next_tasks( if chan in channels: channels[chan].update(vals) updated_channels.add(chan) - elif chan != PregelIO.OUT: + else: logger.warning(f"Skipping write for channel {chan} which has no readers") tasks: list[tuple[Runnable, Any, str]] = [] diff --git a/permchain/pregel/constants.py b/permchain/pregel/constants.py index 5d572c3b6..6e5f9c37f 100644 --- a/permchain/pregel/constants.py +++ b/permchain/pregel/constants.py @@ -1,5 +1,2 @@ -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 index 737276fbf..636512bcc 100644 --- a/permchain/pregel/debug.py +++ b/permchain/pregel/debug.py @@ -1,10 +1,10 @@ from pprint import pformat -from typing import Any +from typing import Any, Iterator, Mapping from langchain.schema.runnable import Runnable from langchain.utils.input import get_bolded_text, get_colored_text -from permchain.pregel.constants import CHAINS_MAIN +from permchain.channels.base import Channel, EmptyChannelError def print_step_start(step: int, next_tasks: list[tuple[Runnable, Any, str]]) -> None: @@ -14,15 +14,21 @@ def print_step_start(step: int, next_tasks: list[tuple[Runnable, Any, str]]) -> + 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} - ) + + "\n".join(f"- {name}({pformat(val)})" for _, val, name in next_tasks) ) -class MainChainStr: - def __repr__(cls) -> str: - return "main" +def print_checkpoint(step: int, channels: Mapping[str, Channel]) -> None: + print( + f"{get_colored_text('[pregel/checkpoint]', color='blue')} " + + get_bolded_text(f"Finishing step {step}. Channel values:\n") + + pformat({name: val for name, val in _read_channels(channels)}, depth=1) + ) -_main = MainChainStr() +def _read_channels(channels: Mapping[str, Channel]) -> Iterator[tuple[str, Any]]: + for name, channel in channels.items(): + try: + yield (name, channel.get()) + except EmptyChannelError: + pass diff --git a/permchain/pregel/io.py b/permchain/pregel/io.py new file mode 100644 index 000000000..ddb0a4688 --- /dev/null +++ b/permchain/pregel/io.py @@ -0,0 +1,34 @@ +from typing import Any, Iterator, Mapping, Sequence + +from permchain.channels.base import Channel +from permchain.pregel.log import logger + + +def map_input( + input_channels: str | Sequence[str], chunk: dict[str, Any] | Any +) -> Iterator[tuple[str, Any]]: + """Map input chunk to a sequence of pending writes in the form (channel, value).""" + if isinstance(input_channels, str): + yield (input_channels, chunk) + else: + if not isinstance(chunk, dict): + raise TypeError(f"Expected chunk to be a dict, got {type(chunk).__name__}") + for k in chunk: + if k in input_channels: + yield (k, chunk[k]) + else: + logger.warning(f"Input channel {k} not found in {input_channels}") + + +def map_output( + output_channels: str | Sequence[str], + pending_writes: Sequence[tuple[str, Any]], + channels: Mapping[str, Channel], +) -> Iterator[dict[str, Any] | Any]: + """Map pending writes (a sequence of tuples (channel, value)) to output chunk.""" + if isinstance(output_channels, str): + if any(chan == output_channels for chan, _ in pending_writes): + yield channels[output_channels].get() + else: + if updated := {c for c, _ in pending_writes if c in output_channels}: + yield {chan: channels[chan].get() for chan in updated} diff --git a/permchain/pregel/log.py b/permchain/pregel/log.py new file mode 100644 index 000000000..eea436a37 --- /dev/null +++ b/permchain/pregel/log.py @@ -0,0 +1,3 @@ +import logging + +logger = logging.getLogger(__name__) diff --git a/permchain/pregel/write.py b/permchain/pregel/write.py index e04d93f2b..6d2b973ac 100644 --- a/permchain/pregel/write.py +++ b/permchain/pregel/write.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import Any, Callable, Optional, Sequence +from typing import Any, Callable, Sequence from langchain.schema.runnable import ( Runnable, @@ -9,7 +9,7 @@ from langchain.schema.runnable import ( ) from langchain.schema.runnable.utils import ConfigurableFieldSpec -from permchain.pregel.constants import CONFIG_KEY_SEND, CONFIG_KEY_STEP +from permchain.pregel.constants import CONFIG_KEY_SEND TYPE_SEND = Callable[[Sequence[tuple[str, Any]]], None] @@ -22,28 +22,17 @@ class PregelSink(RunnableLambda): or None to skip writing. """ - max_steps: Optional[int] - def __init__( self, *, channels: Sequence[tuple[str, Runnable]], - max_steps: Optional[int] = None, ): super().__init__(func=self._write, afunc=self._awrite) # type: ignore[arg-type] self.channels = channels - self.max_steps = max_steps @property def config_specs(self) -> Sequence[ConfigurableFieldSpec]: return [ - ConfigurableFieldSpec( - id=CONFIG_KEY_STEP, - name=CONFIG_KEY_STEP, - description=None, - default=None, - annotation=int, - ), ConfigurableFieldSpec( id=CONFIG_KEY_SEND, name=CONFIG_KEY_SEND, @@ -54,11 +43,6 @@ class PregelSink(RunnableLambda): ] def _write(self, input: Any, config: RunnableConfig) -> None: - step: int = config["configurable"][CONFIG_KEY_STEP] - - if self.max_steps is not None and step >= self.max_steps: - return - write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND] values = [(chan, r.invoke(input, config)) for chan, r in self.channels] @@ -68,11 +52,6 @@ class PregelSink(RunnableLambda): return input async def _awrite(self, input: Any, config: RunnableConfig) -> None: - step: int = config["configurable"][CONFIG_KEY_STEP] - - if self.max_steps is not None and step >= self.max_steps: - return - write: TYPE_SEND = config["configurable"][CONFIG_KEY_SEND] values = [(chan, await r.ainvoke(input, config)) for chan, r in self.channels] diff --git a/tests/test_pregel.py b/tests/test_pregel.py index 6c9bfede2..bb8ab9f12 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -4,7 +4,7 @@ from contextlib import contextmanager from typing import Generator import pytest -from langchain.schema.runnable import RunnableLambda, RunnablePassthrough +from langchain.schema.runnable import RunnablePassthrough from pytest_mock import MockerFixture from permchain import Pregel, channels @@ -31,18 +31,6 @@ 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") diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index b29d4e41a..7eb8e381f 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -3,7 +3,7 @@ from contextlib import asynccontextmanager, contextmanager from typing import Any, AsyncGenerator, AsyncIterator, Generator import pytest -from langchain.schema.runnable import RunnableLambda, RunnablePassthrough +from langchain.schema.runnable import RunnablePassthrough from pytest_mock import MockerFixture from permchain import Pregel, channels @@ -28,18 +28,6 @@ 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") From 8284a6b6eaa24122e1e7bf7ad960ba7870374490 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 23 Oct 2023 10:29:46 +0100 Subject: [PATCH 3/4] Rename send_to to write_to --- README.md | 2 +- examples/draft-revise-loop.py | 6 +-- examples/readme.py | 2 +- examples/recursive-web-loader.py | 4 +- permchain/pregel/__init__.py | 2 +- tests/test_pregel.py | 64 ++++++++++++++++---------------- tests/test_pregel_async.py | 62 +++++++++++++++---------------- 7 files changed, 71 insertions(+), 71 deletions(-) diff --git a/README.md b/README.md index 2d2f3beef..e74289dd7 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ from permchain import Pregel, channels grow_value = ( Pregel.subscribe_to("value") | (lambda x: x + x) - | Pregel.send_to(value=lambda x: x if len(x) < 10 else None) + | Pregel.write_to(value=lambda x: x if len(x) < 10 else None) ) app = Pregel( diff --git a/examples/draft-revise-loop.py b/examples/draft-revise-loop.py index e9d6b534d..11b073185 100644 --- a/examples/draft-revise-loop.py +++ b/examples/draft-revise-loop.py @@ -84,14 +84,14 @@ drafter = ( # subscribe to question channel as a dict with a single key, "question" Pregel.subscribe_to(["question"]) | drafter_chain - | Pregel.send_to("draft") + | Pregel.write_to("draft") ) editor = ( # subscribe to draft channel as a dict with a single key, "draft" Pregel.subscribe_to(["draft"]) | editor_chain - | Pregel.send_to( + | Pregel.write_to( # send to "notes" channel if the editor does not accept the draft notes=lambda x: x["arguments"]["notes"] if x["name"] == "revise" @@ -104,7 +104,7 @@ reviser = ( # and join them with the input value (question) and "draft" Pregel.subscribe_to(["notes"]).join(["question", "draft"]) | reviser_chain - | Pregel.send_to("draft") + | Pregel.write_to("draft") ) draft_revise_loop = Pregel( diff --git a/examples/readme.py b/examples/readme.py index c014537a4..97998b3ce 100644 --- a/examples/readme.py +++ b/examples/readme.py @@ -3,7 +3,7 @@ from permchain import Pregel, channels grow_value = ( Pregel.subscribe_to("value") | (lambda x: x + x) - | Pregel.send_to(value=lambda x: x if len(x) < 10 else None) + | Pregel.write_to(value=lambda x: x if len(x) < 10 else None) ) app = Pregel( diff --git a/examples/recursive-web-loader.py b/examples/recursive-web-loader.py index 125933cd1..cbba99a7f 100644 --- a/examples/recursive-web-loader.py +++ b/examples/recursive-web-loader.py @@ -92,7 +92,7 @@ def recursive_web_loader( ) # load the url (with sync and async implementations) | RunnablePassthrough.assign(body=RunnableLambda(load_url, load_url_async)) - | Pregel.send_to( + | Pregel.write_to( # send this url to the visited set visited=lambda x: x["url"], # send a new document to the documents stream @@ -114,7 +114,7 @@ def recursive_web_loader( return Pregel( chains={ # use the base_url as the first url to visit - "input": Pregel.subscribe_to("base_url") | Pregel.send_to("next_urls"), + "input": Pregel.subscribe_to("base_url") | Pregel.write_to("next_urls"), # add the main chain "visitor": visitor, }, diff --git a/permchain/pregel/__init__.py b/permchain/pregel/__init__.py index 1474b3ce0..06784eb37 100644 --- a/permchain/pregel/__init__.py +++ b/permchain/pregel/__init__.py @@ -115,7 +115,7 @@ class Pregel(RunnableSerializable[dict[str, Any] | Any, dict[str, Any] | Any]): return PregelBatch(channel=inbox, key=key) @classmethod - def send_to( + def write_to( cls, *channels: str, **kwargs: RunnableLike, diff --git a/tests/test_pregel.py b/tests/test_pregel.py index bb8ab9f12..975a10748 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -12,7 +12,7 @@ from permchain import Pregel, channels 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") + chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") app = Pregel( chains={ @@ -33,7 +33,7 @@ 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.send_to("output") + chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") app = Pregel( chains={ @@ -58,7 +58,7 @@ 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.send_to("output") + chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") app = Pregel( chains={ @@ -87,8 +87,8 @@ 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.send_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") + chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, @@ -106,8 +106,8 @@ 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.send_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") + chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, @@ -129,10 +129,10 @@ def test_batch_two_processes_in_out() -> None: return inp + 1 chain_one = ( - Pregel.subscribe_to("input") | add_one_with_delay | Pregel.send_to("one") + Pregel.subscribe_to("input") | add_one_with_delay | Pregel.write_to("one") ) chain_two = ( - Pregel.subscribe_to("one") | add_one_with_delay | Pregel.send_to("output") + Pregel.subscribe_to("one") | add_one_with_delay | Pregel.write_to("output") ) app = Pregel( @@ -158,13 +158,13 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: "output": channels.LastValue(int), "-1": channels.LastValue(int), } - chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1")} + chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")} for i in range(test_size - 2): chans[str(i)] = channels.LastValue(int) chains[str(i)] = ( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i)) + Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i)) ) - chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output") + chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output") app = Pregel(chains=chains, channels=chans, input="input", output="output") @@ -186,13 +186,13 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: "output": channels.LastValue(int), "-1": channels.LastValue(int), } - chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1")} + chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")} for i in range(test_size - 2): chans[str(i)] = channels.LastValue(int) chains[str(i)] = ( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i)) + Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i)) ) - chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output") + chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output") app = Pregel(chains=chains, channels=chans, input="input", output="output") @@ -218,8 +218,8 @@ 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.send_to("output") - chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, @@ -239,8 +239,8 @@ 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.send_to("output") - chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, @@ -260,9 +260,9 @@ 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") + 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") app = Pregel( chains={ @@ -295,7 +295,7 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: inner_app = Pregel( chains={ - "one": Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") + "one": Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") }, channels={ "input": channels.LastValue(int), @@ -306,15 +306,15 @@ def test_invoke_join_then_call_other_app(mocker: MockerFixture) -> None: ) chain_one = ( - Pregel.subscribe_to("input") | add_10_each | Pregel.send_to("inbox_one").map() + Pregel.subscribe_to("input") | add_10_each | Pregel.write_to("inbox_one").map() ) chain_two = ( Pregel.subscribe_to("inbox_one") | inner_app.map() | sorted - | Pregel.send_to("outbox_one") + | Pregel.write_to("outbox_one") ) - chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.send_to("output") + chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.write_to("output") app = Pregel( chains={ @@ -345,9 +345,9 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: chain_one = ( Pregel.subscribe_to("input") | add_one - | Pregel.send_to(output=RunnablePassthrough(), between=RunnablePassthrough()) + | Pregel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough()) ) - chain_two = Pregel.subscribe_to("between") | add_one | Pregel.send_to("output") + chain_two = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output") app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, @@ -365,7 +365,7 @@ 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.send_to("between") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("between") chain_two = Pregel.subscribe_to("between") | add_one app = Pregel( @@ -387,7 +387,7 @@ 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.send_to("output") + chain_one = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output") chain_two = Pregel.subscribe_to("between") | add_one with pytest.raises(ValueError): @@ -416,8 +416,8 @@ 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.send_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") + chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 7eb8e381f..135192f06 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -11,7 +11,7 @@ from permchain import Pregel, channels 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") + chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") app = Pregel( chains={ @@ -30,7 +30,7 @@ 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.send_to("output") + chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") app = Pregel( chains={ @@ -55,7 +55,7 @@ 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.send_to("output") + chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") app = Pregel( chains={ @@ -84,8 +84,8 @@ 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.send_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") + chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, @@ -103,8 +103,8 @@ 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.send_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") + chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") pubsub = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, @@ -127,10 +127,10 @@ async def test_batch_two_processes_in_out() -> None: return inp + 1 chain_one = ( - Pregel.subscribe_to("input") | add_one_with_delay | Pregel.send_to("one") + Pregel.subscribe_to("input") | add_one_with_delay | Pregel.write_to("one") ) chain_two = ( - Pregel.subscribe_to("one") | add_one_with_delay | Pregel.send_to("output") + Pregel.subscribe_to("one") | add_one_with_delay | Pregel.write_to("output") ) app = Pregel( @@ -156,13 +156,13 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: "output": channels.LastValue(int), "-1": channels.LastValue(int), } - chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1")} + chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")} for i in range(test_size - 2): chans[str(i)] = channels.LastValue(int) chains[str(i)] = ( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i)) + Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i)) ) - chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output") + chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output") app = Pregel(chains=chains, channels=chans, input="input", output="output") @@ -185,13 +185,13 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: "output": channels.LastValue(int), "-1": channels.LastValue(int), } - chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.send_to("-1")} + chains = {"-1": Pregel.subscribe_to("input") | add_one | Pregel.write_to("-1")} for i in range(test_size - 2): chans[str(i)] = channels.LastValue(int) chains[str(i)] = ( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i)) + Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_to(str(i)) ) - chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.send_to("output") + chains["last"] = Pregel.subscribe_to(str(i)) | add_one | Pregel.write_to("output") app = Pregel(chains=chains, channels=chans, input="input", output="output") @@ -223,8 +223,8 @@ 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.send_to("output") - chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, @@ -244,8 +244,8 @@ 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.send_to("output") - chain_two = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") + chain_two = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, @@ -265,9 +265,9 @@ 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.send_to("inbox") - chain_three = Pregel.subscribe_to("input") | add_one | Pregel.send_to("inbox") - chain_four = Pregel.subscribe_to("inbox") | add_10_each | Pregel.send_to("output") + chain_one = Pregel.subscribe_to("input") | 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") app = Pregel( chains={ @@ -301,7 +301,7 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None inner_app = Pregel( chains={ - "one": Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") + "one": Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") }, channels={ "input": channels.LastValue(int), @@ -312,15 +312,15 @@ async def test_invoke_join_then_call_other_pubsub(mocker: MockerFixture) -> None ) chain_one = ( - Pregel.subscribe_to("input") | add_10_each | Pregel.send_to("inbox_one").map() + Pregel.subscribe_to("input") | add_10_each | Pregel.write_to("inbox_one").map() ) chain_two = ( Pregel.subscribe_to("inbox_one") | inner_app.map() | sorted - | Pregel.send_to("outbox_one") + | Pregel.write_to("outbox_one") ) - chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.send_to("output") + chain_three = Pregel.subscribe_to("outbox_one") | sum | Pregel.write_to("output") app = Pregel( chains={ @@ -353,9 +353,9 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non chain_one = ( Pregel.subscribe_to("input") | add_one - | Pregel.send_to(output=RunnablePassthrough(), between=RunnablePassthrough()) + | Pregel.write_to(output=RunnablePassthrough(), between=RunnablePassthrough()) ) - chain_two = Pregel.subscribe_to("between") | add_one | Pregel.send_to("output") + chain_two = Pregel.subscribe_to("between") | add_one | Pregel.write_to("output") app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, @@ -374,7 +374,7 @@ 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.send_to("between") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("between") chain_two = Pregel.subscribe_to("between") | add_one app = Pregel( @@ -417,8 +417,8 @@ 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.send_to("inbox") - chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.send_to("output") + chain_one = Pregel.subscribe_to("input") | add_one | Pregel.write_to("inbox") + chain_two = Pregel.subscribe_to_each("inbox") | add_one | Pregel.write_to("output") app = Pregel( chains={"chain_one": chain_one, "chain_two": chain_two}, From 948ab93563dacc706ffe54454eeb81760a661439 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 23 Oct 2023 10:30:28 +0100 Subject: [PATCH 4/4] Lint --- examples/draft-revise-loop.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/examples/draft-revise-loop.py b/examples/draft-revise-loop.py index 11b073185..0fd6a5697 100644 --- a/examples/draft-revise-loop.py +++ b/examples/draft-revise-loop.py @@ -63,12 +63,7 @@ gpt4 = ChatOpenAI(model="gpt-4") # chains -drafter_chain = ( - {"question": RunnablePassthrough(input_type=str)} - | drafter_prompt - | gpt3 - | StrOutputParser() -) +drafter_chain = drafter_prompt | gpt3 | StrOutputParser() editor_chain = ( editor_prompt