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

Pregel: Add NodeBuilder class to replace Channel.subscribe_to (#4819)

Browse Source
This commit is contained in:
Nuno Campos
2025-05-28 14:04:08 -07:00
committed by GitHub
parent 8386eaceab 85522ec6a6
commit 3737d886e7
14 changed files with 494 additions and 616 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/docs/concepts/pregel.md
+37 -49
View File
@@ -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__': <langgraph.pregel.read.PregelNode at 0x7d05e3ba1810>,
'write_essay': <langgraph.pregel.read.PregelNode at 0x7d05e3ba14d0>,
'score_essay': <langgraph.pregel.read.PregelNode at 0x7d05e3ba1710>}
@@ -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': <langgraph.pregel.read.PregelNode object at 0x7d05e2f9aad0>}
Channels:
Channels:
{'__start__': <langgraph.channels.ephemeral_value.EphemeralValue object at 0x7d05e2c906c0>, '__end__': <langgraph.channels.last_value.LastValue object at 0x7d05e2c90c40>, '__previous__': <langgraph.channels.last_value.LastValue object at 0x7d05e1007280>}
```
docs/docs/reference/pregel.md
+15
View File
@@ -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
libs/langgraph/langgraph/graph/graph.py
+4 -2
View File
@@ -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)
libs/langgraph/langgraph/pregel/__init__.py
+209 -52
View File
@@ -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
libs/langgraph/langgraph/pregel/read.py
+19 -37
View File
@@ -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]:
libs/langgraph/langgraph/pregel/validate.py
+1 -1
View File
@@ -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:
libs/langgraph/langgraph/pregel/write.py
+2 -16
View File
@@ -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:
libs/langgraph/langgraph/types.py
+1 -1
View File
@@ -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
libs/langgraph/tests/__snapshots__/test_large_cases.ambr
+9 -9
View File
@@ -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__"
}
libs/langgraph/tests/__snapshots__/test_pregel.ambr
+39 -39
View File
@@ -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__',
}),
libs/langgraph/tests/test_large_cases.py
+5 -4
View File
@@ -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},
libs/langgraph/tests/test_large_cases_async.py
+3 -3
View File
@@ -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={
libs/langgraph/tests/test_pregel.py
+108 -178
View File
@@ -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,
libs/langgraph/tests/test_pregel_async.py
+42 -225
View File
@@ -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__",
},
},