From f5c3c7ac7dd3db4014e594f0cdc6e5714ab68e3d Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 13 Mar 2024 11:17:49 -0700 Subject: [PATCH] Remove subscribe_to_each() method in pregel api - can be replaced by piping to batch method instead --- langgraph/pregel/__init__.py | 93 ++++++++++++++---------------------- langgraph/pregel/read.py | 60 +---------------------- langgraph/pregel/validate.py | 8 ++-- tests/test_pregel.py | 14 ++++-- tests/test_pregel_async.py | 14 ++++-- 5 files changed, 61 insertions(+), 128 deletions(-) diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 5826be8a2..5ad9c1922 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -65,7 +65,7 @@ from langgraph.constants import CONFIG_KEY_READ, CONFIG_KEY_SEND, INTERRUPT from langgraph.pregel.debug import print_checkpoint, print_step_start from langgraph.pregel.io import map_input, map_output from langgraph.pregel.log import logger -from langgraph.pregel.read import ChannelBatch, ChannelInvoke +from langgraph.pregel.read import ChannelInvoke from langgraph.pregel.reserved import AllReservedChannels, ReservedChannels from langgraph.pregel.validate import validate_graph, validate_keys from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry @@ -140,11 +140,6 @@ class Channel: tags=tags, ) - @classmethod - def subscribe_to_each(cls, inbox: str, key: Optional[str] = None) -> ChannelBatch: - """Runs process.batch() with the content of inbox each time it is updated.""" - return ChannelBatch(channel=inbox, key=key) - @classmethod def write_to( cls, @@ -173,7 +168,7 @@ class StateSnapshot(NamedTuple): class Pregel( RunnableSerializable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]] ): - nodes: Mapping[str, Union[ChannelInvoke, ChannelBatch]] + nodes: Mapping[str, ChannelInvoke] channels: Mapping[str, BaseChannel] = Field(default_factory=dict) @@ -1023,7 +1018,7 @@ def _apply_writes_from_view( def _prepare_next_tasks( checkpoint: Checkpoint, - processes: Mapping[str, Union[ChannelInvoke, ChannelBatch]], + processes: Mapping[str, ChannelInvoke], channels: Mapping[str, BaseChannel], update_seen: bool = True, ) -> tuple[Checkpoint, list[tuple[Runnable, Any, str]]]: @@ -1033,59 +1028,41 @@ def _prepare_next_tasks( # If so, prepare the values to be passed to them for name, proc in processes.items(): seen = checkpoint["versions_seen"][name] - if isinstance(proc, ChannelInvoke): - # If any of the channels read by this process were updated - if any( - checkpoint["channel_versions"][chan] > seen[chan] - for chan in proc.triggers - ): - # If all trigger channels subscribed by this process are not empty - # then invoke the process with the values of all non-empty channels - try: - val: Any = { - k: _read_channel( - channels, chan, catch=chan not in proc.triggers - ) - for k, chan in proc.channels.items() + # If any of the channels read by this process were updated + if any( + checkpoint["channel_versions"][chan] > seen[chan] for chan in proc.triggers + ): + # If all trigger channels subscribed by this process are not empty + # then invoke the process with the values of all non-empty channels + try: + val: Any = { + k: _read_channel(channels, chan, catch=chan not in proc.triggers) + for k, chan in proc.channels.items() + } + except EmptyChannelError: + continue + + # If the process has a mapper, apply it to the value + if proc.mapper is not None: + val = proc.mapper(val) + + # Processes that subscribe to a single keyless channel get + # the value directly, instead of a dict + if list(proc.channels.keys()) == [None]: + val = val[None] + + # update seen versions + if update_seen: + seen.update( + { + chan: checkpoint["channel_versions"][chan] + for chan in proc.triggers } - except EmptyChannelError: - continue - - # If the process has a mapper, apply it to the value - if proc.mapper is not None: - val = proc.mapper(val) - - # Processes that subscribe to a single keyless channel get - # the value directly, instead of a dict - if list(proc.channels.keys()) == [None]: - val = val[None] - - # update seen versions - if update_seen: - seen.update( - { - chan: checkpoint["channel_versions"][chan] - for chan in proc.triggers - } - ) - - # skip if condition is not met - if proc.when is None or proc.when(val): - tasks.append((proc, val, name)) - elif isinstance(proc, ChannelBatch): - # If the channel read by this process was updated - if checkpoint["channel_versions"][proc.channel] > seen[proc.channel]: - # If the channel subscribed by this process is not empty - try: - val = channels[proc.channel].get() - except EmptyChannelError: - continue - if proc.key is not None: - val = [{proc.key: v} for v in val] + ) + # skip if condition is not met + if proc.when is None or proc.when(val): tasks.append((proc, val, name)) - if update_seen: - seen[proc.channel] = checkpoint["channel_versions"][proc.channel] return checkpoint, tasks diff --git a/langgraph/pregel/read.py b/langgraph/pregel/read.py index 9006b639a..15a7dd6d6 100644 --- a/langgraph/pregel/read.py +++ b/langgraph/pregel/read.py @@ -10,12 +10,7 @@ from langchain_core.runnables import ( RunnablePassthrough, RunnableSerializable, ) -from langchain_core.runnables.base import ( - Other, - RunnableBindingBase, - RunnableEach, - coerce_to_runnable, -) +from langchain_core.runnables.base import Other, RunnableBindingBase, coerce_to_runnable from langchain_core.runnables.config import merge_configs from langchain_core.runnables.utils import ConfigurableFieldSpec @@ -170,56 +165,3 @@ class ChannelInvoke(RunnableBindingBase): ], ) -> RunnableSerializable: raise NotImplementedError() - - -class ChannelBatch(RunnableEach): - channel: str - - key: Optional[str] - - bound: Runnable[Any, Any] = Field(default=default_bound) - - def join(self, channels: Sequence[str]) -> ChannelBatch: - if self.key is None: - raise ValueError( - "Cannot join() additional channels without a key." - " Pass a key arg to Channel.subscribe_to_each()." - ) - - joiner = RunnablePassthrough.assign( - **{chan: ChannelRead(chan) for chan in channels} - ) - if self.bound is default_bound: - return ChannelBatch(channel=self.channel, key=self.key, bound=joiner) - else: - return ChannelBatch( - channel=self.channel, key=self.key, bound=self.bound | joiner - ) - - def __or__( # type: ignore[override] - self, - other: Union[ - Runnable[Any, Other], - Callable[[Any], Other], - Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]], - ], - ) -> ChannelBatch: - if self.bound is default_bound: - return ChannelBatch( - channel=self.channel, key=self.key, bound=coerce_to_runnable(other) - ) - else: - # delegate to __or__ in self.bound - return ChannelBatch( - channel=self.channel, key=self.key, bound=self.bound | other - ) - - def __ror__( - self, - other: Union[ - Runnable[Other, Any], - Callable[[Any], Other], - Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]], - ], - ) -> RunnableSerializable: - raise NotImplementedError() diff --git a/langgraph/pregel/validate.py b/langgraph/pregel/validate.py index 498d96aa2..302937689 100644 --- a/langgraph/pregel/validate.py +++ b/langgraph/pregel/validate.py @@ -3,12 +3,12 @@ from typing import Any, Mapping, Sequence, Union from langgraph.channels.base import BaseChannel from langgraph.channels.last_value import LastValue from langgraph.constants import INTERRUPT -from langgraph.pregel.read import ChannelBatch, ChannelInvoke +from langgraph.pregel.read import ChannelInvoke from langgraph.pregel.reserved import ReservedChannels def validate_graph( - nodes: Mapping[str, Union[ChannelInvoke, ChannelBatch]], + nodes: Mapping[str, ChannelInvoke], channels: dict[str, BaseChannel], input: Union[str, Sequence[str]], output: Union[str, Sequence[str]], @@ -22,11 +22,9 @@ def validate_graph( raise ValueError(f"Node name {INTERRUPT} is reserved") if isinstance(node, ChannelInvoke): subscribed_channels.update(node.channels.values()) - elif isinstance(node, ChannelBatch): - subscribed_channels.add(node.channel) else: raise TypeError( - f"Invalid node type {type(node)}, expected Channel.subscribe_to() or Channel.subscribe_to_each()" + f"Invalid node type {type(node)}, expected Channel.subscribe_to()" ) for chan in subscribed_channels: diff --git a/tests/test_pregel.py b/tests/test_pregel.py index c5d935285..a85fa080e 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -7,7 +7,7 @@ from contextlib import contextmanager from typing import Annotated, Generator, Optional, TypedDict, Union import pytest -from langchain_core.runnables import RunnablePassthrough +from langchain_core.runnables import RunnableLambda, RunnablePassthrough from pytest_mock import MockerFixture from syrupy import SnapshotAssertion @@ -304,7 +304,11 @@ def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> None: def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") + two = ( + Channel.subscribe_to("inbox") + | RunnableLambda(add_one).batch + | Channel.write_to("output").batch + ) app = Pregel( nodes={"one": one, "two": two}, @@ -653,7 +657,11 @@ def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") + two = ( + Channel.subscribe_to("inbox") + | RunnableLambda(add_one).batch + | Channel.write_to("output").batch + ) app = Pregel( nodes={"one": one, "two": two}, diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index f40c02338..eb06b9e74 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -14,7 +14,7 @@ from typing import ( ) import pytest -from langchain_core.runnables import RunnablePassthrough +from langchain_core.runnables import RunnableLambda, RunnablePassthrough from pytest_mock import MockerFixture from langgraph.channels.base import InvalidUpdateError @@ -311,7 +311,11 @@ async def test_invoke_two_processes_in_out_interrupt(mocker: MockerFixture) -> N async def test_invoke_two_processes_in_dict_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") + two = ( + Channel.subscribe_to("inbox") + | RunnableLambda(add_one).abatch + | Channel.write_to("output").abatch + ) pubsub = Pregel( nodes={"one": one, "two": two}, @@ -680,7 +684,11 @@ async def test_channel_enter_exit_timing(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to_each("inbox") | add_one | Channel.write_to("output") + two = ( + Channel.subscribe_to("inbox") + | RunnableLambda(add_one).abatch + | Channel.write_to("output").abatch + ) app = Pregel( nodes={"one": one, "two": two},