WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-30 19:59:40 +02:00
Code Issues Packages Projects Releases Wiki Activity

Remove subscribe_to_each() method in pregel api

Browse Source 
- can be replaced by piping to batch method instead
This commit is contained in:
Nuno Campos
2024-03-13 11:17:49 -07:00
parent 8b764f0a99
commit f5c3c7ac7d
5 changed files with 61 additions and 128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
langgraph/pregel/__init__.py
+35 -58
View File
@@ -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
langgraph/pregel/read.py
+1 -59
View File
@@ -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()
langgraph/pregel/validate.py
+3 -5
View File
@@ -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:
tests/test_pregel.py
+11 -3
View File
@@ -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},
tests/test_pregel_async.py
+11 -3
View File
@@ -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},