From 274600436b61f51bf86e9653bd3bc204e4453f93 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 23 Oct 2023 10:26:10 +0100 Subject: [PATCH] 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")