Compare commits

..
Author SHA1 Message Date
Igor Soarez d9f65ef5ec feat(langgraph)!: remove Pregel construction and NodeBuilder from the public API
Building a graph by constructing `Pregel(...)` with `NodeBuilder` nodes
was never meant to be public. Graphs are built with `StateGraph` or
`@entrypoint`.

- Delete `NodeBuilder`; `langgraph.pregel` now exports only `Pregel`.
- `Pregel(...)` raises `TypeError` pointing to `StateGraph` and
  `@entrypoint`. `Pregel` stays public as the type that `compile()` and
  `@entrypoint` return, for type annotations and `isinstance` checks.
- `@entrypoint` builds its graph through a private `_EntrypointPregel`
  subclass, the same way `StateGraph` uses `CompiledStateGraph`.
- Rewrite the `Pregel` docstring to describe the compiled graph type and
  drop the hand-built graph examples.
- Engine tests that build graphs by hand now use a test-only copy of
  `NodeBuilder` and a `RawPregel` subclass in `tests/raw_pregel.py`.

BREAKING CHANGE: `from langgraph.pregel import NodeBuilder` and calling
`Pregel(...)` directly no longer work. Build graphs with `StateGraph` or
`@entrypoint`.
2026-10-27 14:43:20 +02:00
25 changed files with 446 additions and 1107 deletions

No files matched your search

