From 85522ec6a69d5476688debb2d072ab126752a521 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Sat, 24 May 2025 15:39:20 -0700 Subject: [PATCH] Pregel: Add NodeBuilder class to replace Channel.subscribe_to - For now keeping Channel.subscribe_to given it was used in docs --- docs/docs/concepts/pregel.md | 86 +++--- docs/docs/reference/pregel.md | 15 + libs/langgraph/langgraph/graph/graph.py | 6 +- libs/langgraph/langgraph/pregel/__init__.py | 261 ++++++++++++---- libs/langgraph/langgraph/pregel/read.py | 56 ++-- libs/langgraph/langgraph/pregel/validate.py | 2 +- libs/langgraph/langgraph/pregel/write.py | 18 +- libs/langgraph/langgraph/types.py | 2 +- .../tests/__snapshots__/test_large_cases.ambr | 18 +- .../tests/__snapshots__/test_pregel.ambr | 78 ++--- libs/langgraph/tests/test_large_cases.py | 9 +- .../langgraph/tests/test_large_cases_async.py | 6 +- libs/langgraph/tests/test_pregel.py | 286 +++++++----------- libs/langgraph/tests/test_pregel_async.py | 267 +++------------- 14 files changed, 494 insertions(+), 616 deletions(-) diff --git a/docs/docs/concepts/pregel.md b/docs/docs/concepts/pregel.md index 5157848db..d5a87d87a 100644 --- a/docs/docs/concepts/pregel.md +++ b/docs/docs/concepts/pregel.md @@ -25,7 +25,7 @@ Each step consists of three phases: Repeat until no **actors** are selected for execution, or a maximum number of steps is reached. -## Actors +## Actors An **actor** is a `PregelNode`. It subscribes to channels, reads data from them, and writes data to them. It can be thought of as an **actor** in the Pregel algorithm. `PregelNodes` implement LangChain's Runnable interface. @@ -39,7 +39,7 @@ Channels are used to communicate between actors (PregelNodes). Each channel has ## Examples -While most users will interact with Pregel through the [StateGraph][langgraph.graph.StateGraph] API or +While most users will interact with Pregel through the [StateGraph][langgraph.graph.StateGraph] API or the [entrypoint][langgraph.func.entrypoint] decorator, it is possible to interact with Pregel directly. Below are a few different examples to give you a sense of the Pregel API. @@ -49,12 +49,12 @@ Below are a few different examples to give you a sense of the Pregel API. ```python from langgraph.channels import EphemeralValue - from langgraph.pregel import Pregel, Channel + from langgraph.pregel import Pregel, NodeBuilder node1 = ( - Channel.subscribe_to("a") - | (lambda x: x + x) - | Channel.write_to("b") + NodeBuilder().subscribe_only("a") + .do(lambda x: x + x) + .write_to("b") ) app = Pregel( @@ -78,18 +78,18 @@ Below are a few different examples to give you a sense of the Pregel API. ```python from langgraph.channels import LastValue, EphemeralValue - from langgraph.pregel import Pregel, Channel + from langgraph.pregel import Pregel, NodeBuilder node1 = ( - Channel.subscribe_to("a") - | (lambda x: x + x) - | Channel.write_to("b") + NodeBuilder().subscribe_only("a") + .do(lambda x: x + x) + .write_to("b") ) node2 = ( - Channel.subscribe_to("b") - | (lambda x: x + x) - | Channel.write_to("c") + NodeBuilder().subscribe_only("b") + .do(lambda x: x + x) + .write_to("c") ) @@ -115,23 +115,18 @@ Below are a few different examples to give you a sense of the Pregel API. ```python from langgraph.channels import EphemeralValue, Topic - from langgraph.pregel import Pregel, Channel + from langgraph.pregel import Pregel, NodeBuilder node1 = ( - Channel.subscribe_to("a") - | (lambda x: x + x) - | { - "b": Channel.write_to("b"), - "c": Channel.write_to("c") - } + NodeBuilder().subscribe_only("a") + .do(lambda x: x + x) + .write_to("b", "c") ) node2 = ( - Channel.subscribe_to("b") - | (lambda x: x + x) - | { - "c": Channel.write_to("c"), - } + NodeBuilder().subscribe_to("b") + .do(lambda x: x["b"] + x["b"]) + .write_to("c") ) app = Pregel( @@ -158,24 +153,19 @@ Below are a few different examples to give you a sense of the Pregel API. ```python from langgraph.channels import EphemeralValue, BinaryOperatorAggregate - from langgraph.pregel import Pregel, Channel + from langgraph.pregel import Pregel, NodeBuilder node1 = ( - Channel.subscribe_to("a") - | (lambda x: x + x) - | { - "b": Channel.write_to("b"), - "c": Channel.write_to("c") - } + NodeBuilder().subscribe_only("a") + .do(lambda x: x + x) + .write_to("b", "c") ) node2 = ( - Channel.subscribe_to("b") - | (lambda x: x + x) - | { - "c": Channel.write_to("c"), - } + NodeBuilder().subscribe_only("b") + .do(lambda x: x + x) + .write_to("c") ) def reducer(current, update): @@ -197,8 +187,7 @@ Below are a few different examples to give you a sense of the Pregel API. app.invoke({"a": "foo"}) ``` - - + === "Cycle" This example demonstrates how to introduce a cycle in the graph, by having @@ -207,12 +196,12 @@ Below are a few different examples to give you a sense of the Pregel API. ```python from langgraph.channels import EphemeralValue - from langgraph.pregel import Pregel, Channel, ChannelWrite, ChannelWriteEntry + from langgraph.pregel import Pregel, NodeBuilder, ChannelWriteEntry example_node = ( - Channel.subscribe_to("value") - | (lambda x: x + x if len(x) < 10 else None) - | ChannelWrite(writes=[ChannelWriteEntry(channel="value", skip_none=True)]) + NodeBuilder().subscribe_only("value") + .do(lambda x: x + x if len(x) < 10 else None) + .write_to(ChannelWriteEntry("value", skip_none=True)) ) app = Pregel( @@ -235,7 +224,6 @@ Below are a few different examples to give you a sense of the Pregel API. LangGraph provides two high-level APIs for creating a Pregel application: the [StateGraph (Graph API)](./low_level.md) and the [Functional API](functional_api.md). - === "StateGraph (Graph API)" The [StateGraph (Graph API)][langgraph.graph.StateGraph] is a higher-level abstraction that simplifies the creation of Pregel applications. It allows you to define a graph of nodes and edges. When you compile the graph, the StateGraph API automatically creates the Pregel application for you. @@ -266,7 +254,7 @@ LangGraph provides two high-level APIs for creating a Pregel application: the [S builder.add_node(score_essay) builder.add_edge(START, "write_essay") - # Compile the graph. + # Compile the graph. # This will return a Pregel instance. graph = builder.compile() ``` @@ -279,7 +267,7 @@ LangGraph provides two high-level APIs for creating a Pregel application: the [S You will see something like this: - ```pycon + ```pycon {'__start__': , 'write_essay': , 'score_essay': } @@ -310,7 +298,7 @@ LangGraph provides two high-level APIs for creating a Pregel application: the [S === "Functional API" In the [Functional API](functional_api.md), you can use an [`entrypoint`][langgraph.func.entrypoint] to create - a Pregel application. The `entrypoint` decorator allows you to define a function that takes input and returns output. + a Pregel application. The `entrypoint` decorator allows you to define a function that takes input and returns output. ```python from typing import TypedDict, Optional @@ -339,8 +327,8 @@ LangGraph provides two high-level APIs for creating a Pregel application: the [S ``` ```pycon - Nodes: + Nodes: {'write_essay': } - Channels: + Channels: {'__start__': , '__end__': , '__previous__': } ``` diff --git a/docs/docs/reference/pregel.md b/docs/docs/reference/pregel.md index a0ada895a..a115b1fde 100644 --- a/docs/docs/reference/pregel.md +++ b/docs/docs/reference/pregel.md @@ -1,5 +1,20 @@ # Pregel +::: langgraph.pregel.NodeBuilder + options: + show_if_no_docstring: true + show_root_heading: true + show_root_full_path: false + members: + - subscribe_single + - subscribe_to + - do + - write_to + - meta + - retry + - cache + - build + ::: langgraph.pregel.Pregel options: show_if_no_docstring: true diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index 6c4c069b7..bb96c95f6 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -26,7 +26,7 @@ from langgraph.constants import ( Send, ) from langgraph.graph.branch import Branch -from langgraph.pregel import Channel, Pregel +from langgraph.pregel import NodeBuilder, Pregel from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore @@ -428,7 +428,9 @@ class CompiledGraph(Pregel): # add hidden start node if start == START and start not in self.nodes: - self.nodes[start] = Channel.subscribe_to(START, tags=[TAG_HIDDEN]) + self.nodes[start] = ( + NodeBuilder().subscribe_only(START).meta(TAG_HIDDEN).build() + ) # attach branch writer self.nodes[start] |= branch.run(get_writes) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 4550c56d0..746836e01 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -8,14 +8,7 @@ import weakref from collections import defaultdict, deque from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from functools import partial -from typing import ( - Any, - Callable, - Union, - cast, - get_type_hints, - overload, -) +from typing import Any, Callable, Union, cast, get_type_hints, overload from uuid import UUID, uuid5 from langchain_core.globals import get_debug @@ -99,7 +92,7 @@ from langgraph.pregel.io import map_input, read_channels from langgraph.pregel.loop import AsyncPregelLoop, StreamProtocol, SyncPregelLoop from langgraph.pregel.messages import StreamMessagesHandler from langgraph.pregel.protocol import PregelProtocol -from langgraph.pregel.read import PregelNode +from langgraph.pregel.read import DEFAULT_BOUND, PregelNode from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.runner import PregelRunner from langgraph.pregel.utils import get_new_channel_versions @@ -127,6 +120,12 @@ from langgraph.utils.config import ( from langgraph.utils.fields import get_enhanced_type_hints from langgraph.utils.pydantic import create_model, is_supported_by_pydantic from langgraph.utils.queue import AsyncQueue, SyncQueue # type: ignore[attr-defined] +from langgraph.utils.runnable import ( + Runnable, + RunnableLike, + RunnableSeq, + coerce_to_runnable, +) try: from langchain_core.tracers._streaming import _StreamingCallbackHandler @@ -136,6 +135,172 @@ except ImportError: WriteValue = Union[Callable[[Input], Output], Any] +class NodeBuilder: + __slots__ = ( + "_channels", + "_triggers", + "_tags", + "_metadata", + "_writes", + "_bound", + "_retries", + "_cache", + ) + + _channels: list[str] | dict[str, str] + _triggers: list[str] + _tags: list[str] + _metadata: dict[str, Any] + _writes: list[ChannelWriteEntry] + _bound: Runnable + _retries: list[RetryPolicy] + _cache: CachePolicy | None + + def __init__( + self, + ) -> None: + self._channels = {} + self._triggers = [] + self._tags = [] + self._metadata = {} + self._writes = [] + self._bound = DEFAULT_BOUND + self._retries = [] + self._cache = None + + def subscribe_only( + self, + channel: str, + ) -> Self: + """Subscribe to a single channel.""" + if isinstance(self._channels, list): + self._channels.append(channel) + elif not self._channels: + self._channels = [channel] + else: + raise ValueError( + "Cannot subscribe to single channels when other channels are already subscribed to" + ) + + self._triggers.append(channel) + + return self + + def subscribe_to( + self, + *channels: str, + read: bool = True, + ) -> Self: + """Add channels to subscribe to. Node will be invoked when any of these + channels are updated, with a dict of the channel values as input. + + Args: + channels: Channel name(s) to subscribe to + read: If True, the channels will be included in the input to the node. + Otherwise, they will trigger the node without being sent in input. + + Returns: + Self for chaining + """ + if isinstance(self._channels, list): + raise ValueError( + "Cannot subscribe to channels when subscribed to a single channel" + ) + if read: + if not self._channels: + self._channels = {chan: chan for chan in channels} + else: + self._channels.update({chan: chan for chan in channels}) + + if isinstance(channels, str): + self._triggers.append(channels) + else: + self._triggers.extend(channels) + + return self + + def read_from( + self, + *channels: str, + ) -> Self: + """Adds the specified channels to read from, without subscribing to them.""" + assert self._channels, "Channels must be specified first" + assert isinstance(self._channels, dict), ( + "Cannot read additional channels when subscribed to single channels" + ) + self._channels.update({c: c for c in channels}) + return self + + def do( + self, + node: RunnableLike, + ) -> Self: + """Adds the specified node.""" + if self._bound is not DEFAULT_BOUND: + self._bound = RunnableSeq( + self._bound, coerce_to_runnable(node, name=None, trace=True) + ) + else: + self._bound = coerce_to_runnable(node, name=None, trace=True) + return self + + def write_to( + self, + *channels: str | ChannelWriteEntry, + **kwargs: WriteValue, + ) -> Self: + """Add channel writes. + + Args: + *channels: Channel names to write to + **kwargs: Channel name and value mappings + + Returns: + Self for chaining + """ + self._writes.extend( + ChannelWriteEntry(c) if isinstance(c, str) else c for c in channels + ) + self._writes.extend( + ChannelWriteEntry(k, mapper=v) + if callable(v) + else ChannelWriteEntry(k, value=v) + for k, v in kwargs.items() + ) + + return self + + def meta(self, *tags: str, **metadata: Any) -> Self: + """Add tags or metadata to the node.""" + self._tags.extend(tags) + self._metadata.update(metadata) + return self + + def retry(self, *policies: RetryPolicy) -> Self: + """Adds retry policies to the node.""" + self._retries.extend(policies) + return self + + def cache(self, policy: CachePolicy) -> Self: + """Adds cache policies to the node.""" + self._cache = policy + return self + + def build(self) -> PregelNode: + """Builds the node.""" + return PregelNode( + channels=self._channels, + triggers=self._triggers, + tags=self._tags, + metadata=self._metadata, + writers=[ChannelWrite(self._writes)], + bound=self._bound, + retry_policy=self._retries, + cache_policy=self._cache, + ) + + +# Deprecated, remove in 2.0 class Channel: @overload @classmethod @@ -191,12 +356,12 @@ class Channel: @classmethod def write_to( cls, - *channels: str, + *channels: str | ChannelWriteEntry, **kwargs: WriteValue, ) -> ChannelWrite: """Writes to channels the result of the lambda, or None to skip writing.""" return ChannelWrite( - [ChannelWriteEntry(c) for c in channels] + [ChannelWriteEntry(c) if isinstance(c, str) else c for c in channels] + [ ( ChannelWriteEntry(k, mapper=v) @@ -285,12 +450,12 @@ class Pregel(PregelProtocol): Example: Single node application ```python from langgraph.channels import EphemeralValue - from langgraph.pregel import Pregel, Channel, ChannelWriteEntry + from langgraph.pregel import Pregel, NodeBuilder node1 = ( - Channel.subscribe_to("a") - | (lambda x: x + x) - | Channel.write_to("b") + NodeBuilder().subscribe_only("a") + .do(lambda x: x + x) + .write_to("b") ) app = Pregel( @@ -313,18 +478,18 @@ class Pregel(PregelProtocol): Example: Using multiple nodes and multiple output channels ```python from langgraph.channels import LastValue, EphemeralValue - from langgraph.pregel import Pregel, Channel, ChannelWriteEntry + from langgraph.pregel import Pregel, NodeBuilder node1 = ( - Channel.subscribe_to("a") - | (lambda x: x + x) - | Channel.write_to("b") + NodeBuilder().subscribe_only("a") + .do(lambda x: x + x) + .write_to("b") ) node2 = ( - Channel.subscribe_to("b") - | (lambda x: x + x) - | Channel.write_to("c") + NodeBuilder().subscribe_to("b") + .do(lambda x: x["b"] + x["b"]) + .write_to("c") ) @@ -349,23 +514,18 @@ class Pregel(PregelProtocol): Example: Using a Topic channel ```python from langgraph.channels import LastValue, EphemeralValue, Topic - from langgraph.pregel import Pregel, Channel, ChannelWriteEntry + from langgraph.pregel import Pregel, NodeBuilder node1 = ( - Channel.subscribe_to("a") - | (lambda x: x + x) - | { - "b": Channel.write_to("b"), - "c": Channel.write_to("c") - } + NodeBuilder().subscribe_only("a") + .do(lambda x: x + x) + .write_to("b", "c") ) node2 = ( - Channel.subscribe_to("b") - | (lambda x: x + x) - | { - "c": Channel.write_to("c"), - } + NodeBuilder().subscribe_only("b") + .do(lambda x: x + x) + .write_to("c") ) @@ -390,24 +550,19 @@ class Pregel(PregelProtocol): Example: Using a BinaryOperatorAggregate channel ```python from langgraph.channels import EphemeralValue, BinaryOperatorAggregate - from langgraph.pregel import Pregel, Channel + from langgraph.pregel import Pregel, NodeBuilder node1 = ( - Channel.subscribe_to("a") - | (lambda x: x + x) - | { - "b": Channel.write_to("b"), - "c": Channel.write_to("c") - } + NodeBuilder().subscribe_only("a") + .do(lambda x: x + x) + .write_to("b", "c") ) node2 = ( - Channel.subscribe_to("b") - | (lambda x: x + x) - | { - "c": Channel.write_to("c"), - } + NodeBuilder().subscribe_only("b") + .do(lambda x: x + x) + .write_to("c") ) @@ -442,12 +597,12 @@ class Pregel(PregelProtocol): ```python from langgraph.channels import EphemeralValue - from langgraph.pregel import Pregel, Channel, ChannelWrite, ChannelWriteEntry + from langgraph.pregel import Pregel, NodeBuilder, ChannelWriteEntry example_node = ( - Channel.subscribe_to("value") - | (lambda x: x + x if len(x) < 10 else None) - | ChannelWrite(writes=[ChannelWriteEntry(channel="value", skip_none=True)]) + NodeBuilder().subscribe_only("value") + .do(lambda x: x + x if len(x) < 10 else None) + .write_to(ChannelWriteEntry(channel="value", skip_none=True)) ) app = Pregel( @@ -524,7 +679,7 @@ class Pregel(PregelProtocol): def __init__( self, *, - nodes: dict[str, PregelNode], + nodes: dict[str, PregelNode | NodeBuilder], channels: dict[str, BaseChannel | ManagedValueSpec] | None, auto_validate: bool = True, stream_mode: StreamMode = "values", @@ -547,7 +702,9 @@ class Pregel(PregelProtocol): trigger_to_nodes: Mapping[str, Sequence[str]] | None = None, name: str = "LangGraph", ) -> None: - self.nodes = nodes + self.nodes = { + k: v.build() if isinstance(v, NodeBuilder) else v for k, v in nodes.items() + } self.channels = channels or {} self.stream_mode = stream_mode self.stream_eager = stream_eager diff --git a/libs/langgraph/langgraph/pregel/read.py b/libs/langgraph/langgraph/pregel/read.py index 78bd108ce..91228e992 100644 --- a/libs/langgraph/langgraph/pregel/read.py +++ b/libs/langgraph/langgraph/pregel/read.py @@ -8,14 +8,7 @@ from typing import ( Union, ) -from langchain_core.runnables import ( - Runnable, - RunnableConfig, - RunnablePassthrough, - RunnableSerializable, -) -from langchain_core.runnables.base import Other, coerce_to_runnable -from langchain_core.runnables.utils import ConfigurableFieldSpec, Input +from langchain_core.runnables import Runnable, RunnableConfig from langgraph.constants import CONF, CONFIG_KEY_READ from langgraph.pregel.protocol import PregelProtocol @@ -24,7 +17,7 @@ from langgraph.pregel.utils import find_subgraph_pregel from langgraph.pregel.write import ChannelWrite from langgraph.types import CachePolicy from langgraph.utils.config import merge_configs -from langgraph.utils.runnable import RunnableCallable, RunnableSeq +from langgraph.utils.runnable import RunnableCallable, RunnableSeq, coerce_to_runnable READ_TYPE = Callable[[Union[str, Sequence[str]], bool], Union[Any, dict[str, Any]]] INPUT_CACHE_KEY_TYPE = tuple[Callable[..., Any], tuple[str, ...]] @@ -40,18 +33,6 @@ class ChannelRead(RunnableCallable): mapper: Callable[[Any], Any] | None = None - @property - def config_specs(self) -> list[ConfigurableFieldSpec]: - return [ - ConfigurableFieldSpec( - id=CONFIG_KEY_READ, - name=CONFIG_KEY_READ, - description=None, - default=None, - annotation=None, - ), - ] - def __init__( self, channel: str | list[str], @@ -112,7 +93,7 @@ class ChannelRead(RunnableCallable): return read(select, fresh) -DEFAULT_BOUND: RunnablePassthrough = RunnablePassthrough() +DEFAULT_BOUND = RunnableCallable(lambda input: input) class PregelNode(Runnable): @@ -217,7 +198,6 @@ class PregelNode(Runnable): # careful to not modify the original writers list or ChannelWrite writers[-2] = ChannelWrite( writes=writers[-2].writes + writers[-1].writes, - tags=writers[-2].tags, ) writers.pop() return writers @@ -266,37 +246,39 @@ class PregelNode(Runnable): def __or__( self, - other: Runnable[Any, Other] - | Callable[[Any], Other] - | Mapping[str, Runnable[Any, Other] | Callable[[Any], Other]], + other: Runnable[Any, Any] + | Callable[[Any], Any] + | Mapping[str, Runnable[Any, Any] | Callable[[Any], Any]], ) -> PregelNode: if isinstance(other, Runnable) and ChannelWrite.is_writer(other): return self.copy(update=dict(writers=[*self.writers, other])) elif self.bound is DEFAULT_BOUND: - return self.copy(update=dict(bound=coerce_to_runnable(other))) + return self.copy( + update=dict(bound=coerce_to_runnable(other, name=None, trace=True)) + ) else: return self.copy(update=dict(bound=RunnableSeq(self.bound, other))) def pipe( self, - *others: Runnable[Any, Other] | Callable[[Any], Other], + *others: Runnable[Any, Any] | Callable[[Any], Any], name: str | None = None, - ) -> RunnableSerializable[Any, Other]: + ) -> PregelNode: for other in others: self = self | other return self def __ror__( self, - other: Runnable[Other, Any] - | Callable[[Any], Other] - | Mapping[str, Runnable[Other, Any] | Callable[[Other], Any]], - ) -> RunnableSerializable: + other: Runnable[Any, Any] + | Callable[[Any], Any] + | Mapping[str, Runnable[Any, Any] | Callable[[Any], Any]], + ) -> PregelNode: raise NotImplementedError() def invoke( self, - input: Input, + input: Any, config: RunnableConfig | None = None, **kwargs: Any | None, ) -> Any: @@ -308,7 +290,7 @@ class PregelNode(Runnable): async def ainvoke( self, - input: Input, + input: Any, config: RunnableConfig | None = None, **kwargs: Any | None, ) -> Any: @@ -320,7 +302,7 @@ class PregelNode(Runnable): def stream( self, - input: Input, + input: Any, config: RunnableConfig | None = None, **kwargs: Any | None, ) -> Iterator[Any]: @@ -332,7 +314,7 @@ class PregelNode(Runnable): async def astream( self, - input: Input, + input: Any, config: RunnableConfig | None = None, **kwargs: Any | None, ) -> AsyncIterator[Any]: diff --git a/libs/langgraph/langgraph/pregel/validate.py b/libs/langgraph/langgraph/pregel/validate.py index 7f638903f..169268169 100644 --- a/libs/langgraph/langgraph/pregel/validate.py +++ b/libs/langgraph/langgraph/pregel/validate.py @@ -28,7 +28,7 @@ def validate_graph( subscribed_channels.update(node.triggers) else: raise TypeError( - f"Invalid node type {type(node)}, expected Channel.subscribe_to()" + f"Invalid node type {type(node)}, expected PregelNode or NodeBuilder" ) for chan in subscribed_channels: diff --git a/libs/langgraph/langgraph/pregel/write.py b/libs/langgraph/langgraph/pregel/write.py index 98b6dd587..ee386ad51 100644 --- a/libs/langgraph/langgraph/pregel/write.py +++ b/libs/langgraph/langgraph/pregel/write.py @@ -12,7 +12,6 @@ from typing import ( ) from langchain_core.runnables import Runnable, RunnableConfig -from langchain_core.runnables.utils import ConfigurableFieldSpec from langgraph.constants import CONF, CONFIG_KEY_SEND, MISSING, TASKS, Send from langgraph.errors import InvalidUpdateError @@ -56,13 +55,13 @@ class ChannelWrite(RunnableCallable): self, writes: Sequence[ChannelWriteEntry | ChannelWriteTupleEntry | Send], *, - tags: Sequence[str] | None = None, # ignored - require_at_least_one_of: Sequence[str] | None = None, # ignored + tags: Sequence[str] | None = None, ): super().__init__( func=self._write, afunc=self._awrite, name=None, + tags=tags, trace=False, func_accepts_config=True, ) @@ -75,18 +74,6 @@ class ChannelWrite(RunnableCallable): name = f"ChannelWrite<{','.join(w.channel if isinstance(w, ChannelWriteEntry) else '...' if isinstance(w, ChannelWriteTupleEntry) else w.node for w in self.writes)}>" return super().get_name(suffix, name=name) - @property - def config_specs(self) -> list[ConfigurableFieldSpec]: - return [ - ConfigurableFieldSpec( - id=CONFIG_KEY_SEND, - name=CONFIG_KEY_SEND, - description=None, - default=None, - annotation=None, - ), - ] - def _write(self, input: Any, config: RunnableConfig) -> None: writes = [ ChannelWriteEntry(write.channel, input, write.skip_none, write.mapper) @@ -122,7 +109,6 @@ class ChannelWrite(RunnableCallable): config: RunnableConfig, writes: Sequence[ChannelWriteEntry | ChannelWriteTupleEntry | Send], allow_passthrough: bool = True, - require_at_least_one_of: Sequence[str] | None = None, # ignored ) -> None: # validate for w in writes: diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index cc92e032c..0d02bac38 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -488,6 +488,7 @@ def interrupt(value: Any) -> Any: Raises: GraphInterrupt: On the first invocation within the node, halts execution and surfaces the provided value to the client. """ + from langgraph.config import get_config from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_SCRATCHPAD, @@ -496,7 +497,6 @@ def interrupt(value: Any) -> Any: RESUME, ) from langgraph.errors import GraphInterrupt - from langgraph.utils.config import get_config conf = get_config()["configurable"] # track interrupt index diff --git a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr index 34c117a1b..c200dae98 100644 --- a/libs/langgraph/tests/__snapshots__/test_large_cases.ambr +++ b/libs/langgraph/tests/__snapshots__/test_large_cases.ambr @@ -155,10 +155,10 @@ "type": "runnable", "data": { "id": [ - "langchain", - "schema", + "langgraph", + "utils", "runnable", - "RunnablePassthrough" + "RunnableCallable" ], "name": "__start__" } @@ -243,10 +243,10 @@ "type": "runnable", "data": { "id": [ - "langchain", - "schema", + "langgraph", + "utils", "runnable", - "RunnablePassthrough" + "RunnableCallable" ], "name": "__start__" } @@ -330,10 +330,10 @@ "type": "runnable", "data": { "id": [ - "langchain", - "schema", + "langgraph", + "utils", "runnable", - "RunnablePassthrough" + "RunnableCallable" ], "name": "__start__" } diff --git a/libs/langgraph/tests/__snapshots__/test_pregel.ambr b/libs/langgraph/tests/__snapshots__/test_pregel.ambr index 713a9d351..4b991b513 100644 --- a/libs/langgraph/tests/__snapshots__/test_pregel.ambr +++ b/libs/langgraph/tests/__snapshots__/test_pregel.ambr @@ -40,10 +40,10 @@ "type": "runnable", "data": { "id": [ - "langchain", - "schema", + "langgraph", + "utils", "runnable", - "RunnablePassthrough" + "RunnableCallable" ], "name": "__start__" } @@ -103,10 +103,10 @@ "type": "runnable", "data": { "id": [ - "langchain", - "schema", + "langgraph", + "utils", "runnable", - "RunnablePassthrough" + "RunnableCallable" ], "name": "__start__" } @@ -192,10 +192,10 @@ "type": "runnable", "data": { "id": [ - "langchain", - "schema", + "langgraph", + "utils", "runnable", - "RunnablePassthrough" + "RunnableCallable" ], "name": "__start__" } @@ -248,10 +248,10 @@ "type": "runnable", "data": { "id": [ - "langchain", - "schema", + "langgraph", + "utils", "runnable", - "RunnablePassthrough" + "RunnableCallable" ], "name": "__start__" } @@ -326,10 +326,10 @@ "type": "runnable", "data": { "id": [ - "langchain", - "schema", + "langgraph", + "utils", "runnable", - "RunnablePassthrough" + "RunnableCallable" ], "name": "__start__" } @@ -405,10 +405,10 @@ "type": "runnable", "data": { "id": [ - "langchain", - "schema", + "langgraph", + "utils", "runnable", - "RunnablePassthrough" + "RunnableCallable" ], "name": "__start__" } @@ -460,10 +460,10 @@ "type": "runnable", "data": { "id": [ - "langchain", - "schema", + "langgraph", + "utils", "runnable", - "RunnablePassthrough" + "RunnableCallable" ], "name": "__start__" } @@ -781,10 +781,10 @@ dict({ 'data': dict({ 'id': list([ - 'langchain', - 'schema', + 'langgraph', + 'utils', 'runnable', - 'RunnablePassthrough', + 'RunnableCallable', ]), 'name': '__start__', }), @@ -823,10 +823,10 @@ dict({ 'data': dict({ 'id': list([ - 'langchain', - 'schema', + 'langgraph', + 'utils', 'runnable', - 'RunnablePassthrough', + 'RunnableCallable', ]), 'name': 'tool_two:__start__', }), @@ -1034,10 +1034,10 @@ dict({ 'data': dict({ 'id': list([ - 'langchain', - 'schema', + 'langgraph', + 'utils', 'runnable', - 'RunnablePassthrough', + 'RunnableCallable', ]), 'name': '__start__', }), @@ -1101,10 +1101,10 @@ dict({ 'data': dict({ 'id': list([ - 'langchain', - 'schema', + 'langgraph', + 'utils', 'runnable', - 'RunnablePassthrough', + 'RunnableCallable', ]), 'name': '__start__', }), @@ -1199,10 +1199,10 @@ dict({ 'data': dict({ 'id': list([ - 'langchain', - 'schema', + 'langgraph', + 'utils', 'runnable', - 'RunnablePassthrough', + 'RunnableCallable', ]), 'name': '__start__', }), @@ -1241,10 +1241,10 @@ dict({ 'data': dict({ 'id': list([ - 'langchain', - 'schema', + 'langgraph', + 'utils', 'runnable', - 'RunnablePassthrough', + 'RunnableCallable', ]), 'name': 'conduct_interview:__start__', }), diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index b570f086a..0babd6455 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -22,7 +22,7 @@ from langgraph.graph.graph import Graph from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.prebuilt.chat_agent_executor import create_react_agent from langgraph.prebuilt.tool_node import ToolNode -from langgraph.pregel import Channel, Pregel +from langgraph.pregel import NodeBuilder, Pregel from langgraph.store.base import BaseStore from langgraph.store.memory import InMemoryStore from langgraph.types import ( @@ -49,11 +49,12 @@ from tests.messages import ( def test_invoke_two_processes_in_out_interrupt( - sync_checkpointer: BaseCheckpointSaver, mocker: MockerFixture + sync_checkpointer: BaseCheckpointSaver, + 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("inbox") | add_one | Channel.write_to("output") + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output") app = Pregel( nodes={"one": one, "two": two}, diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index f0ae78813..1de6a2442 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -25,7 +25,7 @@ from langgraph.graph.message import MessageGraph, add_messages from langgraph.graph.state import StateGraph from langgraph.prebuilt.chat_agent_executor import create_react_agent from langgraph.prebuilt.tool_node import ToolNode -from langgraph.pregel import Channel, Pregel +from langgraph.pregel import NodeBuilder, Pregel from langgraph.store.base import BaseStore from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter from tests.any_int import AnyInt @@ -46,8 +46,8 @@ async def test_invoke_two_processes_in_out_interrupt( async_checkpointer: BaseCheckpointSaver, 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("inbox") | add_one | Channel.write_to("output") + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output") app = Pregel( nodes={"one": one, "two": two}, channels={ diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index fe85ba99f..a652ec66f 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -49,8 +49,15 @@ from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessageGraph, MessagesState, add_messages from langgraph.prebuilt.tool_node import ToolNode -from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot +from langgraph.pregel import ( + Channel, + GraphRecursionError, + NodeBuilder, + Pregel, + StateSnapshot, +) from langgraph.pregel.loop import SyncPregelLoop +from langgraph.pregel.read import PregelNode from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.runner import PregelRunner from langgraph.store.base import BaseStore @@ -260,9 +267,15 @@ def test_checkpoint_errors() -> None: graph.invoke("", {"configurable": {"thread_id": "thread-1"}}) -def test_config_json_schema() -> None: +@pytest.mark.parametrize("use_node_builder", [True, False]) +def test_config_json_schema(use_node_builder: bool) -> None: """Test that config json schema is generated properly.""" - chain = Channel.subscribe_to("input") | Channel.write_to("output") + if use_node_builder: + chain: Union[NodeBuilder, PregelNode] = ( + NodeBuilder().subscribe_only("input").write_to("output") + ) + else: + chain = Channel.subscribe_to("input") | Channel.write_to("output") @dataclass class Foo: @@ -469,9 +482,15 @@ def test_reducer_before_first_node() -> None: } -def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: +@pytest.mark.parametrize("use_node_builder", [True, False]) +def test_invoke_single_process_in_out( + mocker: MockerFixture, use_node_builder: bool +) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + if use_node_builder: + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + else: + chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( nodes={ @@ -526,13 +545,24 @@ def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> None: assert gapp.invoke(1) == falsy_value -def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: +@pytest.mark.parametrize("use_node_builder", [True, False]) +def test_invoke_single_process_in_write_kwargs( + mocker: MockerFixture, use_node_builder: bool +) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = ( - Channel.subscribe_to("input") - | add_one - | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) - ) + if use_node_builder: + chain = ( + NodeBuilder() + .subscribe_only("input") + .do(add_one) + .write_to("output", fixed=5, output_plus_one=lambda x: x + 1) + ) + else: + chain = ( + Channel.subscribe_to("input") + | add_one + | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) + ) app = Pregel( nodes={"one": chain}, @@ -566,9 +596,15 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: assert app.invoke(2) == {"output": 3, "fixed": 5, "output_plus_one": 4} -def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: +@pytest.mark.parametrize("use_node_builder", [True, False]) +def test_invoke_single_process_in_out_dict( + mocker: MockerFixture, use_node_builder: bool +) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + if use_node_builder: + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + else: + chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( nodes={"one": chain}, @@ -591,9 +627,15 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: assert app.invoke(2) == {"output": 3} -def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: +@pytest.mark.parametrize("use_node_builder", [True, False]) +def test_invoke_single_process_in_dict_out_dict( + mocker: MockerFixture, use_node_builder: bool +) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + if use_node_builder: + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + else: + chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") app = Pregel( nodes={"one": chain}, @@ -616,10 +658,17 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: assert app.invoke({"input": 2}) == {"output": 3} -def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: +@pytest.mark.parametrize("use_node_builder", [True, False]) +def test_invoke_two_processes_in_out( + mocker: MockerFixture, use_node_builder: bool +) -> 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("inbox") | add_one | Channel.write_to("output") + if use_node_builder: + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output") + else: + one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") app = Pregel( nodes={"one": one, "two": two}, @@ -738,132 +787,13 @@ def test_run_from_checkpoint_id_retains_previous_writes( assert _get_tasks(new_history, 1) == _get_tasks(history, 0) -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("inbox") - | RunnableLambda(add_one).batch - | RunnablePassthrough(lambda _: time.sleep(0.1)) - | Channel.write_to("output").batch - ) - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": Topic(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels=["input", "inbox"], - stream_channels=["output", "inbox"], - output_channels=["output"], - ) - - # [12 + 1, 2 + 1 + 1] - assert [ - *app.stream( - {"input": 2, "inbox": 12}, output_keys="output", stream_mode="updates" - ) - ] == [ - {"one": None}, - {"two": 13}, - {"two": 4}, - ] - assert [*app.stream({"input": 2, "inbox": 12}, output_keys="output")] == [ - 13, - 4, - ] - - assert [*app.stream({"input": 2, "inbox": 12}, stream_mode="updates")] == [ - {"one": {"inbox": 3}}, - {"two": {"output": 13}}, - {"two": {"output": 4}}, - ] - assert [*app.stream({"input": 2, "inbox": 12})] == [ - {"inbox": [3], "output": 13}, - {"output": 4}, - ] - assert [*app.stream({"input": 2, "inbox": 12}, stream_mode="debug")] == [ - { - "type": "task", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "one", - "input": 2, - "triggers": ("input",), - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "two", - "input": [12], - "triggers": ("inbox",), - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "one", - "result": [("inbox", 3)], - "error": None, - "interrupts": [], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "two", - "result": [("output", 13)], - "error": None, - "interrupts": [], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "two", - "input": [3], - "triggers": ("inbox",), - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "two", - "result": [("output", 4)], - "error": None, - "interrupts": [], - }, - }, - ] - - def test_batch_two_processes_in_out() -> None: def add_one_with_delay(inp: int) -> int: time.sleep(inp / 10) return inp + 1 - one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") - two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output") + one = NodeBuilder().subscribe_only("input").do(add_one_with_delay).write_to("one") + two = NodeBuilder().subscribe_only("one").do(add_one_with_delay).write_to("output") app = Pregel( nodes={"one": one, "two": two}, @@ -900,12 +830,12 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} + nodes = {"-1": NodeBuilder().subscribe_only("input").do(add_one).write_to("-1")} for i in range(test_size - 2): nodes[str(i)] = ( - Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) + NodeBuilder().subscribe_only(str(i - 1)).do(add_one).write_to(str(i)) ) - nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") + nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output") app = Pregel( nodes=nodes, @@ -928,12 +858,12 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} + nodes = {"-1": NodeBuilder().subscribe_only("input").do(add_one).write_to("-1")} for i in range(test_size - 2): nodes[str(i)] = ( - Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) + NodeBuilder().subscribe_only(str(i - 1)).do(add_one).write_to(str(i)) ) - nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") + nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output") app = Pregel( nodes=nodes, @@ -965,8 +895,8 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") app = Pregel( nodes={"one": one, "two": two}, @@ -998,8 +928,8 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") app = Pregel( nodes={"one": one, "two": two}, @@ -1034,10 +964,12 @@ def test_invoke_checkpoint_two( return input one = ( - Channel.subscribe_to(["input"]).join(["total"]) - | add_one - | Channel.write_to("output", "total") - | raise_if_above_10 + NodeBuilder() + .subscribe_to("input") + .read_from("total") + .do(add_one) + .write_to("output", "total") + .do(raise_if_above_10) ) app = Pregel( @@ -1681,10 +1613,12 @@ def test_invoke_checkpoint_three( return input one = ( - Channel.subscribe_to(["input"]).join(["total"]) - | adder - | Channel.write_to("output", "total") - | raise_if_above_10 + NodeBuilder() + .subscribe_to("input") + .read_from("total") + .do(adder) + .write_to("output", "total") + .do(raise_if_above_10) ) app = Pregel( @@ -1806,10 +1740,10 @@ 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)) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + chain_three = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") chain_four = ( - Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output") + NodeBuilder().subscribe_only("inbox").do(add_10_each).write_to("output") ) app = Pregel( @@ -1845,7 +1779,7 @@ def test_invoke_join_then_call_other_pregel( inner_app = Pregel( nodes={ - "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") + "one": NodeBuilder().subscribe_only("input").do(add_one).write_to("output") }, channels={ "output": LastValue(int), @@ -1855,18 +1789,14 @@ def test_invoke_join_then_call_other_pregel( output_channels="output", ) - one = ( - Channel.subscribe_to("input") - | add_10_each - | Channel.write_to("inbox_one").map() - ) + one = NodeBuilder().subscribe_only("input").do(add_10_each).write_to("inbox_one") two = ( - Channel.subscribe_to("inbox_one") - | inner_app.map() - | sorted - | Channel.write_to("outbox_one") + NodeBuilder() + .subscribe_only("inbox_one") + .do(inner_app.map()) + .write_to("outbox_one") ) - chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output") + chain_three = NodeBuilder().subscribe_only("outbox_one").do(sum).write_to("output") app = Pregel( nodes={ @@ -1905,9 +1835,9 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = ( - Channel.subscribe_to("input") | add_one | Channel.write_to("output", "between") + NodeBuilder().subscribe_only("input").do(add_one).write_to("output", "between") ) - two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") + two = NodeBuilder().subscribe_only("between").do(add_one).write_to("output") app = Pregel( nodes={"one": one, "two": two}, @@ -1933,8 +1863,8 @@ 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) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") - two = Channel.subscribe_to("between") | add_one + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("between") + two = NodeBuilder().subscribe_only("between").do(add_one) app = Pregel( nodes={"one": one, "two": two}, @@ -1955,8 +1885,8 @@ 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) - one = Channel.subscribe_to("between") | add_one | Channel.write_to("output") - two = Channel.subscribe_to("between") | add_one + one = NodeBuilder().subscribe_only("between").do(add_one).write_to("output") + two = NodeBuilder().subscribe_only("between").do(add_one) with pytest.raises(TypeError): Pregel(nodes={"one": one, "two": two}) @@ -4706,7 +4636,7 @@ def test_xray_lance(snapshot: SnapshotAssertion): def test_channel_values(sync_checkpointer: BaseCheckpointSaver) -> None: config = {"configurable": {"thread_id": "1"}} - chain = Channel.subscribe_to("input") | Channel.write_to("output") + chain = NodeBuilder().subscribe_only("input").write_to("output") app = Pregel( nodes={ "one": chain, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 0b64ebda4..8aa1f8a41 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -47,7 +47,7 @@ from langgraph.func import entrypoint, task from langgraph.graph import END, Graph, StateGraph from langgraph.graph.message import MessagesState, add_messages from langgraph.prebuilt.tool_node import ToolNode -from langgraph.pregel import Channel, GraphRecursionError, Pregel, StateSnapshot +from langgraph.pregel import GraphRecursionError, NodeBuilder, Pregel, StateSnapshot from langgraph.pregel.loop import AsyncPregelLoop from langgraph.pregel.retry import RetryPolicy from langgraph.pregel.runner import PregelRunner @@ -1408,7 +1408,7 @@ async def test_node_schemas_custom_output() -> None: async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") app = Pregel( nodes={ @@ -1457,9 +1457,10 @@ async def test_invoke_single_process_in_out_falsy_values(falsy_value: Any) -> No async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) chain = ( - Channel.subscribe_to("input") - | add_one - | Channel.write_to("output", fixed=5, output_plus_one=lambda x: x + 1) + NodeBuilder() + .subscribe_only("input") + .do(add_one) + .write_to("output", fixed=5, output_plus_one=lambda x: x + 1) ) app = Pregel( @@ -1496,7 +1497,7 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") app = Pregel( nodes={"one": chain}, @@ -1521,7 +1522,7 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - chain = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") app = Pregel( nodes={"one": chain}, @@ -1547,8 +1548,8 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - two = Channel.subscribe_to("inbox") | add_one | Channel.write_to("output") + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output") app = Pregel( nodes={"one": one, "two": two}, @@ -1605,135 +1606,13 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None: assert step == 2 -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("inbox") - | RunnableLambda(add_one).abatch - | Channel.write_to("output").abatch - ) - - app = Pregel( - nodes={"one": one, "two": two}, - channels={ - "inbox": Topic(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels=["input", "inbox"], - stream_channels=["output", "inbox"], - output_channels=["output"], - ) - - # [12 + 1, 2 + 1 + 1] - assert [ - c - async for c in app.astream( - {"input": 2, "inbox": 12}, output_keys="output", stream_mode="updates" - ) - ] == [ - {"one": None}, - {"two": 13}, - {"two": 4}, - ] - assert [ - c async for c in app.astream({"input": 2, "inbox": 12}, output_keys="output") - ] == [13, 4] - - assert [ - c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="updates") - ] == [ - {"one": {"inbox": 3}}, - {"two": {"output": 13}}, - {"two": {"output": 4}}, - ] - assert [c async for c in app.astream({"input": 2, "inbox": 12})] == [ - {"inbox": [3], "output": 13}, - {"output": 4}, - ] - assert [ - c async for c in app.astream({"input": 2, "inbox": 12}, stream_mode="debug") - ] == [ - { - "type": "task", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "one", - "input": 2, - "triggers": ("input",), - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "two", - "input": [12], - "triggers": ("inbox",), - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "one", - "result": [("inbox", 3)], - "error": None, - "interrupts": [], - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 0, - "payload": { - "id": AnyStr(), - "name": "two", - "result": [("output", 13)], - "error": None, - "interrupts": [], - }, - }, - { - "type": "task", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "two", - "input": [3], - "triggers": ("inbox",), - }, - }, - { - "type": "task_result", - "timestamp": AnyStr(), - "step": 1, - "payload": { - "id": AnyStr(), - "name": "two", - "result": [("output", 4)], - "error": None, - "interrupts": [], - }, - }, - ] - - async def test_batch_two_processes_in_out() -> None: async def add_one_with_delay(inp: int) -> int: await asyncio.sleep(inp / 10) return inp + 1 - one = Channel.subscribe_to("input") | add_one_with_delay | Channel.write_to("one") - two = Channel.subscribe_to("one") | add_one_with_delay | Channel.write_to("output") + one = NodeBuilder().subscribe_only("input").do(add_one_with_delay).write_to("one") + two = NodeBuilder().subscribe_only("one").do(add_one_with_delay).write_to("output") app = Pregel( nodes={"one": one, "two": two}, @@ -1770,12 +1649,12 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} + nodes = {"-1": NodeBuilder().subscribe_only("input").do(add_one).write_to("-1")} for i in range(test_size - 2): nodes[str(i)] = ( - Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) + NodeBuilder().subscribe_only(str(i - 1)).do(add_one).write_to(str(i)) ) - nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") + nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output") app = Pregel( nodes=nodes, @@ -1799,12 +1678,12 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None: test_size = 100 add_one = mocker.Mock(side_effect=lambda x: x + 1) - nodes = {"-1": Channel.subscribe_to("input") | add_one | Channel.write_to("-1")} + nodes = {"-1": NodeBuilder().subscribe_only("input").do(add_one).write_to("-1")} for i in range(test_size - 2): nodes[str(i)] = ( - Channel.subscribe_to(str(i - 1)) | add_one | Channel.write_to(str(i)) + NodeBuilder().subscribe_only(str(i - 1)).do(add_one).write_to(str(i)) ) - nodes["last"] = Channel.subscribe_to(str(i)) | add_one | Channel.write_to("output") + nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output") app = Pregel( nodes=nodes, @@ -1839,8 +1718,8 @@ async def test_invoke_two_processes_two_in_two_out_invalid( ) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") app = Pregel( nodes={"one": one, "two": two}, @@ -1857,8 +1736,8 @@ async def test_invoke_two_processes_two_in_two_out_invalid( async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("output") - two = Channel.subscribe_to("input") | add_one | Channel.write_to("output") + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") + two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output") app = Pregel( nodes={"one": one, "two": two}, @@ -1893,10 +1772,12 @@ async def test_invoke_checkpoint( return input one = ( - Channel.subscribe_to(["input"]).join(["total"]) - | add_one - | Channel.write_to("output", "total") - | raise_if_above_10 + NodeBuilder() + .subscribe_to("input") + .read_from("total") + .do(add_one) + .write_to("output", "total") + .do(raise_if_above_10) ) app = Pregel( @@ -3917,10 +3798,12 @@ async def test_invoke_checkpoint_three( return input one = ( - Channel.subscribe_to(["input"]).join(["total"]) - | add_one - | Channel.write_to("output", "total") - | raise_if_above_10 + NodeBuilder() + .subscribe_to("input") + .read_from("total") + .do(add_one) + .write_to("output", "total") + .do(raise_if_above_10) ) app = Pregel( @@ -4041,10 +3924,10 @@ 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)) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") - chain_three = Channel.subscribe_to("input") | add_one | Channel.write_to("inbox") + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") + chain_three = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox") chain_four = ( - Channel.subscribe_to("inbox") | add_10_each | Channel.write_to("output") + NodeBuilder().subscribe_only("inbox").do(add_10_each).write_to("output") ) app = Pregel( @@ -4073,79 +3956,13 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) - ] -async def test_invoke_join_then_call_other_pregel( - mocker: MockerFixture, async_checkpointer: BaseCheckpointSaver -) -> None: - add_one = mocker.Mock(side_effect=lambda x: x + 1) - add_10_each = mocker.Mock(side_effect=lambda x: [y + 10 for y in x]) - - inner_app = Pregel( - nodes={ - "one": Channel.subscribe_to("input") | add_one | Channel.write_to("output") - }, - channels={ - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - ) - - one = ( - Channel.subscribe_to("input") - | add_10_each - | Channel.write_to("inbox_one").map() - ) - two = ( - Channel.subscribe_to("inbox_one") - | inner_app.map() - | sorted - | Channel.write_to("outbox_one") - ) - chain_three = Channel.subscribe_to("outbox_one") | sum | Channel.write_to("output") - - app = Pregel( - nodes={ - "one": one, - "two": two, - "chain_three": chain_three, - }, - channels={ - "inbox_one": Topic(int), - "outbox_one": LastValue(int), - "output": LastValue(int), - "input": LastValue(int), - }, - input_channels="input", - output_channels="output", - ) - - # Then invoke pubsub - for _ in range(10): - assert await app.ainvoke([2, 3]) == 27 - - assert await asyncio.gather(*(app.ainvoke([2, 3]) for _ in range(10))) == [ - 27 for _ in range(10) - ] - - # add checkpointer - app.checkpointer = async_checkpointer - # subgraph is called twice, and that works - assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 - - # set inner graph checkpointer NeverCheckpoint - inner_app.checkpointer = False - # subgraph still called twice, but checkpointing for inner graph is disabled - assert await app.ainvoke([2, 3], {"configurable": {"thread_id": "1"}}) == 27 - - async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None: add_one = mocker.Mock(side_effect=lambda x: x + 1) one = ( - Channel.subscribe_to("input") | add_one | Channel.write_to("output", "between") + NodeBuilder().subscribe_only("input").do(add_one).write_to("output", "between") ) - two = Channel.subscribe_to("between") | add_one | Channel.write_to("output") + two = NodeBuilder().subscribe_only("between").do(add_one).write_to("output") app = Pregel( nodes={"one": one, "two": two}, @@ -4168,8 +3985,8 @@ 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) - one = Channel.subscribe_to("input") | add_one | Channel.write_to("between") - two = Channel.subscribe_to("between") | add_one + one = NodeBuilder().subscribe_only("input").do(add_one).write_to("between") + two = NodeBuilder().subscribe_only("between").do(add_one) app = Pregel( nodes={"one": one, "two": two}, @@ -8796,7 +8613,7 @@ async def test_draw_invalid(): "id": "__start__", "type": "runnable", "data": { - "id": ["langchain", "schema", "runnable", "RunnablePassthrough"], + "id": ["langgraph", "utils", "runnable", "RunnableCallable"], "name": "__start__", }, },