mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-10-28 01:41:27 +01:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d9f65ef5ec |
No files matched your search
@@ -255,6 +255,11 @@ R = TypeVar("R")
|
||||
S = TypeVar("S")
|
||||
|
||||
|
||||
class _EntrypointPregel(Pregel):
|
||||
"""The graph that `@entrypoint` builds. `Pregel` itself cannot be constructed
|
||||
directly."""
|
||||
|
||||
|
||||
# The decorator was wrapped in a class to support the `final` attribute.
|
||||
# In this form, the `final` attribute should play nicely with IDE autocompletion,
|
||||
# and type checking tools.
|
||||
@@ -573,7 +578,7 @@ class entrypoint(Generic[ContextT]):
|
||||
else:
|
||||
output_type = save_type = sig.return_annotation
|
||||
|
||||
graph: Pregel[Any, ContextT, Any, Any] = Pregel(
|
||||
graph: Pregel[Any, ContextT, Any, Any] = _EntrypointPregel(
|
||||
nodes={
|
||||
func.__name__: PregelNode(
|
||||
bound=bound,
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
from langgraph.pregel.main import NodeBuilder, Pregel
|
||||
from langgraph.pregel.main import Pregel
|
||||
|
||||
__all__ = ("Pregel", "NodeBuilder")
|
||||
__all__ = ("Pregel",)
|
||||
@@ -47,9 +47,7 @@ def validate_graph(
|
||||
f"not in known channels: '{repr(sorted(channels))[:100]}'"
|
||||
)
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Invalid node type {type(node)}, expected PregelNode or NodeBuilder"
|
||||
)
|
||||
raise TypeError(f"Invalid node type {type(node)}, expected PregelNode")
|
||||
|
||||
for chan in subscribed_channels:
|
||||
if chan not in channels:
|
||||
|
||||
@@ -17,7 +17,6 @@ from collections.abc import (
|
||||
Sequence,
|
||||
)
|
||||
from dataclasses import is_dataclass, replace
|
||||
from datetime import timedelta
|
||||
from functools import partial
|
||||
from inspect import isclass
|
||||
from typing import (
|
||||
@@ -35,7 +34,6 @@ from langchain_core.globals import get_debug
|
||||
from langchain_core.runnables import (
|
||||
RunnableSequence,
|
||||
)
|
||||
from langchain_core.runnables.base import Input, Output
|
||||
from langchain_core.runnables.config import (
|
||||
RunnableConfig,
|
||||
get_async_callback_manager_for_config,
|
||||
@@ -93,13 +91,6 @@ from langgraph._internal._queue import ( # type: ignore[attr-defined]
|
||||
AsyncQueue,
|
||||
SyncQueue,
|
||||
)
|
||||
from langgraph._internal._runnable import (
|
||||
Runnable,
|
||||
RunnableLike,
|
||||
RunnableSeq,
|
||||
coerce_to_runnable,
|
||||
)
|
||||
from langgraph._internal._timeout import coerce_timeout_policy
|
||||
from langgraph._internal._typing import MISSING, DeprecatedKwargs
|
||||
from langgraph.callbacks import (
|
||||
GraphInterruptEvent,
|
||||
@@ -155,7 +146,7 @@ from langgraph.pregel._messages import (
|
||||
StreamMessagesHandlerV2,
|
||||
ensure_message_ids,
|
||||
)
|
||||
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
|
||||
from langgraph.pregel._read import PregelNode
|
||||
from langgraph.pregel._retry import RetryPolicy
|
||||
from langgraph.pregel._runner import PregelRunner
|
||||
from langgraph.pregel._task_status import read_task_statuses
|
||||
@@ -165,7 +156,6 @@ from langgraph.pregel._utils import (
|
||||
validate_timeout_supported,
|
||||
)
|
||||
from langgraph.pregel._validate import validate_graph, validate_keys
|
||||
from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.pregel.debug import get_bolded_text, get_colored_text, tasks_w_writes
|
||||
from langgraph.pregel.protocol import PregelProtocol, StreamChunk, StreamProtocol
|
||||
from langgraph.runtime import (
|
||||
@@ -197,7 +187,6 @@ from langgraph.types import (
|
||||
StateUpdate,
|
||||
StreamMode,
|
||||
StreamPart,
|
||||
TimeoutPolicy,
|
||||
ensure_valid_checkpointer,
|
||||
)
|
||||
from langgraph.typing import ContextT, InputT, OutputT, StateT
|
||||
@@ -208,182 +197,7 @@ try:
|
||||
except ImportError:
|
||||
_StreamingCallbackHandler = None # type: ignore
|
||||
|
||||
__all__ = ("NodeBuilder", "Pregel")
|
||||
|
||||
_WriteValue = Callable[[Input], Output] | Any
|
||||
|
||||
|
||||
class NodeBuilder:
|
||||
__slots__ = (
|
||||
"_channels",
|
||||
"_triggers",
|
||||
"_tags",
|
||||
"_metadata",
|
||||
"_writes",
|
||||
"_bound",
|
||||
"_retry_policy",
|
||||
"_cache_policy",
|
||||
"_timeout",
|
||||
)
|
||||
|
||||
_channels: str | list[str]
|
||||
_triggers: list[str]
|
||||
_tags: list[str]
|
||||
_metadata: dict[str, Any]
|
||||
_writes: list[ChannelWriteEntry]
|
||||
_bound: Runnable
|
||||
_retry_policy: list[RetryPolicy]
|
||||
_cache_policy: CachePolicy | None
|
||||
_timeout: TimeoutPolicy | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
) -> None:
|
||||
self._channels = []
|
||||
self._triggers = []
|
||||
self._tags = []
|
||||
self._metadata = {}
|
||||
self._writes = []
|
||||
self._bound = DEFAULT_BOUND
|
||||
self._retry_policy = []
|
||||
self._cache_policy = None
|
||||
self._timeout = None
|
||||
|
||||
def subscribe_only(
|
||||
self,
|
||||
channel: str,
|
||||
) -> Self:
|
||||
"""Subscribe to a single channel."""
|
||||
if 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, str):
|
||||
raise ValueError(
|
||||
"Cannot subscribe to channels when subscribed to a single channel"
|
||||
)
|
||||
if read:
|
||||
if not self._channels:
|
||||
self._channels = list(channels)
|
||||
else:
|
||||
self._channels.extend(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 isinstance(self._channels, list), (
|
||||
"Cannot read additional channels when subscribed to single channels"
|
||||
)
|
||||
self._channels.extend(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 add_retry_policies(self, *policies: RetryPolicy) -> Self:
|
||||
"""Adds retry policies to the node."""
|
||||
self._retry_policy.extend(policies)
|
||||
return self
|
||||
|
||||
def add_cache_policy(self, policy: CachePolicy) -> Self:
|
||||
"""Adds cache policies to the node."""
|
||||
self._cache_policy = policy
|
||||
return self
|
||||
|
||||
def set_timeout(self, timeout: float | timedelta | TimeoutPolicy | None) -> Self:
|
||||
"""Set the per-attempt timeout policy for this node."""
|
||||
self._timeout = coerce_timeout_policy(timeout)
|
||||
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._retry_policy,
|
||||
cache_policy=self._cache_policy,
|
||||
timeout=self._timeout,
|
||||
)
|
||||
__all__ = ("Pregel",)
|
||||
|
||||
|
||||
# Kwargs that ``stream_events(version="v3")`` / ``astream_events(version="v3")``
|
||||
@@ -461,255 +275,22 @@ class Pregel(
|
||||
PregelProtocol[StateT, ContextT, InputT, OutputT],
|
||||
Generic[StateT, ContextT, InputT, OutputT],
|
||||
):
|
||||
"""Pregel manages the runtime behavior for LangGraph applications.
|
||||
"""The compiled graph that runs a LangGraph application.
|
||||
|
||||
## Overview
|
||||
`StateGraph.compile()` and the `@entrypoint` decorator return an instance of
|
||||
this class. It provides the methods that run and inspect the application,
|
||||
such as `invoke`, `stream`, `get_state` and `update_state`.
|
||||
|
||||
Pregel combines [**actors**](https://en.wikipedia.org/wiki/Actor_model)
|
||||
and **channels** into a single application.
|
||||
**Actors** read data from channels and write data to channels.
|
||||
Pregel organizes the execution of the application into multiple steps,
|
||||
following the **Pregel Algorithm**/**Bulk Synchronous Parallel** model.
|
||||
The runtime is named after Google's
|
||||
[Pregel algorithm](https://research.google/pubs/pub37252/). It runs the
|
||||
application in steps. In each step it plans which nodes to run, runs them in
|
||||
parallel, and then applies their writes to the graph's state, so writes
|
||||
become visible only in the next step. The run ends when no node is
|
||||
scheduled, or when the recursion limit is reached.
|
||||
|
||||
Each step consists of three phases:
|
||||
|
||||
- **Plan**: Determine which **actors** to execute in this step. For example,
|
||||
in the first step, select the **actors** that subscribe to the special
|
||||
**input** channels; in subsequent steps,
|
||||
select the **actors** that subscribe to channels updated in the previous step.
|
||||
- **Execution**: Execute all selected **actors** in parallel,
|
||||
until all complete, or one fails, or a timeout is reached. During this
|
||||
phase, channel updates are invisible to actors until the next step.
|
||||
- **Update**: Update the channels with the values written by the **actors**
|
||||
in this step.
|
||||
|
||||
Repeat until no **actors** are selected for execution, or a maximum number of
|
||||
steps is reached.
|
||||
|
||||
## 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.
|
||||
|
||||
## Channels
|
||||
|
||||
Channels are used to communicate between actors (`PregelNodes`).
|
||||
Each channel has a value type, an update type, and an update function – which
|
||||
takes a sequence of updates and
|
||||
modifies the stored value. Channels can be used to send data from one chain to
|
||||
another, or to send data from a chain to itself in a future step. LangGraph
|
||||
provides a number of built-in channels:
|
||||
|
||||
### Basic channels: LastValue and Topic
|
||||
|
||||
- `LastValue`: The default channel, stores the last value sent to the channel,
|
||||
useful for input and output values, or for sending data from one step to the next
|
||||
- `Topic`: A configurable PubSub Topic, useful for sending multiple values
|
||||
between *actors*, or for accumulating output. Can be configured to deduplicate
|
||||
values, and/or to accumulate values over the course of multiple steps.
|
||||
|
||||
### Advanced channels: Context and BinaryOperatorAggregate
|
||||
|
||||
- `Context`: exposes the value of a context manager, managing its lifecycle.
|
||||
Useful for accessing external resources that require setup and/or teardown. e.g.
|
||||
`client = Context(httpx.Client)`
|
||||
- `BinaryOperatorAggregate`: stores a persistent value, updated by applying
|
||||
a binary operator to the current value and each update
|
||||
sent to the channel, useful for computing aggregates over multiple steps. e.g.
|
||||
`total = BinaryOperatorAggregate(int, operator.add)`
|
||||
|
||||
## Examples
|
||||
|
||||
Most users will interact with Pregel via a
|
||||
[StateGraph (Graph API)][langgraph.graph.StateGraph] or via an
|
||||
[entrypoint (Functional API)][langgraph.func.entrypoint].
|
||||
|
||||
However, for **advanced** use cases, Pregel can be used directly. If you're
|
||||
not sure whether you need to use Pregel directly, then the answer is probably no
|
||||
- you should use the Graph API or Functional API instead. These are higher-level
|
||||
interfaces that will compile down to Pregel under the hood.
|
||||
|
||||
Here are some examples to give you a sense of how it works:
|
||||
|
||||
Example: Single node application
|
||||
```python
|
||||
from langgraph.channels import EphemeralValue
|
||||
from langgraph.pregel import Pregel, NodeBuilder
|
||||
|
||||
node1 = (
|
||||
NodeBuilder().subscribe_only("a")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("b")
|
||||
)
|
||||
|
||||
app = Pregel(
|
||||
nodes={"node1": node1},
|
||||
channels={
|
||||
"a": EphemeralValue(str),
|
||||
"b": EphemeralValue(str),
|
||||
},
|
||||
input_channels=["a"],
|
||||
output_channels=["b"],
|
||||
)
|
||||
|
||||
app.invoke({"a": "foo"})
|
||||
```
|
||||
|
||||
```con
|
||||
{'b': 'foofoo'}
|
||||
```
|
||||
|
||||
Example: Using multiple nodes and multiple output channels
|
||||
```python
|
||||
from langgraph.channels import LastValue, EphemeralValue
|
||||
from langgraph.pregel import Pregel, NodeBuilder
|
||||
|
||||
node1 = (
|
||||
NodeBuilder().subscribe_only("a")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("b")
|
||||
)
|
||||
|
||||
node2 = (
|
||||
NodeBuilder().subscribe_to("b")
|
||||
.do(lambda x: x["b"] + x["b"])
|
||||
.write_to("c")
|
||||
)
|
||||
|
||||
|
||||
app = Pregel(
|
||||
nodes={"node1": node1, "node2": node2},
|
||||
channels={
|
||||
"a": EphemeralValue(str),
|
||||
"b": LastValue(str),
|
||||
"c": EphemeralValue(str),
|
||||
},
|
||||
input_channels=["a"],
|
||||
output_channels=["b", "c"],
|
||||
)
|
||||
|
||||
app.invoke({"a": "foo"})
|
||||
```
|
||||
|
||||
```con
|
||||
{'b': 'foofoo', 'c': 'foofoofoofoo'}
|
||||
```
|
||||
|
||||
Example: Using a Topic channel
|
||||
```python
|
||||
from langgraph.channels import LastValue, EphemeralValue, Topic
|
||||
from langgraph.pregel import Pregel, NodeBuilder
|
||||
|
||||
node1 = (
|
||||
NodeBuilder().subscribe_only("a")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("b", "c")
|
||||
)
|
||||
|
||||
node2 = (
|
||||
NodeBuilder().subscribe_only("b")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("c")
|
||||
)
|
||||
|
||||
|
||||
app = Pregel(
|
||||
nodes={"node1": node1, "node2": node2},
|
||||
channels={
|
||||
"a": EphemeralValue(str),
|
||||
"b": EphemeralValue(str),
|
||||
"c": Topic(str, accumulate=True),
|
||||
},
|
||||
input_channels=["a"],
|
||||
output_channels=["c"],
|
||||
)
|
||||
|
||||
app.invoke({"a": "foo"})
|
||||
```
|
||||
|
||||
```pycon
|
||||
{"c": ["foofoo", "foofoofoofoo"]}
|
||||
```
|
||||
|
||||
Example: Using a `BinaryOperatorAggregate` channel
|
||||
```python
|
||||
from langgraph.channels import EphemeralValue, BinaryOperatorAggregate
|
||||
from langgraph.pregel import Pregel, NodeBuilder
|
||||
|
||||
|
||||
node1 = (
|
||||
NodeBuilder().subscribe_only("a")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("b", "c")
|
||||
)
|
||||
|
||||
node2 = (
|
||||
NodeBuilder().subscribe_only("b")
|
||||
.do(lambda x: x + x)
|
||||
.write_to("c")
|
||||
)
|
||||
|
||||
|
||||
def reducer(current, update):
|
||||
if current:
|
||||
return current + " | " + update
|
||||
else:
|
||||
return update
|
||||
|
||||
|
||||
app = Pregel(
|
||||
nodes={"node1": node1, "node2": node2},
|
||||
channels={
|
||||
"a": EphemeralValue(str),
|
||||
"b": EphemeralValue(str),
|
||||
"c": BinaryOperatorAggregate(str, operator=reducer),
|
||||
},
|
||||
input_channels=["a"],
|
||||
output_channels=["c"],
|
||||
)
|
||||
|
||||
app.invoke({"a": "foo"})
|
||||
```
|
||||
|
||||
```con
|
||||
{'c': 'foofoo | foofoofoofoo'}
|
||||
```
|
||||
|
||||
Example: Introducing a cycle
|
||||
This example demonstrates how to introduce a cycle in the graph, by having
|
||||
a chain write to a channel it subscribes to.
|
||||
|
||||
Execution will continue until a `None` value is written to the channel.
|
||||
|
||||
```python
|
||||
from langgraph.channels import EphemeralValue
|
||||
from langgraph.pregel import Pregel, NodeBuilder, ChannelWriteEntry
|
||||
|
||||
example_node = (
|
||||
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(
|
||||
nodes={"example_node": example_node},
|
||||
channels={
|
||||
"value": EphemeralValue(str),
|
||||
},
|
||||
input_channels=["value"],
|
||||
output_channels=["value"],
|
||||
)
|
||||
|
||||
app.invoke({"value": "a"})
|
||||
```
|
||||
|
||||
```con
|
||||
{'value': 'aaaaaaaaaaaaaaaa'}
|
||||
```
|
||||
`Pregel` is not constructed directly. Build graphs with
|
||||
[`StateGraph`][langgraph.graph.StateGraph] or
|
||||
[`@entrypoint`][langgraph.func.entrypoint].
|
||||
"""
|
||||
|
||||
nodes: dict[str, PregelNode]
|
||||
@@ -768,7 +349,7 @@ class Pregel(
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
nodes: dict[str, PregelNode | NodeBuilder],
|
||||
nodes: dict[str, PregelNode],
|
||||
channels: dict[str, BaseChannel | ManagedValueSpec] | None,
|
||||
auto_validate: bool = True,
|
||||
stream_mode: StreamMode = "values",
|
||||
@@ -793,6 +374,13 @@ class Pregel(
|
||||
stream_transformers: Sequence[Callable[[tuple[str, ...]], Any]] | None = None,
|
||||
**deprecated_kwargs: Unpack[DeprecatedKwargs],
|
||||
) -> None:
|
||||
# LangGraph builds graphs through subclasses: `CompiledStateGraph` and
|
||||
# the graph `@entrypoint` returns.
|
||||
if type(self) is Pregel:
|
||||
raise TypeError(
|
||||
"`Pregel` cannot be constructed directly. "
|
||||
"Build graphs with `StateGraph` or `@entrypoint`."
|
||||
)
|
||||
if (
|
||||
config_type := deprecated_kwargs.get("config_type", MISSING)
|
||||
) is not MISSING:
|
||||
@@ -807,9 +395,7 @@ class Pregel(
|
||||
|
||||
checkpointer = ensure_valid_checkpointer(checkpointer)
|
||||
|
||||
self.nodes = {
|
||||
k: v.build() if isinstance(v, NodeBuilder) else v for k, v in nodes.items()
|
||||
}
|
||||
self.nodes = dict(nodes)
|
||||
self.channels = channels or {}
|
||||
if TASKS in self.channels and not isinstance(self.channels[TASKS], Topic):
|
||||
raise ValueError(
|
||||
|
||||
@@ -0,0 +1,222 @@
|
||||
"""Build graphs directly on the Pregel engine, for tests of engine behavior.
|
||||
|
||||
LangGraph does not expose this: applications build graphs with `StateGraph` or
|
||||
`@entrypoint`. Tests use it to set up channels, triggers and writes that those
|
||||
APIs do not produce on their own.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from datetime import timedelta
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables.base import Input, Output
|
||||
from typing_extensions import Self
|
||||
|
||||
from langgraph._internal._runnable import (
|
||||
Runnable,
|
||||
RunnableLike,
|
||||
RunnableSeq,
|
||||
coerce_to_runnable,
|
||||
)
|
||||
from langgraph._internal._timeout import coerce_timeout_policy
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
|
||||
from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.types import CachePolicy, RetryPolicy, TimeoutPolicy
|
||||
|
||||
__all__ = ("NodeBuilder", "RawPregel")
|
||||
|
||||
_WriteValue = Callable[[Input], Output] | Any
|
||||
|
||||
|
||||
class NodeBuilder:
|
||||
"""Builds a `PregelNode` that reads and writes named channels."""
|
||||
|
||||
__slots__ = (
|
||||
"_channels",
|
||||
"_triggers",
|
||||
"_tags",
|
||||
"_metadata",
|
||||
"_writes",
|
||||
"_bound",
|
||||
"_retry_policy",
|
||||
"_cache_policy",
|
||||
"_timeout",
|
||||
)
|
||||
|
||||
_channels: str | list[str]
|
||||
_triggers: list[str]
|
||||
_tags: list[str]
|
||||
_metadata: dict[str, Any]
|
||||
_writes: list[ChannelWriteEntry]
|
||||
_bound: Runnable
|
||||
_retry_policy: list[RetryPolicy]
|
||||
_cache_policy: CachePolicy | None
|
||||
_timeout: TimeoutPolicy | None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
) -> None:
|
||||
self._channels = []
|
||||
self._triggers = []
|
||||
self._tags = []
|
||||
self._metadata = {}
|
||||
self._writes = []
|
||||
self._bound = DEFAULT_BOUND
|
||||
self._retry_policy = []
|
||||
self._cache_policy = None
|
||||
self._timeout = None
|
||||
|
||||
def subscribe_only(
|
||||
self,
|
||||
channel: str,
|
||||
) -> Self:
|
||||
"""Subscribe to a single channel."""
|
||||
if 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, str):
|
||||
raise ValueError(
|
||||
"Cannot subscribe to channels when subscribed to a single channel"
|
||||
)
|
||||
if read:
|
||||
if not self._channels:
|
||||
self._channels = list(channels)
|
||||
else:
|
||||
self._channels.extend(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 isinstance(self._channels, list), (
|
||||
"Cannot read additional channels when subscribed to single channels"
|
||||
)
|
||||
self._channels.extend(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 add_retry_policies(self, *policies: RetryPolicy) -> Self:
|
||||
"""Adds retry policies to the node."""
|
||||
self._retry_policy.extend(policies)
|
||||
return self
|
||||
|
||||
def add_cache_policy(self, policy: CachePolicy) -> Self:
|
||||
"""Adds cache policies to the node."""
|
||||
self._cache_policy = policy
|
||||
return self
|
||||
|
||||
def set_timeout(self, timeout: float | timedelta | TimeoutPolicy | None) -> Self:
|
||||
"""Set the per-attempt timeout policy for this node."""
|
||||
self._timeout = coerce_timeout_policy(timeout)
|
||||
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._retry_policy,
|
||||
cache_policy=self._cache_policy,
|
||||
timeout=self._timeout,
|
||||
)
|
||||
|
||||
|
||||
class RawPregel(Pregel):
|
||||
"""A `Pregel` built from hand-written nodes and channels.
|
||||
|
||||
Accepts `NodeBuilder` values in `nodes`, and builds them.
|
||||
"""
|
||||
|
||||
def __init__(self, *, nodes: dict[str, PregelNode | NodeBuilder], **kwargs: Any):
|
||||
super().__init__(
|
||||
nodes={
|
||||
k: v.build() if isinstance(v, NodeBuilder) else v
|
||||
for k, v in nodes.items()
|
||||
},
|
||||
**kwargs,
|
||||
)
|
||||
@@ -19,8 +19,8 @@ from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Command, Durability, Overwrite
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -279,7 +279,7 @@ def _pregel_graph(seen: list[list[str | None]]) -> Any:
|
||||
.do(look)
|
||||
.write_to("messages")
|
||||
)
|
||||
return Pregel(
|
||||
return RawPregel(
|
||||
nodes={"look": node},
|
||||
channels={
|
||||
"messages": DeltaChannel(_messages_delta_reducer),
|
||||
|
||||
@@ -7,8 +7,9 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
from langgraph.channels.binop import BinaryOperatorAggregate
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.types import Durability
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -19,7 +20,7 @@ def _sorted_extend(current: list, writes: list) -> list:
|
||||
|
||||
def _delta_input_graph() -> Pregel:
|
||||
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain")
|
||||
return Pregel(
|
||||
return RawPregel(
|
||||
nodes={"n": node},
|
||||
channels={
|
||||
"log": DeltaChannel(_sorted_extend),
|
||||
|
||||
@@ -3,7 +3,8 @@ from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.pregel import Pregel
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -17,7 +18,7 @@ def _extend(current: list, writes: list) -> list:
|
||||
def _graph(reads: list) -> Pregel:
|
||||
writer = NodeBuilder().subscribe_only("a").do(lambda _: [1]).write_to("d")
|
||||
reader = NodeBuilder().subscribe_only("d").do(lambda d: reads.append(list(d)))
|
||||
return Pregel(
|
||||
return RawPregel(
|
||||
nodes={"writer": writer, "reader": reader},
|
||||
channels={"a": LastValue(str), "d": DeltaChannel(_extend)},
|
||||
input_channels="a",
|
||||
|
||||
@@ -13,8 +13,9 @@ from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.types import Overwrite
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -36,7 +37,7 @@ def _extend(current: list, writes: list) -> list:
|
||||
|
||||
def _delta_input_graph(saver: InMemorySaver) -> Pregel:
|
||||
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log")
|
||||
return Pregel(
|
||||
return RawPregel(
|
||||
nodes={"n": node},
|
||||
channels={"log": DeltaChannel(_extend), "go": LastValue(int)},
|
||||
input_channels=["log", "go"],
|
||||
|
||||
@@ -30,8 +30,8 @@ from langgraph.channels.delta import DeltaChannel
|
||||
from langgraph.channels.last_value import LastValue
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.graph.message import _messages_delta_reducer
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import Overwrite, StateSnapshot, StateUpdate
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -549,7 +549,7 @@ def _sorted_extend(current: list, writes: list) -> list:
|
||||
|
||||
def _delta_input_graph(snapshot_frequency: int = 1000) -> Any:
|
||||
node = NodeBuilder().subscribe_only("go").do(lambda _: [2]).write_to("log", "plain")
|
||||
return Pregel(
|
||||
return RawPregel(
|
||||
nodes={"n": node},
|
||||
channels={
|
||||
"log": DeltaChannel(_sorted_extend, snapshot_frequency=snapshot_frequency),
|
||||
@@ -649,7 +649,7 @@ def _messages_input_graph() -> Any:
|
||||
.do(lambda _: [HumanMessage("n", id="n")])
|
||||
.write_to("messages")
|
||||
)
|
||||
return Pregel(
|
||||
return RawPregel(
|
||||
nodes={"n": node},
|
||||
channels={
|
||||
"messages": DeltaChannel(_messages_delta_reducer),
|
||||
|
||||
@@ -13,13 +13,13 @@ from langgraph.errors import NodeInterrupt
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.graph.message import MessageGraph
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import GraphOutput, Interrupt, RetryPolicy
|
||||
from langgraph.warnings import (
|
||||
LangGraphDeprecatedSinceV05,
|
||||
LangGraphDeprecatedSinceV10,
|
||||
LangGraphDeprecatedSinceV11,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
|
||||
class PlainState(TypedDict): ...
|
||||
@@ -157,7 +157,7 @@ def test_config_type_deprecation_pregel(mocker: MockerFixture) -> None:
|
||||
LangGraphDeprecatedSinceV10,
|
||||
match="`config_type` is deprecated and will be removed. Please use `context_schema` instead.",
|
||||
):
|
||||
instance = Pregel(
|
||||
instance = RawPregel(
|
||||
nodes={
|
||||
"one": chain,
|
||||
},
|
||||
|
||||
@@ -39,7 +39,6 @@ from langgraph.channels.untracked_value import UntrackedValue
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.graph.message import MessagesState, add_messages
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import (
|
||||
Command,
|
||||
Durability,
|
||||
@@ -62,6 +61,7 @@ from tests.messages import (
|
||||
_AnyIdHumanMessage,
|
||||
_AnyIdToolMessage,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
|
||||
def test_invoke_two_processes_in_out_interrupt(
|
||||
@@ -72,7 +72,7 @@ def test_invoke_two_processes_in_out_interrupt(
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
|
||||
two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"inbox": LastValue(int),
|
||||
|
||||
@@ -39,7 +39,6 @@ from langgraph.constants import END, START
|
||||
from langgraph.graph import MessagesState
|
||||
from langgraph.graph.message import add_messages
|
||||
from langgraph.graph.state import StateGraph
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.types import PregelTask, Send, StateSnapshot, StreamWriter
|
||||
from tests.any_int import AnyInt
|
||||
from tests.any_str import AnyDict, AnyStr, UnsortedSequence
|
||||
@@ -51,6 +50,7 @@ from tests.messages import (
|
||||
_AnyIdHumanMessage,
|
||||
_AnyIdToolMessage,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -61,7 +61,7 @@ async def test_invoke_two_processes_in_out_interrupt(
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
|
||||
two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output")
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"inbox": LastValue(int),
|
||||
|
||||
@@ -66,10 +66,7 @@ from langgraph.errors import GraphRecursionError, InvalidUpdateError, ParentComm
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import MessagesState, _messages_delta_reducer, add_messages
|
||||
from langgraph.pregel import (
|
||||
NodeBuilder,
|
||||
Pregel,
|
||||
)
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel._loop import PregelLoop, SyncPregelLoop
|
||||
from langgraph.pregel._runner import PregelRunner
|
||||
from langgraph.runtime import RunControl
|
||||
@@ -95,6 +92,7 @@ from tests.messages import (
|
||||
_AnyIdHumanMessage,
|
||||
_AnyIdToolMessage,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
pytestmark = pytest.mark.anyio
|
||||
|
||||
@@ -257,7 +255,7 @@ def test_context_json_schema() -> None:
|
||||
x: int
|
||||
y: str = field(default="foo")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={
|
||||
"one": chain,
|
||||
},
|
||||
@@ -449,7 +447,7 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={
|
||||
"one": chain,
|
||||
},
|
||||
@@ -485,7 +483,7 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None:
|
||||
.write_to("output", fixed=5, output_plus_one=lambda x: x + 1)
|
||||
)
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": chain},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -521,7 +519,7 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": chain},
|
||||
channels={"input": LastValue(int), "output": LastValue(int)},
|
||||
input_channels="input",
|
||||
@@ -546,7 +544,7 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": chain},
|
||||
channels={"input": LastValue(int), "output": LastValue(int)},
|
||||
input_channels=["input"],
|
||||
@@ -572,7 +570,7 @@ def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
|
||||
two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"inbox": LastValue(int),
|
||||
@@ -677,7 +675,7 @@ def test_batch_two_processes_in_out() -> None:
|
||||
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(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"one": LastValue(int),
|
||||
@@ -709,7 +707,7 @@ def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
)
|
||||
nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes=nodes,
|
||||
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
|
||||
| {"input": LastValue(int), "output": LastValue(int)},
|
||||
@@ -737,7 +735,7 @@ def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
)
|
||||
nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes=nodes,
|
||||
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
|
||||
| {"input": LastValue(int), "output": LastValue(int)},
|
||||
@@ -770,7 +768,7 @@ def test_invoke_two_processes_two_in_two_out_invalid(mocker: MockerFixture) -> N
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={"output": LastValue(int), "input": LastValue(int)},
|
||||
input_channels="input",
|
||||
@@ -803,7 +801,7 @@ def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture) -> Non
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -844,7 +842,7 @@ def test_invoke_checkpoint_two(
|
||||
.do(raise_if_above_10)
|
||||
)
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one},
|
||||
channels={
|
||||
"total": BinaryOperatorAggregate(int, operator.add),
|
||||
@@ -1451,7 +1449,7 @@ def test_invoke_checkpoint_three(
|
||||
.do(raise_if_above_10)
|
||||
)
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one},
|
||||
channels={
|
||||
"total": BinaryOperatorAggregate(int, operator.add),
|
||||
@@ -1576,7 +1574,7 @@ def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -> None
|
||||
NodeBuilder().subscribe_only("inbox").do(add_10_each).write_to("output")
|
||||
)
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={
|
||||
"one": one,
|
||||
"chain_three": chain_three,
|
||||
@@ -1607,7 +1605,7 @@ def test_invoke_join_then_call_other_pregel(
|
||||
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(
|
||||
inner_app = RawPregel(
|
||||
nodes={
|
||||
"one": NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
},
|
||||
@@ -1628,7 +1626,7 @@ def test_invoke_join_then_call_other_pregel(
|
||||
)
|
||||
chain_three = NodeBuilder().subscribe_only("outbox_one").do(sum).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={
|
||||
"one": one,
|
||||
"two": two,
|
||||
@@ -1669,7 +1667,7 @@ def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> None:
|
||||
)
|
||||
two = NodeBuilder().subscribe_only("between").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -1696,7 +1694,7 @@ def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("between")
|
||||
two = NodeBuilder().subscribe_only("between").do(add_one)
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -1719,7 +1717,7 @@ def test_invoke_two_processes_no_in(mocker: MockerFixture) -> None:
|
||||
two = NodeBuilder().subscribe_only("between").do(add_one)
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
Pregel(nodes={"one": one, "two": two})
|
||||
RawPregel(nodes={"one": one, "two": two})
|
||||
|
||||
|
||||
def test_conditional_entrypoint_to_multiple_state_graph(
|
||||
@@ -4128,7 +4126,7 @@ def test_xray_lance(snapshot: SnapshotAssertion):
|
||||
def test_channel_values(sync_checkpointer: BaseCheckpointSaver) -> None:
|
||||
config = {"configurable": {"thread_id": "1"}}
|
||||
chain = NodeBuilder().subscribe_only("input").write_to("output")
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={
|
||||
"one": chain,
|
||||
},
|
||||
@@ -9507,7 +9505,7 @@ class _IntVersionSaver(InMemorySaver):
|
||||
|
||||
|
||||
def _build_chain_after_a_delta_channel() -> Pregel:
|
||||
return Pregel(
|
||||
return RawPregel(
|
||||
nodes={
|
||||
"a": NodeBuilder().subscribe_only("inp").do(lambda _: ["a"]).write_to("d"),
|
||||
"b": NodeBuilder().subscribe_only("d").do(lambda _: "b").write_to("x"),
|
||||
@@ -9552,7 +9550,7 @@ async def test_aupdate_state_after_an_exit_snapshot_infers_the_last_writer() ->
|
||||
|
||||
|
||||
def test_update_state_after_a_fork_seal_infers_the_subscriber_that_ran() -> None:
|
||||
graph = Pregel(
|
||||
graph = RawPregel(
|
||||
nodes={
|
||||
"a": NodeBuilder()
|
||||
.subscribe_only("inp")
|
||||
@@ -9587,7 +9585,7 @@ def test_update_state_after_a_fork_seal_infers_the_subscriber_that_ran() -> None
|
||||
def test_update_state_after_a_fork_seal_infers_the_node_whose_read_it_advanced() -> (
|
||||
None
|
||||
):
|
||||
graph = Pregel(
|
||||
graph = RawPregel(
|
||||
nodes={
|
||||
"a": NodeBuilder().subscribe_only("inp").do(lambda _: ["a"]).write_to("d"),
|
||||
"b": NodeBuilder().subscribe_only("d").do(lambda _: "b").write_to("out"),
|
||||
|
||||
@@ -69,7 +69,6 @@ from langgraph.errors import (
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import END, START, StateGraph
|
||||
from langgraph.graph.message import MessagesState, _messages_delta_reducer, add_messages
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.pregel._loop import AsyncPregelLoop, PregelLoop
|
||||
from langgraph.pregel._runner import PregelRunner
|
||||
from langgraph.runtime import RunControl
|
||||
@@ -95,6 +94,7 @@ from tests.messages import (
|
||||
_AnyIdHumanMessage,
|
||||
_AnyIdToolMessage,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -1446,7 +1446,7 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={
|
||||
"one": chain,
|
||||
},
|
||||
@@ -1479,7 +1479,7 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N
|
||||
.write_to("output", fixed=5, output_plus_one=lambda x: x + 1)
|
||||
)
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": chain},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -1515,7 +1515,7 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None:
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": chain},
|
||||
channels={"input": LastValue(int), "output": LastValue(int)},
|
||||
input_channels="input",
|
||||
@@ -1540,7 +1540,7 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) ->
|
||||
add_one = mocker.Mock(side_effect=lambda x: x + 1)
|
||||
chain = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": chain},
|
||||
channels={"input": LastValue(int), "output": LastValue(int)},
|
||||
input_channels=["input"],
|
||||
@@ -1567,7 +1567,7 @@ async def test_invoke_two_processes_in_out(mocker: MockerFixture) -> None:
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("inbox")
|
||||
two = NodeBuilder().subscribe_only("inbox").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"inbox": LastValue(int),
|
||||
@@ -1607,7 +1607,7 @@ async def test_batch_two_processes_in_out() -> None:
|
||||
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(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"one": LastValue(int),
|
||||
@@ -1639,7 +1639,7 @@ async def test_invoke_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
)
|
||||
nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes=nodes,
|
||||
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
|
||||
| {"input": LastValue(int), "output": LastValue(int)},
|
||||
@@ -1668,7 +1668,7 @@ async def test_batch_many_processes_in_out(mocker: MockerFixture) -> None:
|
||||
)
|
||||
nodes["last"] = NodeBuilder().subscribe_only(str(i)).do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes=nodes,
|
||||
channels={str(i): LastValue(int) for i in range(-1, test_size - 2)}
|
||||
| {"input": LastValue(int), "output": LastValue(int)},
|
||||
@@ -1704,7 +1704,7 @@ async def test_invoke_two_processes_two_in_two_out_invalid(
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={"output": LastValue(int), "input": LastValue(int)},
|
||||
input_channels="input",
|
||||
@@ -1722,7 +1722,7 @@ async def test_invoke_two_processes_two_in_two_out_valid(mocker: MockerFixture)
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
two = NodeBuilder().subscribe_only("input").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -1763,7 +1763,7 @@ async def test_invoke_checkpoint(
|
||||
.do(raise_if_above_10)
|
||||
)
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one},
|
||||
channels={
|
||||
"total": BinaryOperatorAggregate(int, operator.add),
|
||||
@@ -3641,7 +3641,7 @@ async def test_invoke_checkpoint_three(
|
||||
.do(raise_if_above_10)
|
||||
)
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one},
|
||||
channels={
|
||||
"total": BinaryOperatorAggregate(int, operator.add),
|
||||
@@ -3765,7 +3765,7 @@ async def test_invoke_two_processes_two_in_join_two_out(mocker: MockerFixture) -
|
||||
NodeBuilder().subscribe_only("inbox").do(add_10_each).write_to("output")
|
||||
)
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={
|
||||
"one": one,
|
||||
"chain_three": chain_three,
|
||||
@@ -3799,7 +3799,7 @@ async def test_invoke_two_processes_one_in_two_out(mocker: MockerFixture) -> Non
|
||||
)
|
||||
two = NodeBuilder().subscribe_only("between").do(add_one).write_to("output")
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
@@ -3823,7 +3823,7 @@ async def test_invoke_two_processes_no_out(mocker: MockerFixture) -> None:
|
||||
one = NodeBuilder().subscribe_only("input").do(add_one).write_to("between")
|
||||
two = NodeBuilder().subscribe_only("between").do(add_one)
|
||||
|
||||
app = Pregel(
|
||||
app = RawPregel(
|
||||
nodes={"one": one, "two": two},
|
||||
channels={
|
||||
"input": LastValue(int),
|
||||
|
||||
@@ -0,0 +1,56 @@
|
||||
"""`Pregel` is the type of compiled graphs, built only by `StateGraph` and
|
||||
`@entrypoint`."""
|
||||
|
||||
import pytest
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
import langgraph.pregel
|
||||
from langgraph.channels.ephemeral_value import EphemeralValue
|
||||
from langgraph.func import entrypoint
|
||||
from langgraph.graph import START, StateGraph
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.pregel._read import PregelNode
|
||||
|
||||
|
||||
def test_pregel_cannot_be_constructed_directly() -> None:
|
||||
with pytest.raises(TypeError, match="`Pregel` cannot be constructed directly"):
|
||||
Pregel(
|
||||
nodes={"node": PregelNode(channels="a", triggers=["a"])},
|
||||
channels={"a": EphemeralValue(str)},
|
||||
input_channels="a",
|
||||
output_channels="a",
|
||||
)
|
||||
|
||||
|
||||
def test_node_builder_is_not_exported() -> None:
|
||||
assert langgraph.pregel.__all__ == ("Pregel",)
|
||||
assert not hasattr(langgraph.pregel, "NodeBuilder")
|
||||
|
||||
|
||||
def test_compiled_graphs_are_pregel_instances() -> None:
|
||||
class State(TypedDict):
|
||||
value: int
|
||||
|
||||
def increment(state: State) -> State:
|
||||
return {"value": state["value"] + 1}
|
||||
|
||||
builder = StateGraph(State)
|
||||
builder.add_node("increment", increment)
|
||||
builder.add_edge(START, "increment")
|
||||
graph = builder.compile(checkpointer=InMemorySaver())
|
||||
|
||||
@entrypoint(checkpointer=InMemorySaver())
|
||||
def workflow(value: int) -> int:
|
||||
return value + 1
|
||||
|
||||
assert isinstance(graph, Pregel)
|
||||
assert isinstance(workflow, Pregel)
|
||||
|
||||
# `with_config` copies the graph by calling its constructor again.
|
||||
config: RunnableConfig = {"configurable": {"thread_id": "1"}}
|
||||
assert graph.with_config(recursion_limit=5).invoke({"value": 1}, config) == {
|
||||
"value": 2
|
||||
}
|
||||
assert workflow.with_config(recursion_limit=5).invoke(1, config) == 2
|
||||
@@ -47,7 +47,6 @@ from langgraph.errors import (
|
||||
)
|
||||
from langgraph.func import entrypoint, task
|
||||
from langgraph.graph import END, START, StateGraph, add_messages
|
||||
from langgraph.pregel import NodeBuilder, Pregel
|
||||
from langgraph.pregel._read import PregelNode
|
||||
from langgraph.pregel._retry import (
|
||||
_checkpoint_ns_for_parent_command,
|
||||
@@ -67,6 +66,7 @@ from langgraph.types import (
|
||||
TimeoutPolicy,
|
||||
interrupt,
|
||||
)
|
||||
from tests.raw_pregel import NodeBuilder, RawPregel
|
||||
|
||||
NEEDS_CONTEXTVARS = pytest.mark.skipif(
|
||||
sys.version_info < (3, 11),
|
||||
@@ -1260,7 +1260,7 @@ def test_pregel_validate_rejects_sync_writer_timeout():
|
||||
return value
|
||||
|
||||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||||
Pregel(
|
||||
RawPregel(
|
||||
nodes={
|
||||
"slow": PregelNode(
|
||||
channels="input",
|
||||
@@ -1284,7 +1284,7 @@ def test_pregel_validate_rejects_wrapped_sync_runnable_lambda_timeout():
|
||||
return value + 1
|
||||
|
||||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||||
Pregel(
|
||||
RawPregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
@@ -1307,7 +1307,7 @@ def test_pregel_validate_accepts_wrapped_async_runnable_lambda_timeout():
|
||||
async def slow(value: int) -> int:
|
||||
return value + 1
|
||||
|
||||
Pregel(
|
||||
RawPregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
@@ -1334,7 +1334,7 @@ def test_pregel_validate_rejects_parallel_sync_branch_timeout():
|
||||
return value + 1
|
||||
|
||||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||||
Pregel(
|
||||
RawPregel(
|
||||
nodes={
|
||||
"parallel": (
|
||||
NodeBuilder()
|
||||
@@ -1363,7 +1363,7 @@ def test_pregel_validate_rejects_sync_node_timeout():
|
||||
return value + 1
|
||||
|
||||
with pytest.raises(ValueError, match="only supported for async nodes"):
|
||||
Pregel(
|
||||
RawPregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
@@ -1388,7 +1388,7 @@ async def test_pregel_validate_accepts_async_runnable_lambda_timeout():
|
||||
await asyncio.sleep(0.2)
|
||||
return value + 1
|
||||
|
||||
graph = Pregel(
|
||||
graph = RawPregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
@@ -1418,7 +1418,7 @@ async def test_pregel_validate_accepts_runnable_callable_with_sync_and_async_tim
|
||||
await asyncio.sleep(0.2)
|
||||
return value + 1
|
||||
|
||||
graph = Pregel(
|
||||
graph = RawPregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
@@ -1632,7 +1632,7 @@ async def test_node_builder_timeout_e2e():
|
||||
await asyncio.sleep(0.2)
|
||||
return value + 1
|
||||
|
||||
graph = Pregel(
|
||||
graph = RawPregel(
|
||||
nodes={
|
||||
"slow": (
|
||||
NodeBuilder()
|
||||
|
||||
Reference in new issue
Block a user