diff --git a/README.md b/README.md index cc1c3bd09..e74289dd7 100644 --- a/README.md +++ b/README.md @@ -12,11 +12,11 @@ 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( - 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 48f8573da..0fd6a5697 100644 --- a/examples/draft-revise-loop.py +++ b/examples/draft-revise-loop.py @@ -4,6 +4,7 @@ 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 @@ -64,28 +65,28 @@ gpt4 = ChatOpenAI(model="gpt-4") drafter_chain = drafter_prompt | gpt3 | StrOutputParser() -reviser_chain = reviser_prompt | gpt3 | StrOutputParser() - editor_chain = ( editor_prompt | gpt4.bind(functions=editor_functions) | 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") + | 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" @@ -95,23 +96,27 @@ editor = ( reviser = ( # subscribe to new values of "notes" channel, - # and join them with the current values of "question" and "draft" + # 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( - [drafter, reviser, editor], 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"], + 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 debug=True, ) diff --git a/examples/readme.py b/examples/readme.py index 67160e445..97998b3ce 100644 --- a/examples/readme.py +++ b/examples/readme.py @@ -3,11 +3,11 @@ 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( - 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..cbba99a7f 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] @@ -66,7 +83,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 @@ -75,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 @@ -92,31 +109,34 @@ def recursive_web_loader( ) if url not in x["visited"] and url != x["url"] ], - _max_steps=max_depth, ) ) 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.write_to("next_urls"), + # add the main chain + "visitor": visitor, + }, # define the channels channels={ "base_url": channels.LastValue(str), "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 35bd56d25..06784eb37 100644 --- a/permchain/pregel/__init__.py +++ b/permchain/pregel/__init__.py @@ -2,7 +2,6 @@ from __future__ import annotations import asyncio import concurrent.futures -import logging from collections import defaultdict, deque from typing import Any, AsyncIterator, Iterator, Mapping, Optional, Sequence, Type, cast @@ -11,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, @@ -23,7 +22,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,18 +29,19 @@ from permchain.channels.base import ( ChannelsManager, EmptyChannelError, ) -from permchain.pregel.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND, CONFIG_KEY_STEP +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__) - 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] @@ -50,37 +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, - *chains: Sequence[PregelInvoke | PregelBatch] | PregelInvoke | PregelBatch, - channels: Mapping[str, Channel], - output: str | Sequence[str], - input: str | Sequence[str], - 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)) - - validate_chains_channels(chains_flat, channels, input, output) - - super().__init__( - chains=chains_flat, - 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: @@ -136,10 +115,9 @@ 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, - _max_steps: Optional[int] = None, **kwargs: RunnableLike, ) -> PregelSink: """Writes to channels the result of the lambda, or None to skip writing.""" @@ -147,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( @@ -157,7 +134,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 @@ -165,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: @@ -188,16 +158,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]]() @@ -216,46 +177,29 @@ 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, }, ), ) - for proc, input in next_tasks + for proc, input, _ in next_tasks ), return_when=concurrent.futures.FIRST_EXCEPTION, timeout=self.step_timeout, ) - while done: - # if any task failed - if exc := done.pop().exception(): - # cancel all pending tasks - while inflight: - inflight.pop().cancel() - # raise the exception - raise exc - # TODO this is where retry of an entire step would happen - - if inflight: - # if we got here means we timed out - while inflight: - # cancel all pending tasks - inflight.pop().cancel() - # raise timeout error - raise TimeoutError(f"Timed out at step {step}") + # 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 is self.output for chan, _ in pending_writes): - yield channels[self.output].get() - else: - if updated := {c for c, _ in pending_writes if c in self.output}: - yield {chan: channels[chan].get() for chan in updated} + if 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 @@ -269,22 +213,13 @@ 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( 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: @@ -300,16 +235,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]]() @@ -330,47 +256,32 @@ 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, }, ), ) ) - for proc, input in next_tasks + for proc, input, _ in next_tasks ], return_when=asyncio.FIRST_EXCEPTION, timeout=self.step_timeout, ) - while done: - # if any task failed - if exc := done.pop().exception(): - # cancel all pending tasks - while inflight: - inflight.pop().cancel() - # raise the exception - raise exc - # TODO this is where retry of an entire step would happen - - if inflight: - # if we got here means we timed out - while inflight: - # cancel all pending tasks - inflight.pop().cancel() - # raise timeout error - raise TimeoutError(f"Timed out at step {step}") + # 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 is self.output for chan, _ in pending_writes): - yield channels[self.output].get() - else: - if updated := {c for c, _ in pending_writes if c in self.output}: - yield {chan: channels[chan].get() for chan in updated} + if 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: @@ -440,11 +351,35 @@ 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: 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: @@ -459,10 +394,10 @@ def _apply_writes_and_prepare_next_tasks( else: 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 +410,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 +422,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..6e5f9c37f 100644 --- a/permchain/pregel/constants.py +++ b/permchain/pregel/constants.py @@ -1,3 +1,2 @@ -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..636512bcc --- /dev/null +++ b/permchain/pregel/debug.py @@ -0,0 +1,34 @@ +from pprint import pformat +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.channels.base import Channel, EmptyChannelError + + +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" + ) + + "\n".join(f"- {name}({pformat(val)})" for _, val, name in next_tasks) + ) + + +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) + ) + + +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/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/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 f9ce13708..975a10748 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -8,15 +8,16 @@ from langchain.schema.runnable import RunnablePassthrough from pytest_mock import MockerFixture from permchain import Pregel, channels -from permchain.pregel import PregelInvoke def test_invoke_single_process_in_out(mocker: MockerFixture) -> 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( - (chain,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -32,10 +33,12 @@ 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( - (chain,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -55,10 +58,12 @@ 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( - (chain,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -82,11 +87,11 @@ 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( - [chain_one, chain_two], + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -101,11 +106,11 @@ 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( - [chain_one, chain_two], + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -124,15 +129,14 @@ 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( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -154,17 +158,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.write_to("-1")} for i in range(test_size - 2): chans[str(i)] = channels.LastValue(int) - chains.append( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i)) + chains[str(i)] = ( + Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_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.write_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 +186,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.write_to("-1")} for i in range(test_size - 2): chans[str(i)] = channels.LastValue(int) - chains.append( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i)) + chains[str(i)] = ( + Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_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.write_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}) == [ @@ -218,12 +218,11 @@ 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( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -240,12 +239,11 @@ 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( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.Inbox(int), @@ -262,14 +260,16 @@ 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( - 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 +294,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.write_to("output") + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -304,20 +306,22 @@ 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( - 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), @@ -341,13 +345,12 @@ 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( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -362,12 +365,11 @@ 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( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -385,13 +387,12 @@ 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): Pregel( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -415,11 +416,11 @@ 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( - [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..135192f06 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -7,15 +7,16 @@ from langchain.schema.runnable import RunnablePassthrough from pytest_mock import MockerFixture from permchain import Pregel, channels -from permchain.pregel import PregelInvoke async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Pregel.subscribe_to("input") | add_one | Pregel.send_to("output") + chain = Pregel.subscribe_to("input") | add_one | Pregel.write_to("output") app = Pregel( - (chain,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -29,10 +30,12 @@ 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( - (chain,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -52,10 +55,12 @@ 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( - (chain,), + chains={ + "one": chain, + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -79,11 +84,11 @@ 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( - [chain_one, chain_two], + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -98,11 +103,11 @@ 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( - [chain_one, chain_two], + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -122,15 +127,14 @@ 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( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -152,17 +156,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.write_to("-1")} for i in range(test_size - 2): chans[str(i)] = channels.LastValue(int) - chains.append( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i)) + chains[str(i)] = ( + Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_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.write_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 +185,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.write_to("-1")} for i in range(test_size - 2): chans[str(i)] = channels.LastValue(int) - chains.append( - Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.send_to(str(i)) + chains[str(i)] = ( + Pregel.subscribe_to(str(i - 1)) | add_one | Pregel.write_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.write_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): @@ -223,12 +223,11 @@ 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( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -245,12 +244,11 @@ 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( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.Inbox(int), @@ -267,14 +265,16 @@ 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( - 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 +300,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.write_to("output") + }, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -310,20 +312,22 @@ 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( - 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), @@ -349,13 +353,12 @@ 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( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -371,12 +374,11 @@ 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( - chain_one, - chain_two, + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int), @@ -415,11 +417,11 @@ 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( - [chain_one, chain_two], + chains={"chain_one": chain_one, "chain_two": chain_two}, channels={ "input": channels.LastValue(int), "output": channels.LastValue(int),