+6 -1
View File
@@ -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,
+11 -8
View File
@@ -22,7 +22,7 @@ from langchain_core.messages import (
)
from typing_extensions import TypedDict, deprecated
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND
from langgraph._internal._constants import CONF, CONFIG_KEY_SEND, NS_SEP
from langgraph.graph.state import StateGraph
from langgraph.pregel._messages import (
_messages_delta_reducer as _messages_delta_reducer,
@@ -343,7 +343,7 @@ def push_message(
)
from langgraph.config import get_config
from langgraph.pregel._messages import StreamMessagesHandler, graph_namespace
from langgraph.pregel._messages import StreamMessagesHandler
config = get_config()
message = next(x for x in convert_to_messages([message]))
@@ -359,12 +359,15 @@ def push_message(
):
handlers = config["callbacks"]
metadata = config["metadata"]
for handler in handlers:
if isinstance(handler, StreamMessagesHandler):
ns = graph_namespace(metadata)
if handler.subgraphs or not ns or ns == handler.parent_ns:
handler._emit((ns, metadata), message, dedupe=False)
if stream_handler := next(
(x for x in handlers if isinstance(x, StreamMessagesHandler)), None
):
metadata = config["metadata"]
message_meta = (
tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),
metadata,
)
stream_handler._emit(message_meta, message, dedupe=False)
if state_key:
config[CONF][CONFIG_KEY_SEND]([(state_key, message)])
+2 -2
View File
@@ -1,3 +1,3 @@
from langgraph.pregel.main import NodeBuilder, Pregel
from langgraph.pregel.main import Pregel
__all__ = ("Pregel", "NodeBuilder")
__all__ = ("Pregel",)
+8 -21
View File
@@ -23,7 +23,7 @@ from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResul
from pydantic import BaseModel
from langgraph._internal._config import filter_to_user_tags
from langgraph._internal._constants import NS_SEP, PUSH
from langgraph._internal._constants import NS_SEP
from langgraph.channels.binop import _get_overwrite
from langgraph.channels.delta import DeltaChannel
from langgraph.constants import TAG_HIDDEN, TAG_NOSTREAM
@@ -55,23 +55,6 @@ def _state_values(obj: Any) -> Sequence[Any]:
return ()
def graph_namespace(metadata: Mapping[str, Any]) -> tuple[str, ...]:
"""Namespace of the graph a task runs in, from the task's run metadata.
Drops the task's own segment, plus one per enclosing `@task` call: a call
runs in its caller's graph but nests its segment under the caller's. A
call's `langgraph_path` is `(PUSH, <caller's path>, <index>, True)`, while
a `Send` task's ends in `False`.
"""
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))
path = metadata.get("langgraph_path", ())
drop = 1
while len(path) == 4 and path[0] == PUSH and path[-1] is True:
drop += 1
path = path[1]
return ns[:-drop]
class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
"""A callback handler that implements stream_mode=messages.
@@ -165,7 +148,9 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
**kwargs: Any,
) -> Any:
if metadata and (not tags or (TAG_NOSTREAM not in tags)):
ns = graph_namespace(metadata)
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
:-1
]
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
return
if (filtered_tags := filter_to_user_tags(tags)) is not None:
@@ -228,8 +213,10 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
and kwargs.get("name") == metadata.get("langgraph_node")
and (not tags or TAG_HIDDEN not in tags)
):
ns = graph_namespace(metadata)
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[
:-1
]
if not self.subgraphs and len(ns) > 0:
return
self.metadata[run_id] = (ns, metadata)
for value in _state_values(inputs):
+12 -7
View File
@@ -7,9 +7,9 @@ from uuid import UUID
from langchain_core.callbacks import BaseCallbackHandler
from langgraph._internal._constants import NS_SEP
from langgraph.constants import TAG_NOSTREAM
from langgraph.errors import GraphDrained, GraphInterrupt
from langgraph.pregel._messages import graph_namespace
from langgraph.pregel.protocol import StreamChunk
try:
@@ -92,9 +92,12 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
) -> tuple[str, ...] | None:
"""Resolve the namespace this tool call should emit at, or `None` to skip.
Mirrors `StreamMessagesHandler.on_chat_model_start`: the namespace
is that of the graph the tool runs in (`graph_namespace`). Returns
`None` when the call should be silently suppressed:
Mirrors `StreamMessagesHandler.on_chat_model_start`'s namespace
derivation: parses `langgraph_checkpoint_ns` (which ends with
the `node_name:task_id` of the calling node), drops that
trailing segment, and returns the containing subgraph's own
namespace. Returns `None` when the call should be silently
suppressed:
- `metadata` is missing — handler is attached to a context
without Pregel routing info.
@@ -107,9 +110,11 @@ class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
return None
if tags and TAG_NOSTREAM in tags:
return None
ns = (
graph_namespace(metadata) if metadata.get("langgraph_checkpoint_ns") else ()
)
nskey = metadata.get("langgraph_checkpoint_ns")
if not nskey:
ns: tuple[str, ...] = ()
else:
ns = tuple(cast(str, nskey).split(NS_SEP))[:-1]
if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns:
return None
return ns
+1 -3
View File
@@ -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:
+57 -482
View File
@@ -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(
@@ -2880,18 +2466,6 @@ class Pregel(
for h in callback_manager.inheritable_handlers
if not isinstance(h, StreamMessagesHandlerV2)
]
if CONFIG_KEY_TASK_ID not in config[CONF]:
# Not run as a task (a child invoked with its own thread, say),
# this graph's namespace starts at the root, so a caller's
# handler would take this graph's events for its own graph's.
callback_manager.handlers = [
h for h in callback_manager.handlers if not _is_stream_handler(h)
]
callback_manager.inheritable_handlers = [
h
for h in callback_manager.inheritable_handlers
if not _is_stream_handler(h)
]
if "ls_integration" not in callback_manager.metadata:
callback_manager.add_metadata({"ls_integration": "langgraph"})
run_manager = callback_manager.on_chain_start(
@@ -2910,9 +2484,9 @@ class Pregel(
)
# set up subgraph checkpointing
config = self._own_checkpoint_config(config)
stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
# set up messages stream mode
if "messages" in stream_modes:
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
use_stream_messages_v2 = bool(
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
)
@@ -2921,12 +2495,24 @@ class Pregel(
if use_stream_messages_v2
else StreamMessagesHandler
)
stream_handlers.append(messages_handler_cls(stream.put, subgraphs))
run_manager.inheritable_handlers.append(
messages_handler_cls(
stream.put,
subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
)
)
# set up tools stream mode
if "tools" in stream_modes:
stream_handlers.append(StreamToolCallHandler(stream.put, subgraphs))
run_manager.inheritable_handlers.extend(stream_handlers)
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
run_manager.inheritable_handlers.append(
StreamToolCallHandler(
stream.put,
subgraphs,
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
)
)
# set up custom stream mode
if "custom" in stream_modes:
@@ -3009,10 +2595,6 @@ class Pregel(
cache_policy=self.cache_policy,
has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
) as loop:
# Known only now: the loop numbers the namespace of every
# subgraph a task runs after its first.
for handler in stream_handlers:
handler.parent_ns = loop.checkpoint_ns
emit_graph_lifecycle_events(loop)
# create runner
runner = PregelRunner(
@@ -3295,18 +2877,6 @@ class Pregel(
for h in callback_manager.inheritable_handlers
if not isinstance(h, StreamMessagesHandlerV2)
]
if CONFIG_KEY_TASK_ID not in config[CONF]:
# Not run as a task (a child invoked with its own thread, say),
# this graph's namespace starts at the root, so a caller's
# handler would take this graph's events for its own graph's.
callback_manager.handlers = [
h for h in callback_manager.handlers if not _is_stream_handler(h)
]
callback_manager.inheritable_handlers = [
h
for h in callback_manager.inheritable_handlers
if not _is_stream_handler(h)
]
if "ls_integration" not in callback_manager.metadata:
callback_manager.add_metadata({"ls_integration": "langgraph"})
run_manager = await callback_manager.on_chain_start(
@@ -3341,9 +2911,10 @@ class Pregel(
)
# set up subgraph checkpointing
config = self._own_checkpoint_config(config)
stream_handlers: list[StreamMessagesHandler | StreamToolCallHandler] = []
# set up messages stream mode
if "messages" in stream_modes:
# namespace can be None in a root level graph?
ns_ = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
use_stream_messages_v2 = bool(
version == "v2" and config[CONF].get(CONFIG_KEY_STREAM_MESSAGES_V2)
)
@@ -3352,12 +2923,24 @@ class Pregel(
if use_stream_messages_v2
else StreamMessagesHandler
)
stream_handlers.append(messages_handler_cls(stream_put, subgraphs))
run_manager.inheritable_handlers.append(
messages_handler_cls(
stream_put,
subgraphs,
parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None,
)
)
# set up tools stream mode
if "tools" in stream_modes:
stream_handlers.append(StreamToolCallHandler(stream_put, subgraphs))
run_manager.inheritable_handlers.extend(stream_handlers)
ns_tools = cast(str | None, config[CONF].get(CONFIG_KEY_CHECKPOINT_NS))
run_manager.inheritable_handlers.append(
StreamToolCallHandler(
stream_put,
subgraphs,
parent_ns=tuple(ns_tools.split(NS_SEP)) if ns_tools else None,
)
)
# set up custom stream mode
def stream_writer(c: Any) -> None:
@@ -3466,10 +3049,6 @@ class Pregel(
cache_policy=self.cache_policy,
has_graph_lifecycle_callbacks=bool(graph_callback_manager.handlers),
) as loop:
# Known only now: the loop numbers the namespace of every
# subgraph a task runs after its first.
for handler in stream_handlers:
handler.parent_ns = loop.checkpoint_ns
await aemit_graph_lifecycle_events(loop)
# create runner
runner = PregelRunner(
@@ -4467,10 +4046,6 @@ def _coerce_checkpoint_values(payload: Any, mapper: Callable[[Any], Any]) -> Non
payload["values"] = mapper(payload["values"])
def _is_stream_handler(handler: object) -> bool:
return isinstance(handler, (StreamMessagesHandler, StreamToolCallHandler))
def _resolve_parent_ns(
graph_config: RunnableConfig | None, call_config: RunnableConfig | None
) -> tuple[str, ...]:
-10
View File
@@ -847,11 +847,6 @@ class RemoteGraph(PregelProtocol):
# filter for what was actually requested
if mode not in requested:
continue
# subgraph events may have been requested only for the parent stream.
# messages are exempt: the server dedupes them across namespaces, so
# dropping them here would also drop the root-level message
if not subgraphs and NS_SEP in chunk.event and mode != "messages":
continue
if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data))
@@ -1002,11 +997,6 @@ class RemoteGraph(PregelProtocol):
# filter for what was actually requested
if mode not in requested:
continue
# subgraph events may have been requested only for the parent stream.
# messages are exempt: the server dedupes them across namespaces, so
# dropping them here would also drop the root-level message
if not subgraphs and NS_SEP in chunk.event and mode != "messages":
continue
if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data))
+222
View File
@@ -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),
+2 -2
View File
@@ -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,
},
+2 -2
View File
@@ -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),
+25 -60
View File
@@ -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,
},
@@ -8950,39 +8948,6 @@ def test_subgraph_streaming_sync() -> None:
assert result["num_chunks"] == 9
def test_subgraph_streaming_node_output_messages_sync() -> None:
"""Messages returned by a node (not streamed by an LLM) are emitted when a
graph is explicitly streamed with `stream_mode="messages"` inside a node."""
def respond(state: MessagesState) -> dict:
return {"messages": [AIMessage("hi there", id="ai1")]}
subgraph = (
StateGraph(MessagesState)
.add_node("respond", respond)
.add_edge(START, "respond")
.compile()
)
msgs = {"messages": [("user", "hello")]}
top_level = [msg for msg, _ in subgraph.stream(msgs, stream_mode="messages")]
def parent_node(state: MessagesState, config: RunnableConfig) -> dict:
nonlocal in_node
in_node = [
msg for msg, _ in subgraph.stream(msgs, config, stream_mode="messages")
]
return {}
in_node: list = []
StateGraph(MessagesState).add_node("parent", parent_node).add_edge(
START, "parent"
).compile().invoke(msgs)
assert [m.content for m in top_level] == ["hi there"]
assert [m.content for m in in_node] == ["hi there"]
def test_get_graph_nonterminal_last_step_source(snapshot: SnapshotAssertion) -> None:
class State(TypedDict):
messages: list[str]
@@ -9540,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"),
@@ -9585,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")
@@ -9620,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"),
+16 -54
View File
@@ -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),
@@ -9503,44 +9503,6 @@ async def test_subgraph_streaming_async() -> None:
assert result["num_chunks"] == 9
async def test_subgraph_streaming_node_output_messages_async() -> None:
"""Messages returned by a node (not streamed by an LLM) are emitted when a
graph is explicitly streamed with `stream_mode="messages"` inside a node."""
async def respond(state: MessagesState) -> dict:
return {"messages": [AIMessage("hi there", id="ai1")]}
subgraph = (
StateGraph(MessagesState)
.add_node("respond", respond)
.add_edge(START, "respond")
.compile()
)
msgs = {"messages": [("user", "hello")]}
top_level = [msg async for msg, _ in subgraph.astream(msgs, stream_mode="messages")]
async def parent_node(state: MessagesState, config: RunnableConfig) -> dict:
nonlocal in_node
in_node = [
msg
async for msg, _ in subgraph.astream(msgs, config, stream_mode="messages")
]
return {}
in_node: list = []
await (
StateGraph(MessagesState)
.add_node("parent", parent_node)
.add_edge(START, "parent")
.compile()
.ainvoke(msgs)
)
assert [m.content for m in top_level] == ["hi there"]
assert [m.content for m in in_node] == ["hi there"]
@NEEDS_CONTEXTVARS
async def test_null_resume_disallowed_with_multiple_interrupts(
async_checkpointer: BaseCheckpointSaver,
@@ -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
-83
View File
@@ -16,11 +16,9 @@ from langgraph_sdk.schema import StreamPart
from pydantic import BaseModel
from typing_extensions import TypedDict
from langgraph._internal._constants import CONFIG_KEY_STREAM
from langgraph.errors import GraphInterrupt
from langgraph.graph import END, START, MessagesState, StateGraph, add_messages
from langgraph.pregel import Pregel
from langgraph.pregel.protocol import StreamProtocol
from langgraph.pregel.remote import RemoteGraph
from langgraph.types import Interrupt, StateSnapshot
from tests.any_str import AnyStr
@@ -865,87 +863,6 @@ async def test_astream_subgraphs_keeps_caller_namespace():
assert stream_parts == _CALLER_PREFIXED
# Subgraph events requested only on behalf of a streaming parent graph
_PARENT_STREAMED_PARTS = [
StreamPart(event="updates|sub:abc", data={"inner": {"chunk": "inner"}}),
StreamPart(event="messages|sub:abc", data=[{"content": "token"}, {}]),
StreamPart(event="updates", data={"sub": {"chunk": "root"}}),
]
def _parent_streamed_config(parent_parts: list) -> RunnableConfig:
return {
"configurable": {
"thread_id": "thread_1",
"checkpoint_ns": "outer:123",
CONFIG_KEY_STREAM: StreamProtocol(
parent_parts.append, {"updates", "messages"}
),
}
}
def test_stream_without_subgraphs_skips_parent_requested_subgraph_events():
mock_sync_client = MagicMock()
mock_sync_client.runs.stream.return_value = _PARENT_STREAMED_PARTS
remote_pregel = RemoteGraph("test_graph_id", sync_client=mock_sync_client)
for stream_mode, expected in [
("updates", [{"sub": {"chunk": "root"}}]),
# messages are not filtered: the server dedupes them across namespaces
("messages", [({"content": "token"}, {})]),
]:
parent_parts: list = []
stream_parts = list(
remote_pregel.stream(
{"input": "data"},
_parent_streamed_config(parent_parts),
stream_mode=stream_mode,
)
)
assert stream_parts == expected
# subgraph events are still requested and forwarded to the parent
assert mock_sync_client.runs.stream.call_args.kwargs["stream_subgraphs"]
assert [ns for ns, _, _ in parent_parts] == [
("outer:123", "sub:abc"),
("outer:123", "sub:abc"),
("outer:123",),
]
async def test_astream_without_subgraphs_skips_parent_requested_subgraph_events():
mock_async_client = MagicMock()
async_iter = MagicMock()
async_iter.__aiter__.return_value = _PARENT_STREAMED_PARTS
mock_async_client.runs.stream.return_value = async_iter
remote_pregel = RemoteGraph("test_graph_id", client=mock_async_client)
for stream_mode, expected in [
("updates", [{"sub": {"chunk": "root"}}]),
# messages are not filtered: the server dedupes them across namespaces
("messages", [({"content": "token"}, {})]),
]:
parent_parts: list = []
stream_parts = [
part
async for part in remote_pregel.astream(
{"input": "data"},
_parent_streamed_config(parent_parts),
stream_mode=stream_mode,
)
]
assert stream_parts == expected
# subgraph events are still requested and forwarded to the parent
assert mock_async_client.runs.stream.call_args.kwargs["stream_subgraphs"]
assert [ns for ns, _, _ in parent_parts] == [
("outer:123", "sub:abc"),
("outer:123", "sub:abc"),
("outer:123",),
]
def test_invoke():
# set up test
mock_sync_client = MagicMock()
+9 -9
View File
@@ -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()
@@ -1,308 +0,0 @@
import sys
from typing import Any
import pytest
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
from langchain_core.messages import AIMessage
from langchain_core.tools import tool
from langgraph.func import entrypoint, task
from langgraph.graph import START, MessagesState, StateGraph
from langgraph.graph.message import push_message
from langgraph.graph.state import CompiledStateGraph
pytestmark = pytest.mark.anyio
NEEDS_CONTEXTVARS = pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
)
@tool
def echo(text: str) -> str:
"""Echo the text."""
return text
def _one_node_graph(node: Any) -> CompiledStateGraph:
return (
StateGraph(MessagesState)
.add_node(node)
.add_edge(START, node.__name__)
.compile()
)
def _chat_and_echo(state: MessagesState) -> dict:
echo.invoke({"text": "done"})
return {"messages": [GenericFakeChatModel(messages=iter(["one two"])).invoke("hi")]}
async def _achat_and_echo(state: MessagesState) -> dict:
await echo.ainvoke({"text": "done"})
model = GenericFakeChatModel(messages=iter(["one two"]))
return {"messages": [await model.ainvoke("hi")]}
def _segment_names(ns: tuple[str, ...]) -> tuple[str, ...]:
return tuple(segment.split(":")[0] for segment in ns)
def _summary(part: dict) -> tuple[tuple[str, ...], str]:
if part["type"] == "messages":
return _segment_names(part["ns"]), part["data"][0].content
return _segment_names(part["ns"]), part["data"]["event"]
_CHAT_AND_ECHO_EVENTS = [
((), "tool-started"),
((), "tool-finished"),
((), "one"),
((), " "),
((), "two"),
]
_CHILD_CONFIGS = pytest.mark.parametrize(
"child_config",
[
None,
{"configurable": {"thread_id": "child"}},
{"configurable": {"thread_id": "parent"}},
],
ids=["no_config", "own_thread", "caller_thread"],
)
_TASK_CALLERS = pytest.mark.parametrize(
"nested", [False, True], ids=["from_entrypoint", "from_task_in_task"]
)
@_TASK_CALLERS
@pytest.mark.parametrize("subgraphs", [False, True])
def test_task_events_stream_under_the_graph_namespace(
nested: bool, subgraphs: bool
) -> None:
@task
def chat_and_echo() -> None:
_chat_and_echo({"messages": []})
@task
def outer() -> None:
chat_and_echo().result()
@entrypoint()
def workflow(inputs: dict) -> None:
(outer() if nested else chat_and_echo()).result()
parts = workflow.stream(
{}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
)
assert [_summary(part) for part in parts] == _CHAT_AND_ECHO_EVENTS
@NEEDS_CONTEXTVARS
@_TASK_CALLERS
@pytest.mark.parametrize("subgraphs", [False, True])
async def test_task_events_stream_under_the_graph_namespace_async(
nested: bool, subgraphs: bool
) -> None:
@task
async def chat_and_echo() -> None:
await _achat_and_echo({"messages": []})
@task
async def outer() -> None:
await chat_and_echo()
@entrypoint()
async def workflow(inputs: dict) -> None:
await (outer() if nested else chat_and_echo())
parts = workflow.astream(
{}, stream_mode=["messages", "tools"], subgraphs=subgraphs, version="v2"
)
assert [_summary(part) async for part in parts] == _CHAT_AND_ECHO_EVENTS
def test_task_events_stream_under_the_namespace_of_their_subgraph() -> None:
@task
def chat_and_echo() -> None:
_chat_and_echo({"messages": []})
def node(state: MessagesState) -> dict:
chat_and_echo().result()
return {}
sub = _one_node_graph(node)
parent = StateGraph(MessagesState).add_node("sub", sub).add_edge(START, "sub")
parts = parent.compile().stream(
{"messages": []},
stream_mode=["messages", "tools"],
subgraphs=True,
version="v2",
)
assert [_summary(part) for part in parts] == [
(("sub",), event) for _, event in _CHAT_AND_ECHO_EVENTS
]
@_CHILD_CONFIGS
def test_child_graph_streams_only_to_the_node_streaming_it(
child_config: dict | None,
) -> None:
child = _one_node_graph(_chat_and_echo)
received = []
def node(state: MessagesState) -> dict:
received.extend(
child.stream(
{"messages": []}, child_config, stream_mode=["messages", "tools"]
)
)
return {}
parts = _one_node_graph(node).stream(
{"messages": []},
{"configurable": {"thread_id": "parent"}},
stream_mode=["messages", "tools"],
)
assert list(parts) == []
assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
@NEEDS_CONTEXTVARS
@_CHILD_CONFIGS
async def test_child_graph_streams_only_to_the_node_streaming_it_async(
child_config: dict | None,
) -> None:
child = _one_node_graph(_achat_and_echo)
received = []
async def node(state: MessagesState) -> dict:
stream = child.astream(
{"messages": []}, child_config, stream_mode=["messages", "tools"]
)
received.extend([part async for part in stream])
return {}
parts = _one_node_graph(node).astream(
{"messages": []},
{"configurable": {"thread_id": "parent"}},
stream_mode=["messages", "tools"],
)
assert [part async for part in parts] == []
assert len(received) == len(_CHAT_AND_ECHO_EVENTS)
@pytest.mark.parametrize(
"child_config",
[
{"configurable": {"thread_id": "child"}},
{"configurable": {"thread_id": "parent"}},
],
ids=["own_thread", "caller_thread"],
)
def test_child_graph_on_its_own_thread_skips_a_caller_streaming_subgraphs(
child_config: dict,
) -> None:
child = _one_node_graph(_chat_and_echo)
def node(state: MessagesState) -> dict:
child.invoke({"messages": []}, child_config)
return {}
parts = _one_node_graph(node).stream(
{"messages": []},
{"configurable": {"thread_id": "parent"}},
stream_mode=["messages", "tools"],
subgraphs=True,
)
assert list(parts) == []
def test_node_gets_the_events_of_every_subgraph_it_streams() -> None:
sub = _one_node_graph(_chat_and_echo)
counts = []
def node(state: MessagesState) -> dict:
for _ in range(2):
stream = sub.stream({"messages": []}, stream_mode=["messages", "tools"])
counts.append(len(list(stream)))
return {}
_one_node_graph(node).invoke({"messages": []})
assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
@NEEDS_CONTEXTVARS
async def test_node_gets_the_events_of_every_subgraph_it_streams_async() -> None:
sub = _one_node_graph(_achat_and_echo)
counts = []
async def node(state: MessagesState) -> dict:
for _ in range(2):
stream = sub.astream({"messages": []}, stream_mode=["messages", "tools"])
counts.append(len([part async for part in stream]))
return {}
await _one_node_graph(node).ainvoke({"messages": []})
assert counts == [len(_CHAT_AND_ECHO_EVENTS)] * 2
def _push(state: MessagesState) -> dict:
push_message(AIMessage(id="pushed", content="hi"), state_key=None)
return {}
def _namespaces(graph: CompiledStateGraph, subgraphs: bool) -> list[tuple[str, ...]]:
parts = graph.stream(
{"messages": []}, stream_mode="messages", subgraphs=subgraphs, version="v2"
)
return [_segment_names(part["ns"]) for part in parts]
def _streams_pushing_graph(received: list) -> Any:
pushing = _one_node_graph(_push)
def node(state: MessagesState) -> dict:
for message, _ in pushing.stream({"messages": []}, stream_mode="messages"):
received.append(message.id)
return {}
return node
def test_pushed_message_streams_under_the_namespace_of_its_graph() -> None:
pushing = _one_node_graph(_push)
parent = StateGraph(MessagesState).add_node("sub", pushing).add_edge(START, "sub")
assert _namespaces(pushing, subgraphs=True) == [()]
assert _namespaces(parent.compile(), subgraphs=False) == []
assert _namespaces(parent.compile(), subgraphs=True) == [("sub",)]
def test_pushed_message_streams_to_the_node_that_streams_its_graph() -> None:
received: list[str] = []
graph = _one_node_graph(_streams_pushing_graph(received))
assert _namespaces(graph, subgraphs=False) == []
assert received == ["pushed"]
def test_pushed_message_streams_to_every_stream_that_takes_it() -> None:
received: list[str] = []
graph = _one_node_graph(_streams_pushing_graph(received))
assert _namespaces(graph, subgraphs=True) == [("node",)]
assert received == ["pushed"]
+1 -1
View File
@@ -811,7 +811,7 @@ class _On:
def __call__(
self,
*,
resources: str | Sequence[str] | None = None,
resources: str | Sequence[str],
actions: str | Sequence[str] | None = None,
) -> Callable[[AHO], AHO]: ...
-41
View File
@@ -1,49 +1,8 @@
import ast
import inspect
import pytest
from langgraph_sdk import Auth
@pytest.mark.parametrize(
"kwargs",
[{}, {"actions": "read"}, {"resources": None}],
)
def test_global_handler_factory_accepts_optional_resources(kwargs) -> None:
auth = Auth()
tree = ast.parse(inspect.getsource(type(auth.on)))
cls = tree.body[0]
assert isinstance(cls, ast.ClassDef)
factory_overload = next(
node
for node in cls.body
if isinstance(node, ast.FunctionDef)
and node.name == "__call__"
and node.decorator_list
and any(arg.arg == "resources" for arg in node.args.kwonlyargs)
)
resources = next(
arg for arg in factory_overload.args.kwonlyargs if arg.arg == "resources"
)
assert resources.annotation is not None
assert ast.unparse(resources.annotation) == "str | Sequence[str] | None"
default = factory_overload.args.kw_defaults[0]
assert isinstance(default, ast.Constant) and default.value is None
@auth.on(**kwargs)
async def handler(ctx, value):
del ctx, value
return None
if kwargs.get("actions"):
assert auth._handlers == {("*", "read"): [handler]}
assert auth._global_handlers == []
else:
assert auth._global_handlers == [handler]
assert auth._handlers == {}
def test_handler_multiple_resources_and_actions() -> None:
auth = Auth()