mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-08 18:57:52 +02:00
Add subgraph lifecycle streaming with scoped per-subagent projections
StreamingHandler now yields SubgraphRunStream handles for each nested Pregel as it spawns. Each handle is a BaseRunStream wrapping a mini-mux built from the same transformer factories as the root — so sub.values, sub.messages, sub.subgraphs are populated by standard ValuesTransformer / MessagesTransformer / SubgraphTransformer instances at that subagent's scope. No routing or ChatModelStream assembly duplicated across transformers. Key pieces: - StreamLifecycleHandler (pregel/_lifecycle.py): callback handler attached at pregel stream / astream sites when "lifecycle" is in stream_modes. Emits started / running / completed / failed / interrupted events per nested Pregel via metadata-based detection (langgraph_checkpoint_ns + name != langgraph_node, excluding __start__/__end__ sentinels). Carries trigger_call_id from the parent task id. Root subgraph terminal state is emitted eagerly at __init__ and by SubgraphTransformer.finalize / fail. - StreamMux.make_child(scope) + factory-based construction. Mux takes a factory list; make_child produces a mini-mux at a new scope with fresh instances. bind_pump / bind_apump cascade through children so any subagent cursor drives the root pump. - StreamTransformer.scope (base attribute) + scope_exact class flag. Mux skips process() for out-of-scope events when scope_exact=True (default), so user transformers get scope-filtered events with no boilerplate. SubgraphTransformer opts out to receive cross-scope events for forwarding. - BaseRunStream shared base for GraphRunStream, AsyncGraphRunStream, and SubgraphRunStream. Provides extensions, native attrs, raw __iter__ / __aiter__, interleave. Root subclasses own the graph iterator and pump; SubgraphRunStream adds lifecycle metadata (path, status, error, checkpoint, trigger_call_id, graph_name). - SubgraphTransformer (stream/transformers.py) is now a thin discovery + forwarding dispatcher. On lifecycle.started at its scope + 1, it creates a SubgraphRunStream via parent_mux.make_child(ns). It forwards each event matching a direct-child's path into that child's mini-mux. Terminal lifecycle closes the mini-mux. - StreamMode literal extended with "lifecycle". pregel/remote.py filters it out of the SDK stream-mode list since wire-protocol support isn't landed yet. Tests: 16 new in test_stream_subgraph_transformer.py (unit + sync/async end-to-end including error + grandchild); existing Values/Messages namespace filter tests migrated through mux.push to reflect the new scope_exact contract.
This commit is contained in:
@@ -0,0 +1,245 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Callable, Iterator
|
||||
from typing import Any, TypeVar, cast
|
||||
from uuid import UUID
|
||||
|
||||
from langchain_core.callbacks import BaseCallbackHandler
|
||||
|
||||
from langgraph._internal._constants import NS_SEP
|
||||
from langgraph.errors import GraphInterrupt
|
||||
from langgraph.pregel.protocol import StreamChunk
|
||||
|
||||
try:
|
||||
from langchain_core.tracers._streaming import _StreamingCallbackHandler
|
||||
except ImportError:
|
||||
_StreamingCallbackHandler = object # type: ignore[assignment,misc]
|
||||
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
_LANGGRAPH_SENTINEL_NODES = frozenset({"__start__", "__end__"})
|
||||
|
||||
|
||||
def _is_nested_pregel_start(name: str | None, metadata: dict[str, Any] | None) -> bool:
|
||||
"""Recognize a nested `Pregel` invocation from its `on_chain_start` metadata.
|
||||
|
||||
When a compiled graph is added as a node, pregel fires two
|
||||
`on_chain_start` callbacks at that task: first for the node chain
|
||||
(whose `name` matches `metadata["langgraph_node"]`) and second for
|
||||
the inner `Pregel` chain (whose `name` is the graph's `name`, not
|
||||
the node name). Both share the same `langgraph_checkpoint_ns`.
|
||||
|
||||
A nested `Pregel` start is therefore identified by having a
|
||||
`langgraph_checkpoint_ns` AND a `name` that does NOT match the
|
||||
owning task's `langgraph_node`. Regular node chains are skipped;
|
||||
the root `Pregel` (which has no `langgraph_node` metadata) isn't
|
||||
observed by this handler because the root's start fires before the
|
||||
handler is attached.
|
||||
|
||||
Metadata-based detection is used because `on_chain_start`'s
|
||||
`serialized` argument is `None` for compiled graphs in this
|
||||
version of langchain-core, so class-based detection via
|
||||
`serialized["id"]` isn't available.
|
||||
|
||||
Sentinel nodes (`__start__` / `__end__`) are excluded: conditional
|
||||
edges from `START` fire an `on_chain_start` with `lg_node=__start__`
|
||||
and the router function's name as `name`, which would otherwise
|
||||
match the discriminator without representing an actual nested
|
||||
`Pregel`.
|
||||
"""
|
||||
if not metadata:
|
||||
return False
|
||||
if not metadata.get("langgraph_checkpoint_ns"):
|
||||
return False
|
||||
lg_node = metadata.get("langgraph_node")
|
||||
if lg_node is None or lg_node in _LANGGRAPH_SENTINEL_NODES:
|
||||
return False
|
||||
return name != lg_node
|
||||
|
||||
|
||||
class StreamLifecycleHandler(BaseCallbackHandler, _StreamingCallbackHandler):
|
||||
"""Callback handler that emits subgraph lifecycle events on the stream.
|
||||
|
||||
Pushes `LifecycleData`-shaped payloads onto the pregel stream under
|
||||
the `"lifecycle"` mode, keyed by the subgraph's namespace tuple.
|
||||
Drives the `started` → `running` → `completed` / `failed` /
|
||||
`interrupted` state machine.
|
||||
|
||||
The handler is attached to `run_manager.inheritable_handlers` inside
|
||||
a `Pregel.stream` / `astream` call, so it sees callbacks for every
|
||||
descendant chain (nodes, nested `Pregel` subgraphs) but *not* for
|
||||
the root `Pregel` whose start event has already fired. The root's
|
||||
`started` event is emitted eagerly at construction; its terminal
|
||||
state is emitted by `SubgraphTransformer.finalize` / `fail`.
|
||||
|
||||
`run_inline = True` keeps event ordering deterministic.
|
||||
"""
|
||||
|
||||
run_inline = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
stream: Callable[[StreamChunk], None],
|
||||
*,
|
||||
root_graph_name: str | None = None,
|
||||
) -> None:
|
||||
"""Initialize the handler and emit the root graph's `started` event.
|
||||
|
||||
Args:
|
||||
stream: Callable that accepts a `StreamChunk` tuple
|
||||
`(namespace, mode, payload)` and enqueues it.
|
||||
root_graph_name: The root `Pregel` instance's `name`, emitted
|
||||
with the root's `started` lifecycle payload.
|
||||
"""
|
||||
self.stream = stream
|
||||
# Namespaces awaiting the started→running transition.
|
||||
self._pending_running: set[tuple[str, ...]] = set()
|
||||
# run_id → subgraph namespace; populated only for Pregel chains.
|
||||
self._run_to_ns: dict[UUID, tuple[str, ...]] = {}
|
||||
|
||||
root_payload: dict[str, Any] = {"event": "started"}
|
||||
if root_graph_name is not None:
|
||||
root_payload["graph_name"] = root_graph_name
|
||||
self.stream(((), "lifecycle", root_payload))
|
||||
self._pending_running.add(())
|
||||
|
||||
@staticmethod
|
||||
def _subgraph_ns_from_metadata(metadata: dict[str, Any] | None) -> tuple[str, ...]:
|
||||
"""Return the running subgraph's own namespace from task metadata.
|
||||
|
||||
For a nested `Pregel` invoked as a node, `langgraph_checkpoint_ns`
|
||||
ends at the node segment (no inner task appended yet), so
|
||||
splitting on `NS_SEP` gives the subgraph's own namespace.
|
||||
"""
|
||||
if not metadata:
|
||||
return ()
|
||||
nskey = metadata.get("langgraph_checkpoint_ns")
|
||||
if not nskey:
|
||||
return ()
|
||||
return tuple(cast(str, nskey).split(NS_SEP))
|
||||
|
||||
@staticmethod
|
||||
def _containing_ns_from_metadata(
|
||||
metadata: dict[str, Any] | None,
|
||||
) -> tuple[str, ...]:
|
||||
"""Return the namespace of the subgraph that contains this task.
|
||||
|
||||
For an inner task with `langgraph_checkpoint_ns`
|
||||
`"seg_a|seg_b"`, the containing subgraph is `("seg_a",)`.
|
||||
"""
|
||||
if not metadata:
|
||||
return ()
|
||||
nskey = metadata.get("langgraph_checkpoint_ns")
|
||||
if not nskey:
|
||||
return ()
|
||||
return tuple(cast(str, nskey).split(NS_SEP))[:-1]
|
||||
|
||||
@staticmethod
|
||||
def _trigger_call_id(metadata: dict[str, Any] | None) -> str | None:
|
||||
"""Extract `trigger_call_id` from task metadata if present.
|
||||
|
||||
The task that spawned a nested `Pregel` has its task id encoded
|
||||
in `langgraph_checkpoint_ns`'s last segment as
|
||||
`node_name:task_id`. Returns the `task_id` portion, which
|
||||
parents can correlate with their `tools` / `tasks` events.
|
||||
"""
|
||||
if not metadata:
|
||||
return None
|
||||
nskey = cast(str | None, metadata.get("langgraph_checkpoint_ns"))
|
||||
if not nskey:
|
||||
return None
|
||||
last = nskey.split(NS_SEP)[-1]
|
||||
_, sep, task_id = last.rpartition(":")
|
||||
return task_id if sep else None
|
||||
|
||||
def _emit(self, ns: tuple[str, ...], payload: dict[str, Any]) -> None:
|
||||
self.stream((ns, "lifecycle", payload))
|
||||
|
||||
def tap_output_aiter(
|
||||
self, run_id: UUID, output: AsyncIterator[T]
|
||||
) -> AsyncIterator[T]:
|
||||
"""Pass-through — required by the `_StreamingCallbackHandler` protocol.
|
||||
|
||||
Returns the iterator unchanged. A missing implementation lets
|
||||
langchain's default `Protocol` body return `None`, which breaks
|
||||
the `_consume_aiter` code path in `_runnable.py:900`.
|
||||
"""
|
||||
return output
|
||||
|
||||
def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]:
|
||||
"""Pass-through — sync counterpart to `tap_output_aiter`."""
|
||||
return output
|
||||
|
||||
def _fire_running_if_pending(self, ns: tuple[str, ...]) -> None:
|
||||
if ns in self._pending_running:
|
||||
self._pending_running.discard(ns)
|
||||
self._emit(ns, {"event": "running"})
|
||||
|
||||
def on_chain_start(
|
||||
self,
|
||||
serialized: dict[str, Any],
|
||||
inputs: dict[str, Any],
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
tags: list[str] | None = None,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
# Any descendant activity transitions the containing subgraph to running.
|
||||
containing = self._containing_ns_from_metadata(metadata)
|
||||
self._fire_running_if_pending(containing)
|
||||
|
||||
name = cast(str | None, kwargs.get("name"))
|
||||
if not _is_nested_pregel_start(name, metadata):
|
||||
return
|
||||
|
||||
ns = self._subgraph_ns_from_metadata(metadata)
|
||||
if not ns:
|
||||
return
|
||||
|
||||
self._run_to_ns[run_id] = ns
|
||||
payload: dict[str, Any] = {"event": "started"}
|
||||
if name:
|
||||
payload["graph_name"] = name
|
||||
trigger_call_id = self._trigger_call_id(metadata)
|
||||
if trigger_call_id:
|
||||
payload["trigger_call_id"] = trigger_call_id
|
||||
self._emit(ns, payload)
|
||||
self._pending_running.add(ns)
|
||||
|
||||
def on_chain_end(
|
||||
self,
|
||||
response: Any,
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
ns = self._run_to_ns.pop(run_id, None)
|
||||
if ns is None:
|
||||
return
|
||||
# Ensure started→running fired even for empty subgraphs.
|
||||
if ns in self._pending_running:
|
||||
self._pending_running.discard(ns)
|
||||
self._emit(ns, {"event": "running"})
|
||||
self._emit(ns, {"event": "completed"})
|
||||
|
||||
def on_chain_error(
|
||||
self,
|
||||
error: BaseException,
|
||||
*,
|
||||
run_id: UUID,
|
||||
parent_run_id: UUID | None = None,
|
||||
**kwargs: Any,
|
||||
) -> Any:
|
||||
ns = self._run_to_ns.pop(run_id, None)
|
||||
if ns is None:
|
||||
return
|
||||
self._pending_running.discard(ns)
|
||||
if isinstance(error, GraphInterrupt):
|
||||
self._emit(ns, {"event": "interrupted"})
|
||||
else:
|
||||
self._emit(ns, {"event": "failed", "error": str(error)})
|
||||
@@ -130,6 +130,7 @@ from langgraph.pregel._checkpoint import (
|
||||
)
|
||||
from langgraph.pregel._draw import draw_graph
|
||||
from langgraph.pregel._io import map_input, read_channels
|
||||
from langgraph.pregel._lifecycle import StreamLifecycleHandler
|
||||
from langgraph.pregel._loop import (
|
||||
AsyncPregelLoop,
|
||||
SyncPregelLoop,
|
||||
@@ -2643,6 +2644,15 @@ class Pregel(
|
||||
)
|
||||
)
|
||||
|
||||
# set up lifecycle stream mode
|
||||
if "lifecycle" in stream_modes:
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamLifecycleHandler(
|
||||
stream.put,
|
||||
root_graph_name=self.name,
|
||||
)
|
||||
)
|
||||
|
||||
# set up custom stream mode
|
||||
if "custom" in stream_modes:
|
||||
|
||||
@@ -3025,6 +3035,15 @@ class Pregel(
|
||||
)
|
||||
)
|
||||
|
||||
# set up lifecycle stream mode
|
||||
if "lifecycle" in stream_modes:
|
||||
run_manager.inheritable_handlers.append(
|
||||
StreamLifecycleHandler(
|
||||
stream_put,
|
||||
root_graph_name=self.name,
|
||||
)
|
||||
)
|
||||
|
||||
# set up custom stream mode
|
||||
def stream_writer(c: Any) -> None:
|
||||
aioloop.call_soon_threadsafe(
|
||||
|
||||
@@ -653,19 +653,24 @@ class RemoteGraph(PregelProtocol):
|
||||
# coerce to list, or add default stream mode
|
||||
if stream_mode:
|
||||
if isinstance(stream_mode, str):
|
||||
updated_stream_modes.append(stream_mode)
|
||||
if stream_mode != "lifecycle":
|
||||
updated_stream_modes.append(cast(StreamModeSDK, stream_mode))
|
||||
else:
|
||||
req_single = False
|
||||
updated_stream_modes.extend(stream_mode)
|
||||
updated_stream_modes.extend(
|
||||
cast(StreamModeSDK, m) for m in stream_mode if m != "lifecycle"
|
||||
)
|
||||
else:
|
||||
updated_stream_modes.append(default)
|
||||
updated_stream_modes.append(default) # type: ignore[arg-type]
|
||||
requested_stream_modes = updated_stream_modes.copy()
|
||||
# add any from parent graph
|
||||
stream: StreamProtocol | None = (
|
||||
(config or {}).get(CONF, {}).get(CONFIG_KEY_STREAM)
|
||||
)
|
||||
if stream:
|
||||
updated_stream_modes.extend(stream.modes)
|
||||
updated_stream_modes.extend(
|
||||
cast(StreamModeSDK, m) for m in stream.modes if m != "lifecycle"
|
||||
)
|
||||
# map "messages" to "messages-tuple"
|
||||
if "messages" in updated_stream_modes:
|
||||
updated_stream_modes.remove("messages")
|
||||
|
||||
@@ -2,7 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._event_log import EventLog
|
||||
@@ -13,6 +13,17 @@ from langgraph.stream._types import (
|
||||
)
|
||||
from langgraph.stream.stream_channel import StreamChannel
|
||||
|
||||
TransformerFactory = Callable[["tuple[str, ...]"], StreamTransformer]
|
||||
"""Factory that builds a scoped transformer for a mux.
|
||||
|
||||
Called once per `StreamMux` (root or mini-mux) with the mux's scope
|
||||
— typically a subgraph's namespace or `()` for the root. Standard
|
||||
transformer classes (`ValuesTransformer`, `MessagesTransformer`,
|
||||
`SubgraphTransformer`) accept a single positional scope argument, so
|
||||
the class itself is a valid factory. User transformers can close over
|
||||
their config: `lambda scope: MyTransformer(scope, foo=...)`.
|
||||
"""
|
||||
|
||||
|
||||
class StreamMux:
|
||||
"""Central event dispatcher for the streaming infrastructure.
|
||||
@@ -40,28 +51,50 @@ class StreamMux:
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
*,
|
||||
is_async: bool = False,
|
||||
factories: list[TransformerFactory] | None = None,
|
||||
scope: tuple[str, ...] = (),
|
||||
) -> None:
|
||||
"""Initialize the mux and register transformers in order.
|
||||
|
||||
Transformers are fixed at construction time — there is no
|
||||
post-init `register()`. Each transformer's `init()` is called,
|
||||
Callers pass either `transformers` (pre-built instances) or
|
||||
`factories` (callables producing fresh instances per mux). A
|
||||
factory list is preferred — mini-muxes built by `make_child()`
|
||||
inherit the factory list, so transformers propagate naturally
|
||||
into every subgraph's scope. `transformers` is kept for
|
||||
back-compat tests that exercise the mux directly.
|
||||
|
||||
Each transformer's `init()` is called once during registration,
|
||||
projections are merged into `extensions`, `_native` keys are
|
||||
recorded in `native_keys`, and any EventLog / StreamChannel
|
||||
instances are bound and wired.
|
||||
|
||||
Args:
|
||||
transformers: Transformers to register, in dispatch order.
|
||||
`None` or empty gives a mux with no projections.
|
||||
transformers: Already-built transformer instances. Mutually
|
||||
exclusive with `factories`.
|
||||
is_async: True for async dispatch (`apush` / `aclose` /
|
||||
`afail`), False for the sync path.
|
||||
factories: Zero-or-one-argument callables producing
|
||||
transformers. Called with this mux's `scope`.
|
||||
scope: The namespace the mux operates within. The root mux
|
||||
is `()`; mini-muxes for subgraphs use the subgraph's
|
||||
namespace tuple.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If any transformer requires an async run but
|
||||
the mux is in sync mode.
|
||||
TypeError: If a transformer's `init()` doesn't return a dict.
|
||||
ValueError: If transformers' projection keys collide.
|
||||
ValueError: If transformers' projection keys collide, or if
|
||||
both `transformers` and `factories` are supplied.
|
||||
"""
|
||||
if transformers is not None and factories is not None:
|
||||
raise ValueError("Pass either `transformers` or `factories`, not both.")
|
||||
|
||||
self._is_async = is_async
|
||||
self._factories: list[TransformerFactory] = list(factories or ())
|
||||
self.scope: tuple[str, ...] = scope
|
||||
self._pump_fn: Callable[[], bool] | None = None
|
||||
self._apump_fn: Callable[[], Awaitable[bool]] | None = None
|
||||
|
||||
self._events: EventLog[ProtocolEvent] = EventLog()
|
||||
self._events._bind(is_async=is_async)
|
||||
self._transformers: list[StreamTransformer] = []
|
||||
@@ -72,9 +105,80 @@ class StreamMux:
|
||||
self.extensions: dict[str, Any] = {}
|
||||
self.native_keys: set[str] = set()
|
||||
self._projection_owners: dict[str, str] = {}
|
||||
self._transformer_by_key: dict[str, StreamTransformer] = {}
|
||||
|
||||
for transformer in transformers or ():
|
||||
self._register(transformer)
|
||||
if factories is not None:
|
||||
for factory in factories:
|
||||
self._register(factory(scope))
|
||||
else:
|
||||
for transformer in transformers or ():
|
||||
self._register(transformer)
|
||||
|
||||
def make_child(self, scope: tuple[str, ...]) -> StreamMux:
|
||||
"""Build a mini-mux with the same factories scoped to `scope`.
|
||||
|
||||
Used by `SubgraphTransformer` to attach a fresh transformer
|
||||
pipeline to each discovered subgraph handle. The child mux
|
||||
inherits the current pump binding (so cursors on its projection
|
||||
logs drive the root pump) and carries the same factory list
|
||||
forward to any grandchild subgraphs.
|
||||
|
||||
Raises:
|
||||
RuntimeError: If the mux was not built from a factory list
|
||||
(i.e., constructed with `transformers=`). Mini-muxes
|
||||
require factories so each scope gets its own fresh
|
||||
transformer instances.
|
||||
"""
|
||||
if not self._factories:
|
||||
raise RuntimeError(
|
||||
"StreamMux.make_child requires the mux to be constructed "
|
||||
"with factories; pre-built transformers can't be cloned "
|
||||
"to a new scope."
|
||||
)
|
||||
child = StreamMux(
|
||||
factories=self._factories,
|
||||
is_async=self._is_async,
|
||||
scope=scope,
|
||||
)
|
||||
if self._pump_fn is not None:
|
||||
child.bind_pump(self._pump_fn)
|
||||
if self._apump_fn is not None:
|
||||
child.bind_apump(self._apump_fn)
|
||||
return child
|
||||
|
||||
def bind_pump(self, fn: Callable[[], bool]) -> None:
|
||||
"""Wire the sync pull callback onto every EventLog in the mux.
|
||||
|
||||
Also propagates to transformers that expose `_bind_pump` so
|
||||
nested handles (e.g., `ChatModelStream` instances produced by
|
||||
`MessagesTransformer`) can drive the graph pump from their
|
||||
projection cursors.
|
||||
"""
|
||||
self._pump_fn = fn
|
||||
self._events._request_more = fn
|
||||
for value in self.extensions.values():
|
||||
if isinstance(value, EventLog):
|
||||
value._request_more = fn
|
||||
elif isinstance(value, StreamChannel):
|
||||
value._log._request_more = fn
|
||||
for transformer in self._transformers:
|
||||
bind = getattr(transformer, "_bind_pump", None)
|
||||
if bind is not None:
|
||||
bind(fn)
|
||||
|
||||
def bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None:
|
||||
"""Async counterpart to `bind_pump`."""
|
||||
self._apump_fn = fn
|
||||
self._events._arequest_more = fn
|
||||
for value in self.extensions.values():
|
||||
if isinstance(value, EventLog):
|
||||
value._arequest_more = fn
|
||||
elif isinstance(value, StreamChannel):
|
||||
value._log._arequest_more = fn
|
||||
for transformer in self._transformers:
|
||||
abind = getattr(transformer, "_bind_apump", None)
|
||||
if abind is not None:
|
||||
abind(fn)
|
||||
|
||||
def _register(self, transformer: StreamTransformer) -> None:
|
||||
"""Register a single transformer.
|
||||
@@ -112,13 +216,27 @@ class StreamMux:
|
||||
owner_name = type(transformer).__name__
|
||||
for key in projection:
|
||||
self._projection_owners[key] = owner_name
|
||||
self._transformer_by_key[key] = transformer
|
||||
if getattr(transformer, "_native", False):
|
||||
self.native_keys.update(projection.keys())
|
||||
on_register = getattr(transformer, "_on_register", None)
|
||||
if on_register is not None:
|
||||
on_register(self)
|
||||
|
||||
def transformer_by_key(self, key: str) -> StreamTransformer | None:
|
||||
"""Return the transformer that owns the projection at `key`, if any."""
|
||||
return self._transformer_by_key.get(key)
|
||||
|
||||
def push(self, event: ProtocolEvent) -> None:
|
||||
"""Route an event through all transformers, then append to the main log.
|
||||
|
||||
Each transformer's `process()` is called in registration order.
|
||||
Each transformer's `process()` is called in registration order
|
||||
— except when the transformer has `scope_exact = True` (the
|
||||
default) and the event's namespace differs from the mux's
|
||||
`scope`, in which case the transformer is skipped. Transformers
|
||||
that need to see cross-scope events opt out by setting
|
||||
`scope_exact = False` (e.g. `SubgraphTransformer`).
|
||||
|
||||
If any transformer returns False, the event is suppressed from
|
||||
the main log, but transformers that already saw it keep their
|
||||
side effects.
|
||||
@@ -132,8 +250,12 @@ class StreamMux:
|
||||
Args:
|
||||
event: The protocol event to dispatch.
|
||||
"""
|
||||
ns = tuple(event["params"]["namespace"])
|
||||
in_scope = ns == self.scope
|
||||
keep = True
|
||||
for transformer in self._transformers:
|
||||
if transformer.scope_exact and not in_scope:
|
||||
continue
|
||||
if not transformer.process(event):
|
||||
keep = False
|
||||
if keep:
|
||||
@@ -205,11 +327,13 @@ class StreamMux:
|
||||
"""Dispatch an event on the async lane.
|
||||
|
||||
Awaits each transformer's `aprocess` in registration order
|
||||
before appending to the main log. A slow `aprocess` serializes
|
||||
the pipeline by design — that's the guarantee that lets a later
|
||||
transformer (or a synchronous consumer) see the result of the
|
||||
async work. For decoupled work, use `schedule()` from inside
|
||||
`process` / `aprocess` instead.
|
||||
before appending to the main log — except when the transformer
|
||||
has `scope_exact = True` and the event's namespace differs from
|
||||
`self.scope`, in which case it is skipped. A slow `aprocess`
|
||||
serializes the pipeline by design — that's the guarantee that
|
||||
lets a later transformer (or a synchronous consumer) see the
|
||||
result of the async work. For decoupled work, use `schedule()`
|
||||
from inside `process` / `aprocess` instead.
|
||||
|
||||
The main log append is a non-blocking `push` — matching v1's
|
||||
`put_nowait` shape. Memory is bounded by caller pace via the
|
||||
@@ -218,8 +342,12 @@ class StreamMux:
|
||||
Args:
|
||||
event: The protocol event to dispatch.
|
||||
"""
|
||||
ns = tuple(event["params"]["namespace"])
|
||||
in_scope = ns == self.scope
|
||||
keep = True
|
||||
for transformer in self._transformers:
|
||||
if transformer.scope_exact and not in_scope:
|
||||
continue
|
||||
if not await transformer.aprocess(event):
|
||||
keep = False
|
||||
if keep:
|
||||
|
||||
@@ -75,6 +75,20 @@ class StreamTransformer(ABC):
|
||||
scoring, cost lookup, external tracing).
|
||||
|
||||
Attributes:
|
||||
scope: Namespace the transformer operates within — `()` for the
|
||||
root mux, a subgraph's namespace tuple inside a mini-mux.
|
||||
Set at construction from the mux's scope (each factory is
|
||||
called as `factory(scope)`). Transformers that only care
|
||||
about events at their own namespace compare against
|
||||
`self.scope`; subgraph-aware transformers can treat it as
|
||||
a parent path.
|
||||
scope_exact: If True (the default), the mux only calls
|
||||
`process` / `aprocess` for events whose namespace equals
|
||||
`self.scope` — user transformers get scope-scoped events
|
||||
for free with no boilerplate. Set False for transformers
|
||||
that need to see events across scopes (e.g.
|
||||
`SubgraphTransformer` forwards deeper events into child
|
||||
mini-muxes).
|
||||
requires_async: Explicit opt-in for transformers that need a
|
||||
running event loop but don't override any async method (for
|
||||
example, transformers that call `schedule()` from a sync
|
||||
@@ -83,6 +97,18 @@ class StreamTransformer(ABC):
|
||||
"""
|
||||
|
||||
requires_async: ClassVar[bool] = False
|
||||
scope_exact: ClassVar[bool] = True
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
"""Initialize the transformer with its mux's scope.
|
||||
|
||||
Args:
|
||||
scope: The namespace tuple the owning mux is scoped to.
|
||||
`()` for the root, the subgraph's namespace inside a
|
||||
mini-mux. Factories receive this at construction time
|
||||
(`factory(scope)` in `StreamMux`).
|
||||
"""
|
||||
self.scope: tuple[str, ...] = scope
|
||||
|
||||
@abstractmethod
|
||||
def init(self) -> dict[str, Any]:
|
||||
|
||||
@@ -3,14 +3,14 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator, Mapping
|
||||
from types import MappingProxyType, TracebackType
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from langgraph.stream._convert import convert_to_protocol_event
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._types import ProtocolEvent
|
||||
from langgraph.stream.stream_channel import StreamChannel
|
||||
from langgraph.stream.transformers import ValuesTransformer
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langgraph.stream.transformers import ValuesTransformer
|
||||
|
||||
|
||||
def _drive_until_done(pump: Callable[[], bool]) -> None:
|
||||
@@ -25,7 +25,99 @@ async def _adrive_until_done(pump: Callable[[], Awaitable[bool]]) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class GraphRunStream:
|
||||
class BaseRunStream:
|
||||
"""Shared shape for any object that wraps a `StreamMux`.
|
||||
|
||||
Both the root `GraphRunStream` / `AsyncGraphRunStream` and the
|
||||
scoped `SubgraphRunStream` compose a `StreamMux` and expose its
|
||||
projections (`values`, `messages`, `subgraphs`, user-registered
|
||||
keys). The root additionally owns the graph iterator and drives
|
||||
the pump; a subgraph's mini-mux borrows the root pump via
|
||||
`make_child`'s pump inheritance.
|
||||
|
||||
Projections registered on the mux show up in `extensions`, and
|
||||
native ones (those with `_native = True`) are also bound directly
|
||||
as attributes (`run.values`, `run.messages`, …).
|
||||
|
||||
Raw iteration (`for event in run:` / `async for event in run`)
|
||||
yields every `ProtocolEvent` that reached this mux's main log —
|
||||
for the root that's every event in the run, for a subgraph that's
|
||||
every event forwarded into its subtree.
|
||||
"""
|
||||
|
||||
def __init__(self, mux: StreamMux) -> None:
|
||||
self._mux = mux
|
||||
self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions)
|
||||
for key in mux.native_keys:
|
||||
setattr(self, key, mux.extensions[key])
|
||||
|
||||
def __iter__(self) -> Iterator[ProtocolEvent]:
|
||||
"""Sync iteration of protocol events on this mux's main log.
|
||||
|
||||
Raises at the EventLog level if the mux is async-bound.
|
||||
"""
|
||||
return iter(self._mux._events)
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[ProtocolEvent]:
|
||||
"""Async iteration of protocol events on this mux's main log.
|
||||
|
||||
Raises at the EventLog level if the mux is sync-bound.
|
||||
"""
|
||||
return self._mux._events.__aiter__()
|
||||
|
||||
def interleave(self, *names: str) -> Iterator[tuple[str, Any]]:
|
||||
"""Iterate multiple projections round-robin, yielding ``(name, item)``.
|
||||
|
||||
Each turn advances one projection's cursor; when a cursor's
|
||||
buffer is empty, pulling from it drives the pump once, which
|
||||
fans out to every subscribed projection log. Projections whose
|
||||
items aren't consumed on this turn sit in their own buffers
|
||||
only until the next turn reaches them, bounding memory by the
|
||||
skew between projection rates rather than letting any single
|
||||
log grow to the full run length.
|
||||
|
||||
Projections are exhausted independently; a projection that
|
||||
finishes early drops out of the rotation while others
|
||||
continue. The overall iterator ends once all named projections
|
||||
are done.
|
||||
|
||||
Args:
|
||||
*names: Projection keys to interleave. Must match keys in
|
||||
`extensions`.
|
||||
|
||||
Yields:
|
||||
`(name, item)` tuples in round-robin order across the named
|
||||
projections.
|
||||
|
||||
Raises:
|
||||
KeyError: If a name doesn't match a registered projection.
|
||||
|
||||
Example:
|
||||
```python
|
||||
for name, item in run.interleave("messages", "values"):
|
||||
if name == "messages":
|
||||
print("msg:", item)
|
||||
else:
|
||||
print("val:", item)
|
||||
```
|
||||
"""
|
||||
cursors: dict[str, Iterator[Any]] = {
|
||||
name: iter(self.extensions[name]) for name in names
|
||||
}
|
||||
done: set[str] = set()
|
||||
while len(done) < len(cursors):
|
||||
for name, cursor in cursors.items():
|
||||
if name in done:
|
||||
continue
|
||||
try:
|
||||
item = next(cursor)
|
||||
except StopIteration:
|
||||
done.add(name)
|
||||
continue
|
||||
yield (name, item)
|
||||
|
||||
|
||||
class GraphRunStream(BaseRunStream):
|
||||
"""Sync run stream with caller-driven pumping.
|
||||
|
||||
The caller's iteration on any projection (`values`, `messages`,
|
||||
@@ -34,10 +126,6 @@ class GraphRunStream:
|
||||
|
||||
Projections are single-consumer — iterating `run.values` twice
|
||||
raises. Use `projection.tee(n)` if you genuinely need fan-out.
|
||||
|
||||
All transformer projections live in `extensions`. Native transformer
|
||||
projections (those with `_native = True`) are also set as direct
|
||||
attributes on this instance (e.g. `run.values`, `run.messages`).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -54,40 +142,18 @@ class GraphRunStream:
|
||||
values_transformer: The built-in values transformer
|
||||
providing `output` / `interrupted` / `interrupts`.
|
||||
"""
|
||||
super().__init__(mux)
|
||||
self._graph_iter = graph_iter
|
||||
self._mux = mux
|
||||
self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions)
|
||||
self._values_transformer = values_transformer
|
||||
self._exhausted = False
|
||||
for key in mux.native_keys:
|
||||
setattr(self, key, mux.extensions[key])
|
||||
self._wire_request_more(mux)
|
||||
|
||||
def _wire_request_more(self, mux: StreamMux) -> None:
|
||||
"""Install `_request_more` on every sync EventLog so cursors can
|
||||
drive the pump when their buffer catches up.
|
||||
|
||||
Also calls `_bind_pump` on any transformer that exposes it, so
|
||||
transformers producing ChatModelStream objects (e.g.
|
||||
MessagesTransformer) can wire the pull callback on each stream
|
||||
as it's created.
|
||||
"""
|
||||
mux._events._request_more = self._pump_next
|
||||
for value in mux.extensions.values():
|
||||
if isinstance(value, EventLog):
|
||||
value._request_more = self._pump_next
|
||||
elif isinstance(value, StreamChannel):
|
||||
value._log._request_more = self._pump_next
|
||||
for transformer in mux._transformers:
|
||||
if hasattr(transformer, "_bind_pump"):
|
||||
transformer._bind_pump(self._pump_next)
|
||||
mux.bind_pump(self._pump_next)
|
||||
|
||||
def _pump_next(self) -> bool:
|
||||
"""Pull one event from the graph and push it through the mux.
|
||||
|
||||
Returns:
|
||||
True if an event was pulled, False if the graph is exhausted
|
||||
or has raised.
|
||||
True if an event was pulled, False if the graph is
|
||||
exhausted or has raised.
|
||||
"""
|
||||
if self._exhausted:
|
||||
return False
|
||||
@@ -141,8 +207,7 @@ class GraphRunStream:
|
||||
|
||||
@property
|
||||
def interrupted(self) -> bool:
|
||||
"""Drive the run to completion, then return whether it was
|
||||
interrupted.
|
||||
"""Drive the run to completion, then return whether it was interrupted.
|
||||
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
@@ -166,70 +231,17 @@ class GraphRunStream:
|
||||
raise err
|
||||
return self._values_transformer._interrupts
|
||||
|
||||
def __iter__(self) -> Iterator[ProtocolEvent]:
|
||||
"""Subscribe to the main event log and iterate protocol events."""
|
||||
return iter(self._mux._events)
|
||||
|
||||
def interleave(self, *names: str) -> Iterator[tuple[str, Any]]:
|
||||
"""Iterate multiple projections round-robin, yielding ``(name, item)``.
|
||||
|
||||
Each turn advances one projection's cursor; when a cursor's buffer
|
||||
is empty, pulling from it drives the pump once, which fans out to
|
||||
every subscribed projection log. Projections whose items aren't
|
||||
consumed on this turn sit in their own buffers only until the next
|
||||
turn reaches them, bounding memory by the skew between projection
|
||||
rates rather than letting any single log grow to the full run
|
||||
length.
|
||||
|
||||
Projections are exhausted independently; a projection that finishes
|
||||
early drops out of the rotation while others continue. The overall
|
||||
iterator ends once all named projections are done.
|
||||
|
||||
Args:
|
||||
*names: Projection keys to interleave. Must match keys in
|
||||
``extensions``.
|
||||
|
||||
Yields:
|
||||
``(name, item)`` tuples in round-robin order across the named
|
||||
projections.
|
||||
|
||||
Raises:
|
||||
KeyError: If a name doesn't match a registered projection.
|
||||
|
||||
Example:
|
||||
```python
|
||||
for name, item in run.interleave("messages", "values"):
|
||||
if name == "messages":
|
||||
print("msg:", item)
|
||||
else:
|
||||
print("val:", item)
|
||||
```
|
||||
"""
|
||||
cursors: dict[str, Iterator[Any]] = {
|
||||
name: iter(self.extensions[name]) for name in names
|
||||
}
|
||||
done: set[str] = set()
|
||||
while len(done) < len(cursors):
|
||||
for name, cursor in cursors.items():
|
||||
if name in done:
|
||||
continue
|
||||
try:
|
||||
item = next(cursor)
|
||||
except StopIteration:
|
||||
done.add(name)
|
||||
continue
|
||||
yield (name, item)
|
||||
|
||||
|
||||
class AsyncGraphRunStream:
|
||||
class AsyncGraphRunStream(BaseRunStream):
|
||||
"""Async run stream with caller-driven pumping.
|
||||
|
||||
Async iteration on any projection drives the graph forward — there
|
||||
is no background task. Concurrent consumers share a single-flight
|
||||
pump via an `asyncio.Lock`, so each awaiting cursor contributes one
|
||||
event per acquisition. Backpressure comes from the logs: when a
|
||||
subscribed log's buffer reaches `maxlen`, `apush` awaits the
|
||||
subscriber to drain, which holds back the pump and paces the graph.
|
||||
pump via an `asyncio.Lock`, so each awaiting cursor contributes
|
||||
one event per acquisition. Backpressure comes from the logs: when
|
||||
a subscribed log's buffer reaches `maxlen`, `apush` awaits the
|
||||
subscriber to drain, which holds back the pump and paces the
|
||||
graph.
|
||||
|
||||
Projections are single-consumer — a second `aiter(run.values)`
|
||||
raises. Use `projection.tee(n)` for fan-out.
|
||||
@@ -258,35 +270,12 @@ class AsyncGraphRunStream:
|
||||
values_transformer: The built-in values transformer
|
||||
providing `output` / `interrupted` / `interrupts`.
|
||||
"""
|
||||
super().__init__(mux)
|
||||
self._graph_aiter = graph_aiter
|
||||
self._mux = mux
|
||||
self.extensions: Mapping[str, Any] = MappingProxyType(mux.extensions)
|
||||
self._values_transformer = values_transformer
|
||||
self._exhausted = False
|
||||
self._pump_lock = asyncio.Lock()
|
||||
for key in mux.native_keys:
|
||||
setattr(self, key, mux.extensions[key])
|
||||
self._wire_arequest_more(mux)
|
||||
|
||||
def _wire_arequest_more(self, mux: StreamMux) -> None:
|
||||
"""Install `_arequest_more` on every async EventLog so cursors
|
||||
can drive the pump when their buffer catches up.
|
||||
|
||||
Also calls `_bind_apump` on any transformer that exposes it,
|
||||
so transformers producing `AsyncChatModelStream` objects (e.g.
|
||||
`MessagesTransformer`) can fan the pull callback out to each
|
||||
stream's projections. Mirrors the sync `_wire_request_more`
|
||||
plumbing.
|
||||
"""
|
||||
mux._events._arequest_more = self._apump_next
|
||||
for value in mux.extensions.values():
|
||||
if isinstance(value, EventLog):
|
||||
value._arequest_more = self._apump_next
|
||||
elif isinstance(value, StreamChannel):
|
||||
value._log._arequest_more = self._apump_next
|
||||
for transformer in mux._transformers:
|
||||
if hasattr(transformer, "_bind_apump"):
|
||||
transformer._bind_apump(self._apump_next)
|
||||
mux.bind_apump(self._apump_next)
|
||||
|
||||
async def _apump_next(self) -> bool:
|
||||
"""Pull one event from the graph and push it through the mux.
|
||||
@@ -324,7 +313,8 @@ class AsyncGraphRunStream:
|
||||
|
||||
Closes the mux and marks the stream exhausted. Any awaiting
|
||||
cursors wake up and see the closed state; any `apush` blocked
|
||||
on backpressure wakes and returns without appending. Idempotent.
|
||||
on backpressure wakes and returns without appending.
|
||||
Idempotent.
|
||||
"""
|
||||
async with self._pump_lock:
|
||||
if self._exhausted:
|
||||
@@ -367,8 +357,7 @@ class AsyncGraphRunStream:
|
||||
return self._values_transformer._latest
|
||||
|
||||
async def interrupted(self) -> bool:
|
||||
"""Drive the run to completion and return whether it was
|
||||
interrupted.
|
||||
"""Drive the run to completion and return whether it was interrupted.
|
||||
|
||||
Raises:
|
||||
BaseException: If the run ended with an error.
|
||||
@@ -388,7 +377,3 @@ class AsyncGraphRunStream:
|
||||
if (err := self._values_transformer.error) is not None:
|
||||
raise err
|
||||
return self._values_transformer._interrupts
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[ProtocolEvent]:
|
||||
"""Subscribe to the main event log and iterate protocol events."""
|
||||
return self._mux._events.__aiter__()
|
||||
|
||||
@@ -7,13 +7,53 @@ from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph._internal._constants import CONF, CONFIG_KEY_STREAM_MESSAGES_V2
|
||||
from langgraph.pregel import Pregel
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._mux import StreamMux, TransformerFactory
|
||||
from langgraph.stream._types import StreamTransformer
|
||||
from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream
|
||||
from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer
|
||||
from langgraph.stream.transformers import (
|
||||
MessagesTransformer,
|
||||
SubgraphTransformer,
|
||||
ValuesTransformer,
|
||||
)
|
||||
from langgraph.types import All, StreamMode
|
||||
|
||||
|
||||
def _coerce_factories(
|
||||
transformers: list[StreamTransformer | TransformerFactory] | None,
|
||||
) -> list[TransformerFactory]:
|
||||
"""Normalize caller-supplied transformers into scope-taking factories.
|
||||
|
||||
Accepts already-built instances (wrapped as single-use factories,
|
||||
with the caveat that they won't be re-instantiated in subgraph
|
||||
mini-muxes) or proper factories (classes / callables taking a
|
||||
scope). The built-in root transformers are always factories so
|
||||
they propagate into every subgraph scope automatically.
|
||||
"""
|
||||
|
||||
def _wrap_instance(t: StreamTransformer) -> TransformerFactory:
|
||||
# Single-use: only wires at root scope. A user that wants
|
||||
# subgraph propagation should pass the class (or a lambda).
|
||||
def _factory(_scope: tuple[str, ...]) -> StreamTransformer:
|
||||
return t
|
||||
|
||||
return _factory
|
||||
|
||||
coerced: list[TransformerFactory] = []
|
||||
for item in transformers or ():
|
||||
if isinstance(item, StreamTransformer):
|
||||
coerced.append(_wrap_instance(item))
|
||||
else:
|
||||
coerced.append(item)
|
||||
return coerced
|
||||
|
||||
|
||||
_BUILTIN_FACTORIES: list[TransformerFactory] = [
|
||||
ValuesTransformer,
|
||||
MessagesTransformer,
|
||||
SubgraphTransformer,
|
||||
]
|
||||
|
||||
|
||||
def _merge_v2_messages_flag(
|
||||
config: RunnableConfig | None,
|
||||
) -> RunnableConfig:
|
||||
@@ -40,6 +80,7 @@ STREAM_V2_MODES: list[StreamMode] = [
|
||||
"checkpoints",
|
||||
"tasks",
|
||||
"debug",
|
||||
"lifecycle",
|
||||
]
|
||||
|
||||
|
||||
@@ -80,7 +121,7 @@ class StreamingHandler:
|
||||
*,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
transformers: list[StreamTransformer | TransformerFactory] | None = None,
|
||||
) -> GraphRunStream:
|
||||
"""Start a sync streaming run.
|
||||
|
||||
@@ -100,11 +141,12 @@ class StreamingHandler:
|
||||
Returns:
|
||||
A GraphRunStream the caller can iterate to drive the run.
|
||||
"""
|
||||
values_t = ValuesTransformer()
|
||||
mux = StreamMux(
|
||||
[values_t, MessagesTransformer(), *(transformers or ())],
|
||||
factories=_BUILTIN_FACTORIES + _coerce_factories(transformers),
|
||||
is_async=False,
|
||||
)
|
||||
values_t = mux.transformer_by_key("values")
|
||||
assert isinstance(values_t, ValuesTransformer)
|
||||
|
||||
graph_iter = iter(
|
||||
self._graph.stream(
|
||||
@@ -127,7 +169,7 @@ class StreamingHandler:
|
||||
*,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
transformers: list[StreamTransformer | TransformerFactory] | None = None,
|
||||
) -> AsyncGraphRunStream:
|
||||
"""Start an async streaming run.
|
||||
|
||||
@@ -149,11 +191,12 @@ class StreamingHandler:
|
||||
concurrently; each subscribed cursor drives the pump when
|
||||
its buffer is empty.
|
||||
"""
|
||||
values_t = ValuesTransformer()
|
||||
mux = StreamMux(
|
||||
[values_t, MessagesTransformer(), *(transformers or ())],
|
||||
factories=_BUILTIN_FACTORIES + _coerce_factories(transformers),
|
||||
is_async=True,
|
||||
)
|
||||
values_t = mux.transformer_by_key("values")
|
||||
assert isinstance(values_t, ValuesTransformer)
|
||||
|
||||
graph_aiter = self._graph.astream(
|
||||
input,
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
from typing import TYPE_CHECKING, Any, Literal, cast
|
||||
|
||||
from langchain_core.language_models._compat_bridge import message_to_events
|
||||
from langchain_core.language_models.chat_model_stream import (
|
||||
@@ -8,14 +8,24 @@ from langchain_core.language_models.chat_model_stream import (
|
||||
ChatModelStream,
|
||||
)
|
||||
from langchain_core.messages import AIMessageChunk, BaseMessage
|
||||
from langchain_protocol.protocol import MessagesData
|
||||
from langchain_protocol.protocol import CheckpointRef, LifecycleData, MessagesData
|
||||
|
||||
from langgraph.errors import GraphInterrupt
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.run_stream import BaseRunStream
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from langgraph.stream._mux import StreamMux
|
||||
|
||||
|
||||
SubgraphStatus = Literal["started", "running", "completed", "failed", "interrupted"]
|
||||
_TERMINAL_STATUSES: frozenset[SubgraphStatus] = frozenset(
|
||||
{"completed", "failed", "interrupted"}
|
||||
)
|
||||
|
||||
|
||||
class ValuesTransformer(StreamTransformer):
|
||||
"""Capture values events as a drainable stream of state snapshots.
|
||||
@@ -28,13 +38,16 @@ class ValuesTransformer(StreamTransformer):
|
||||
Native transformer — projection keys are exposed as direct
|
||||
attributes on the run stream (e.g. `run.values`).
|
||||
|
||||
Only root-namespace values events are captured; subgraph state
|
||||
snapshots are ignored.
|
||||
`scope` (inherited from `StreamTransformer`) is the namespace the
|
||||
transformer captures values for. `()` matches the root graph;
|
||||
subgraph mini-muxes pass their subgraph's namespace, so each
|
||||
instance sees only its own level.
|
||||
"""
|
||||
|
||||
_native = True
|
||||
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
self._log: EventLog[dict[str, Any]] = EventLog()
|
||||
self._latest: dict[str, Any] | None = None
|
||||
self._interrupted = False
|
||||
@@ -52,11 +65,10 @@ class ValuesTransformer(StreamTransformer):
|
||||
return self._log._error
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
# Namespace filtering is handled by the mux via `scope_exact`.
|
||||
if event["method"] != "values":
|
||||
return True
|
||||
params = event["params"]
|
||||
if params["namespace"]:
|
||||
return True
|
||||
self._latest = params["data"]
|
||||
interrupts = params.get("interrupts", ())
|
||||
if interrupts:
|
||||
@@ -95,9 +107,10 @@ class MessagesTransformer(StreamTransformer):
|
||||
`stream()` method still surface their final `AIMessage` via
|
||||
`on_chain_end` when a node returns it as state.
|
||||
|
||||
Only root-namespace events are captured; tokens from subgraphs are
|
||||
dropped. Consumers that need subgraph tokens should iterate the raw
|
||||
event stream or register a custom transformer.
|
||||
`scope` (inherited from `StreamTransformer`) is the namespace the
|
||||
transformer captures messages for. `()` matches the root graph;
|
||||
subgraph mini-muxes pass their subgraph's namespace, so each
|
||||
instance sees only its own level.
|
||||
|
||||
Native transformer — the `messages` projection is exposed as a
|
||||
direct attribute on the run stream.
|
||||
@@ -105,7 +118,8 @@ class MessagesTransformer(StreamTransformer):
|
||||
|
||||
_native = True
|
||||
|
||||
def __init__(self) -> None:
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
self._log: EventLog[ChatModelStream] = EventLog()
|
||||
# Correlate protocol events back to a ChatModelStream by run_id
|
||||
# (attached to the event's metadata by StreamMessagesHandler).
|
||||
@@ -166,11 +180,10 @@ class MessagesTransformer(StreamTransformer):
|
||||
)
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
# Namespace filtering is handled by the mux via `scope_exact`.
|
||||
if event["method"] != "messages":
|
||||
return True
|
||||
params = event["params"]
|
||||
if params["namespace"]:
|
||||
return True
|
||||
|
||||
payload, metadata = params["data"]
|
||||
node: str | None = metadata.get("langgraph_node")
|
||||
@@ -201,7 +214,7 @@ class MessagesTransformer(StreamTransformer):
|
||||
if event_type == "message-start":
|
||||
message_id = event.get("message_id")
|
||||
stream = self._make_stream(
|
||||
namespace=[],
|
||||
namespace=list(self.scope),
|
||||
node=node,
|
||||
message_id=str(message_id) if message_id is not None else None,
|
||||
)
|
||||
@@ -215,7 +228,11 @@ class MessagesTransformer(StreamTransformer):
|
||||
del self._by_run[run_id]
|
||||
|
||||
def _route_whole_message(self, message: BaseMessage, *, node: str | None) -> None:
|
||||
stream = self._make_stream(namespace=[], node=node, message_id=message.id)
|
||||
stream = self._make_stream(
|
||||
namespace=list(self.scope),
|
||||
node=node,
|
||||
message_id=message.id,
|
||||
)
|
||||
for evt in message_to_events(message, message_id=message.id):
|
||||
stream.dispatch(evt)
|
||||
self._log.push(stream)
|
||||
@@ -229,3 +246,213 @@ class MessagesTransformer(StreamTransformer):
|
||||
for stream in list(self._by_run.values()):
|
||||
stream.fail(err)
|
||||
self._by_run.clear()
|
||||
|
||||
|
||||
class SubgraphRunStream(BaseRunStream):
|
||||
"""Scoped view of a single nested subgraph execution.
|
||||
|
||||
Yielded on `run.subgraphs` (or `parent.subgraphs` for grandchildren)
|
||||
when a nested `Pregel` spawns. Wraps a mini-`StreamMux` built with
|
||||
the same transformer factories as the root mux, so `.values`,
|
||||
`.messages`, `.subgraphs` are populated by the standard
|
||||
transformers scoped to this handle's namespace — no duplicated
|
||||
routing logic.
|
||||
|
||||
Inherits `BaseRunStream`, so all shared shape works uniformly:
|
||||
raw iteration (`for event in sub`), `.interleave(...)`, native
|
||||
projection attributes, and the `extensions` mapping — identical
|
||||
to the root `GraphRunStream`. The mini-mux borrows the root's
|
||||
pump via `make_child`'s pump inheritance, so any cursor on a
|
||||
subagent projection drives the whole run forward.
|
||||
|
||||
Lifecycle fields update in place as events arrive:
|
||||
|
||||
- `path`: the namespace tuple — stable for the life of the handle.
|
||||
- `graph_name` / `trigger_call_id`: set once from the `started`
|
||||
payload.
|
||||
- `status`: advances `started` → `running` → `completed` /
|
||||
`failed` / `interrupted`.
|
||||
- `error` / `checkpoint`: set on the terminal event when present.
|
||||
|
||||
`.output` is a snapshot of the latest values seen at this
|
||||
namespace — it doesn't drive the pump (unlike root's
|
||||
`GraphRunStream.output`), because advancing a subgraph to
|
||||
completion is only meaningful as part of advancing the whole run.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
path: tuple[str, ...],
|
||||
mux: StreamMux,
|
||||
*,
|
||||
graph_name: str | None = None,
|
||||
trigger_call_id: str | None = None,
|
||||
) -> None:
|
||||
super().__init__(mux)
|
||||
self.path: tuple[str, ...] = path
|
||||
self.graph_name: str | None = graph_name
|
||||
self.trigger_call_id: str | None = trigger_call_id
|
||||
self.status: SubgraphStatus = "started"
|
||||
self.error: str | None = None
|
||||
self.checkpoint: CheckpointRef | None = None
|
||||
|
||||
@property
|
||||
def output(self) -> dict[str, Any] | None:
|
||||
"""Latest values snapshot at this namespace, or `None`.
|
||||
|
||||
Snapshot-only — iterating other projections or the root's
|
||||
`.output` is what drives the pump.
|
||||
"""
|
||||
values_t = self._mux.transformer_by_key("values")
|
||||
if isinstance(values_t, ValuesTransformer):
|
||||
return values_t._latest
|
||||
return None
|
||||
|
||||
|
||||
class SubgraphTransformer(StreamTransformer):
|
||||
"""Discover subgraphs and route events into per-subgraph mini-muxes.
|
||||
|
||||
Thin state-machine + dispatcher. At its own `scope` (inherited
|
||||
from `StreamTransformer`, determined by the enclosing mux), it
|
||||
watches for `lifecycle` events at exactly one level deeper to
|
||||
discover direct children. Each discovered child gets its own
|
||||
`SubgraphRunStream` backed by a mini-`StreamMux` — built via
|
||||
`parent_mux.make_child(path)`, so the same factory list produces
|
||||
fresh transformer instances at the child's scope.
|
||||
|
||||
Every incoming event that falls under one of the direct children
|
||||
(ns starts with a child's `path`) is forwarded into that child's
|
||||
mini-mux via `push`. The standard transformers in that mini-mux
|
||||
(`ValuesTransformer`, `MessagesTransformer`, and another
|
||||
`SubgraphTransformer` for grandchildren) handle the rest. No
|
||||
duplicated routing or assembly logic.
|
||||
|
||||
Lifecycle state for each handle (running / completed / failed /
|
||||
interrupted) is updated in place as events fire. On terminal
|
||||
events, the handle's mini-mux is closed so any subscribed cursors
|
||||
unblock. `finalize` / `fail` handle dangling handles left mid-run.
|
||||
|
||||
Native transformer — `subgraphs` exposes the direct-children log.
|
||||
|
||||
`scope_exact = False`: this transformer sees events at any
|
||||
namespace, because it forwards out-of-scope events to the matching
|
||||
direct-child mini-mux.
|
||||
"""
|
||||
|
||||
_native = True
|
||||
scope_exact = False
|
||||
|
||||
def __init__(self, scope: tuple[str, ...] = ()) -> None:
|
||||
super().__init__(scope)
|
||||
self._root_log: EventLog[SubgraphRunStream] = EventLog()
|
||||
# Direct children only (namespace = scope + one segment).
|
||||
self._by_ns: dict[tuple[str, ...], SubgraphRunStream] = {}
|
||||
self._mux: StreamMux | None = None
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"subgraphs": self._root_log}
|
||||
|
||||
def _on_register(self, mux: StreamMux) -> None:
|
||||
"""Capture the enclosing mux so we can build child mini-muxes."""
|
||||
self._mux = mux
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
ns = tuple(event["params"]["namespace"])
|
||||
method = event["method"]
|
||||
depth = len(self.scope)
|
||||
|
||||
# 1. On `started` for a direct child (ns depth = mine + 1 and
|
||||
# ns prefix matches mine), register the handle.
|
||||
if method == "lifecycle" and len(ns) == depth + 1 and ns[:-1] == self.scope:
|
||||
data = cast(LifecycleData, event["params"]["data"])
|
||||
if data.get("event") == "started":
|
||||
self._on_started(ns, data)
|
||||
|
||||
# 2. Forward the event to the matching direct-child mini-mux.
|
||||
# Prefix-match: ns must start with some child's path.
|
||||
direct_child_ns = ns[: depth + 1] if len(ns) > depth else None
|
||||
if direct_child_ns is not None and direct_child_ns in self._by_ns:
|
||||
self._by_ns[direct_child_ns]._mux.push(event)
|
||||
|
||||
# 3. Status change for a direct child (ns = child's path, method
|
||||
# = lifecycle). Update handle fields, close mini-mux on
|
||||
# terminal.
|
||||
if (
|
||||
method == "lifecycle"
|
||||
and ns in self._by_ns
|
||||
and len(ns) == depth + 1
|
||||
and ns[:-1] == self.scope
|
||||
):
|
||||
data = cast(LifecycleData, event["params"]["data"])
|
||||
event_type = data.get("event")
|
||||
if event_type in ("running", "completed", "failed", "interrupted"):
|
||||
self._on_status_change(ns, event_type, data)
|
||||
|
||||
return True
|
||||
|
||||
def _on_started(self, ns: tuple[str, ...], data: LifecycleData) -> None:
|
||||
if ns in self._by_ns:
|
||||
# Duplicate started — ignore.
|
||||
return
|
||||
if self._mux is None:
|
||||
# Not registered yet; can't build a mini-mux.
|
||||
return
|
||||
child_mux = self._mux.make_child(ns)
|
||||
handle = SubgraphRunStream(
|
||||
path=ns,
|
||||
mux=child_mux,
|
||||
graph_name=data.get("graph_name"),
|
||||
trigger_call_id=data.get("trigger_call_id"),
|
||||
)
|
||||
self._by_ns[ns] = handle
|
||||
self._root_log.push(handle)
|
||||
|
||||
def _on_status_change(
|
||||
self,
|
||||
ns: tuple[str, ...],
|
||||
event_type: SubgraphStatus,
|
||||
data: LifecycleData,
|
||||
) -> None:
|
||||
handle = self._by_ns[ns]
|
||||
handle.status = event_type
|
||||
err = data.get("error")
|
||||
if err is not None:
|
||||
handle.error = err
|
||||
checkpoint = data.get("checkpoint")
|
||||
if checkpoint is not None:
|
||||
handle.checkpoint = checkpoint
|
||||
if event_type in _TERMINAL_STATUSES:
|
||||
self._close_handle_mux(handle)
|
||||
|
||||
@staticmethod
|
||||
def _close_handle_mux(handle: SubgraphRunStream) -> None:
|
||||
# Idempotent close — mux.close() runs finalize on its transformers
|
||||
# (which cascades through grandchildren) and closes projection logs.
|
||||
if not handle._mux._events._closed:
|
||||
try:
|
||||
handle._mux.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def finalize(self) -> None:
|
||||
"""Transition any still-open direct children to `completed`."""
|
||||
for handle in self._by_ns.values():
|
||||
if handle.status not in _TERMINAL_STATUSES:
|
||||
handle.status = "completed"
|
||||
self._close_handle_mux(handle)
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
"""Transition any still-open direct children to `failed` / `interrupted`."""
|
||||
is_interrupt = isinstance(err, GraphInterrupt)
|
||||
terminal: SubgraphStatus = "interrupted" if is_interrupt else "failed"
|
||||
error_str = None if is_interrupt else str(err)
|
||||
for handle in self._by_ns.values():
|
||||
if handle.status not in _TERMINAL_STATUSES:
|
||||
handle.status = terminal
|
||||
if error_str is not None and handle.error is None:
|
||||
handle.error = error_str
|
||||
if not handle._mux._events._closed:
|
||||
try:
|
||||
handle._mux.fail(err)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
@@ -116,7 +116,14 @@ def ensure_valid_checkpointer(checkpointer: Checkpointer) -> Checkpointer:
|
||||
|
||||
|
||||
StreamMode = Literal[
|
||||
"values", "updates", "checkpoints", "tasks", "debug", "messages", "custom"
|
||||
"values",
|
||||
"updates",
|
||||
"checkpoints",
|
||||
"tasks",
|
||||
"debug",
|
||||
"messages",
|
||||
"custom",
|
||||
"lifecycle",
|
||||
]
|
||||
"""How the stream method should emit outputs.
|
||||
|
||||
@@ -129,6 +136,7 @@ StreamMode = Literal[
|
||||
- `"checkpoints"`: Emit an event when a checkpoint is created, in the same format as returned by `get_state()`.
|
||||
- `"tasks"`: Emit events when tasks start and finish, including their results and errors.
|
||||
- `"debug"`: Emit `"checkpoints"` and `"tasks"` events for debugging purposes.
|
||||
- `"lifecycle"`: Emit subgraph lifecycle events (`started`, `running`, `completed`, `failed`, `interrupted`) with payloads matching `LifecycleData`.
|
||||
"""
|
||||
|
||||
StreamWriter = Callable[[Any], None]
|
||||
|
||||
@@ -318,8 +318,16 @@ class TestFiltering:
|
||||
assert t.process(values_event) is True
|
||||
|
||||
def test_subgraph_namespace_dropped(self) -> None:
|
||||
t, log = _make_sync_transformer()
|
||||
t.process(
|
||||
"""Root MessagesTransformer (via the mux) ignores non-root events."""
|
||||
from langgraph.stream._mux import StreamMux
|
||||
|
||||
mux = StreamMux([MessagesTransformer()], is_async=False)
|
||||
t = mux.transformer_by_key("messages")
|
||||
assert isinstance(t, MessagesTransformer)
|
||||
t._log._subscribed = True
|
||||
t._bind_pump(lambda: False)
|
||||
|
||||
mux.push(
|
||||
{
|
||||
"type": "event",
|
||||
"method": "messages",
|
||||
@@ -333,8 +341,8 @@ class TestFiltering:
|
||||
},
|
||||
}
|
||||
)
|
||||
log.close()
|
||||
assert list(log._items) == []
|
||||
t._log.close()
|
||||
assert list(t._log._items) == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -0,0 +1,367 @@
|
||||
"""Tests for subgraph lifecycle events and the SubgraphTransformer."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import operator
|
||||
import time
|
||||
from typing import Annotated, Any
|
||||
|
||||
import pytest
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from langgraph.constants import END, START
|
||||
from langgraph.errors import GraphInterrupt
|
||||
from langgraph.graph import StateGraph
|
||||
from langgraph.stream import StreamingHandler
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._types import ProtocolEvent
|
||||
from langgraph.stream.transformers import (
|
||||
MessagesTransformer,
|
||||
SubgraphRunStream,
|
||||
SubgraphTransformer,
|
||||
ValuesTransformer,
|
||||
)
|
||||
|
||||
TS = int(time.time() * 1000)
|
||||
|
||||
|
||||
def _lifecycle(
|
||||
event: str,
|
||||
*,
|
||||
namespace: list[str] | None = None,
|
||||
graph_name: str | None = None,
|
||||
trigger_call_id: str | None = None,
|
||||
error: str | None = None,
|
||||
) -> ProtocolEvent:
|
||||
data: dict[str, Any] = {"event": event}
|
||||
if graph_name is not None:
|
||||
data["graph_name"] = graph_name
|
||||
if trigger_call_id is not None:
|
||||
data["trigger_call_id"] = trigger_call_id
|
||||
if error is not None:
|
||||
data["error"] = error
|
||||
return {
|
||||
"type": "event",
|
||||
"method": "lifecycle",
|
||||
"params": {
|
||||
"namespace": namespace or [],
|
||||
"timestamp": TS,
|
||||
"data": data,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _values(payload: dict[str, Any], *, namespace: list[str]) -> ProtocolEvent:
|
||||
return {
|
||||
"type": "event",
|
||||
"method": "values",
|
||||
"params": {
|
||||
"namespace": namespace,
|
||||
"timestamp": TS,
|
||||
"data": payload,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _subscribe(log: EventLog) -> None:
|
||||
"""Flip `_subscribed = True` so pushes retain items for test inspection."""
|
||||
log._subscribed = True
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Unit tests: feed events directly into the transformer
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
_FACTORIES = [ValuesTransformer, MessagesTransformer, SubgraphTransformer]
|
||||
|
||||
|
||||
def _handle_values_items(handle: SubgraphRunStream) -> list:
|
||||
return list(handle._mux.extensions["values"]._items) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def _handle_subgraphs_items(handle: SubgraphRunStream) -> list:
|
||||
return list(handle._mux.extensions["subgraphs"]._items) # type: ignore[attr-defined]
|
||||
|
||||
|
||||
def _pre_subscribe_handle(handle: SubgraphRunStream) -> None:
|
||||
"""Flip `_subscribed` on every EventLog inside the handle's mini-mux.
|
||||
|
||||
The mini-mux is built via `make_child` with the full factory list,
|
||||
so values / messages / subgraphs logs all exist as projections.
|
||||
Tests that feed events directly need them subscribed so pushes
|
||||
retain items in the deque for `_items` inspection.
|
||||
"""
|
||||
for value in handle._mux.extensions.values():
|
||||
if isinstance(value, EventLog):
|
||||
_subscribe(value)
|
||||
|
||||
|
||||
class TestSubgraphTransformerUnit:
|
||||
def _mux(self) -> tuple[StreamMux, SubgraphTransformer]:
|
||||
mux = StreamMux(factories=_FACTORIES, is_async=False)
|
||||
transformer = mux.transformer_by_key("subgraphs")
|
||||
assert isinstance(transformer, SubgraphTransformer)
|
||||
_subscribe(transformer._root_log)
|
||||
return mux, transformer
|
||||
|
||||
def _handle(self, transformer: SubgraphTransformer) -> SubgraphRunStream:
|
||||
"""Return the single root handle after pushing one lifecycle started."""
|
||||
(handle,) = list(transformer._root_log._items)
|
||||
return handle
|
||||
|
||||
def test_root_started_is_ignored(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(_lifecycle("started", graph_name="root"))
|
||||
assert list(transformer._root_log._items) == []
|
||||
assert transformer._by_ns == {}
|
||||
|
||||
def test_child_started_yields_handle(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(
|
||||
_lifecycle(
|
||||
"started",
|
||||
namespace=["task_a:child"],
|
||||
graph_name="child",
|
||||
trigger_call_id="task_a",
|
||||
)
|
||||
)
|
||||
|
||||
handle = self._handle(transformer)
|
||||
assert handle.path == ("task_a:child",)
|
||||
assert handle.graph_name == "child"
|
||||
assert handle.trigger_call_id == "task_a"
|
||||
assert handle.status == "started"
|
||||
|
||||
def test_status_transitions(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||
mux.push(_lifecycle("running", namespace=["t:c"]))
|
||||
mux.push(_lifecycle("completed", namespace=["t:c"]))
|
||||
|
||||
handle = self._handle(transformer)
|
||||
assert handle.status == "completed"
|
||||
|
||||
def test_grandchild_surfaces_under_child(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(_lifecycle("started", namespace=["t:child"], graph_name="child"))
|
||||
|
||||
child = self._handle(transformer)
|
||||
_pre_subscribe_handle(child)
|
||||
|
||||
mux.push(
|
||||
_lifecycle(
|
||||
"started",
|
||||
namespace=["t:child", "u:grand"],
|
||||
graph_name="grand",
|
||||
)
|
||||
)
|
||||
|
||||
(grand,) = _handle_subgraphs_items(child)
|
||||
assert grand.path == ("t:child", "u:grand")
|
||||
assert grand.graph_name == "grand"
|
||||
|
||||
def test_failed_stores_error(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||
mux.push(_lifecycle("failed", namespace=["t:c"], error="boom"))
|
||||
|
||||
handle = self._handle(transformer)
|
||||
assert handle.status == "failed"
|
||||
assert handle.error == "boom"
|
||||
|
||||
def test_values_routed_into_handle(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||
|
||||
handle = self._handle(transformer)
|
||||
_pre_subscribe_handle(handle)
|
||||
|
||||
mux.push(_values({"value": 1}, namespace=["t:c"]))
|
||||
mux.push(_values({"value": 2}, namespace=["t:c"]))
|
||||
|
||||
assert _handle_values_items(handle) == [{"value": 1}, {"value": 2}]
|
||||
assert handle.output == {"value": 2}
|
||||
|
||||
def test_root_values_not_routed(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||
handle = self._handle(transformer)
|
||||
_pre_subscribe_handle(handle)
|
||||
|
||||
# Values event at root namespace — must not leak into child handle.
|
||||
mux.push(_values({"value": "root"}, namespace=[]))
|
||||
assert _handle_values_items(handle) == []
|
||||
|
||||
def test_finalize_closes_dangling(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||
handle = self._handle(transformer)
|
||||
|
||||
mux.close()
|
||||
assert handle.status == "completed"
|
||||
assert handle._mux.extensions["values"]._closed
|
||||
assert handle._mux.extensions["subgraphs"]._closed
|
||||
|
||||
def test_fail_with_graph_interrupt_marks_interrupted(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||
handle = self._handle(transformer)
|
||||
|
||||
mux.fail(GraphInterrupt())
|
||||
assert handle.status == "interrupted"
|
||||
|
||||
def test_fail_with_generic_error_marks_failed(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||
handle = self._handle(transformer)
|
||||
|
||||
mux.fail(RuntimeError("explode"))
|
||||
assert handle.status == "failed"
|
||||
assert handle.error == "explode"
|
||||
|
||||
def test_duplicate_started_ignored(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="c"))
|
||||
mux.push(_lifecycle("started", namespace=["t:c"], graph_name="other"))
|
||||
|
||||
handles = list(transformer._root_log._items)
|
||||
assert len(handles) == 1
|
||||
assert handles[0].graph_name == "c"
|
||||
|
||||
def test_non_lifecycle_non_values_passthrough(self) -> None:
|
||||
mux, transformer = self._mux()
|
||||
mux.push(
|
||||
{
|
||||
"type": "event",
|
||||
"method": "messages",
|
||||
"params": {"namespace": ["t:c"], "timestamp": TS, "data": "x"},
|
||||
}
|
||||
)
|
||||
assert list(transformer._root_log._items) == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# End-to-end tests via StreamingHandler on real graphs
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class SimpleState(TypedDict):
|
||||
value: str
|
||||
items: Annotated[list[str], operator.add]
|
||||
|
||||
|
||||
def _build_nested_graph():
|
||||
"""Parent graph with a compiled subgraph node."""
|
||||
|
||||
def inner_node(state: SimpleState) -> dict:
|
||||
return {"value": state["value"] + "X", "items": ["x"]}
|
||||
|
||||
inner_builder = StateGraph(SimpleState)
|
||||
inner_builder.add_node("inner_node", inner_node)
|
||||
inner_builder.add_edge(START, "inner_node")
|
||||
inner_builder.add_edge("inner_node", END)
|
||||
inner = inner_builder.compile()
|
||||
|
||||
def outer_node(state: SimpleState) -> dict:
|
||||
return {"value": state["value"] + "Y", "items": ["y"]}
|
||||
|
||||
outer_builder = StateGraph(SimpleState)
|
||||
outer_builder.add_node("outer_node", outer_node)
|
||||
outer_builder.add_node("sub", inner)
|
||||
outer_builder.add_edge(START, "outer_node")
|
||||
outer_builder.add_edge("outer_node", "sub")
|
||||
outer_builder.add_edge("sub", END)
|
||||
return outer_builder.compile()
|
||||
|
||||
|
||||
class TestSubgraphTransformerEndToEnd:
|
||||
def test_flat_graph_yields_no_subgraphs(self) -> None:
|
||||
builder = StateGraph(SimpleState)
|
||||
builder.add_node("n", lambda s: {"value": s["value"] + "!", "items": ["!"]})
|
||||
builder.add_edge(START, "n")
|
||||
builder.add_edge("n", END)
|
||||
graph = builder.compile()
|
||||
|
||||
handler = StreamingHandler(graph)
|
||||
run = handler.stream({"value": "", "items": []})
|
||||
|
||||
collected: list[SubgraphRunStream] = []
|
||||
for sub in run.subgraphs:
|
||||
collected.append(sub)
|
||||
assert collected == []
|
||||
# Output still resolves.
|
||||
assert run.output is not None
|
||||
|
||||
def test_nested_graph_yields_one_child(self) -> None:
|
||||
graph = _build_nested_graph()
|
||||
handler = StreamingHandler(graph)
|
||||
run = handler.stream({"value": "", "items": []})
|
||||
|
||||
collected: list[SubgraphRunStream] = []
|
||||
for sub in run.subgraphs:
|
||||
collected.append(sub)
|
||||
|
||||
assert len(collected) == 1
|
||||
child = collected[0]
|
||||
assert len(child.path) == 1
|
||||
assert child.path[0].startswith("sub:")
|
||||
assert child.status == "completed"
|
||||
|
||||
def test_error_in_subgraph_fails_child(self) -> None:
|
||||
def boom(state: SimpleState) -> dict:
|
||||
raise RuntimeError("subgraph_failed")
|
||||
|
||||
inner_builder = StateGraph(SimpleState)
|
||||
inner_builder.add_node("inner", boom)
|
||||
inner_builder.add_edge(START, "inner")
|
||||
inner_builder.add_edge("inner", END)
|
||||
inner = inner_builder.compile()
|
||||
|
||||
outer_builder = StateGraph(SimpleState)
|
||||
outer_builder.add_node("sub", inner)
|
||||
outer_builder.add_edge(START, "sub")
|
||||
outer_builder.add_edge("sub", END)
|
||||
graph = outer_builder.compile()
|
||||
|
||||
handler = StreamingHandler(graph)
|
||||
run = handler.stream({"value": "", "items": []})
|
||||
|
||||
collected: list[SubgraphRunStream] = []
|
||||
with pytest.raises(RuntimeError):
|
||||
for sub in run.subgraphs:
|
||||
collected.append(sub)
|
||||
|
||||
assert len(collected) == 1
|
||||
assert collected[0].status == "failed"
|
||||
|
||||
|
||||
class TestSubgraphTransformerAsyncEndToEnd:
|
||||
@pytest.mark.anyio
|
||||
async def test_nested_graph_yields_one_child(self) -> None:
|
||||
async def inner(state: SimpleState) -> dict:
|
||||
return {"value": state["value"] + "X", "items": ["x"]}
|
||||
|
||||
inner_builder = StateGraph(SimpleState)
|
||||
inner_builder.add_node("inner", inner)
|
||||
inner_builder.add_edge(START, "inner")
|
||||
inner_builder.add_edge("inner", END)
|
||||
inner_graph = inner_builder.compile()
|
||||
|
||||
outer_builder = StateGraph(SimpleState)
|
||||
outer_builder.add_node("sub", inner_graph)
|
||||
outer_builder.add_edge(START, "sub")
|
||||
outer_builder.add_edge("sub", END)
|
||||
graph = outer_builder.compile()
|
||||
|
||||
handler = StreamingHandler(graph)
|
||||
run = await handler.astream({"value": "", "items": []})
|
||||
|
||||
collected: list[SubgraphRunStream] = []
|
||||
async for sub in run.subgraphs:
|
||||
collected.append(sub)
|
||||
|
||||
assert len(collected) == 1
|
||||
child = collected[0]
|
||||
assert child.status == "completed"
|
||||
@@ -869,14 +869,18 @@ class TestStreamMux:
|
||||
|
||||
class TestValuesTransformer:
|
||||
def test_ignores_non_root_namespace(self) -> None:
|
||||
"""Values events from subgraphs (non-empty namespace) should be ignored."""
|
||||
t = ValuesTransformer()
|
||||
t.init()
|
||||
t._log._bind(is_async=False)
|
||||
"""The root mux only dispatches root-ns values events to its ValuesTransformer.
|
||||
|
||||
Namespace filtering is enforced by the mux via `scope_exact`
|
||||
— the transformer itself no longer filters.
|
||||
"""
|
||||
mux = StreamMux([ValuesTransformer()], is_async=False)
|
||||
t = mux.transformer_by_key("values")
|
||||
assert isinstance(t, ValuesTransformer)
|
||||
it = iter(t._log)
|
||||
|
||||
t.process(_event("values", {"val": "root"}))
|
||||
t.process(_event("values", {"val": "sub"}, namespace=["sub"]))
|
||||
mux.push(_event("values", {"val": "root"}))
|
||||
mux.push(_event("values", {"val": "sub"}, namespace=["sub"]))
|
||||
|
||||
t._log.close()
|
||||
items = list(it)
|
||||
@@ -935,14 +939,15 @@ class TestMessagesTransformer:
|
||||
assert items[0].message_id == "run-1"
|
||||
|
||||
def test_ignores_non_root_namespace(self) -> None:
|
||||
t = MessagesTransformer()
|
||||
t.init()
|
||||
t._log._bind(is_async=False)
|
||||
"""Namespace filtering is enforced by the mux via `scope_exact`."""
|
||||
mux = StreamMux([MessagesTransformer()], is_async=False)
|
||||
t = mux.transformer_by_key("messages")
|
||||
assert isinstance(t, MessagesTransformer)
|
||||
t._bind_pump(lambda: False)
|
||||
it = iter(t._log)
|
||||
|
||||
meta = {"langgraph_node": "llm", "run_id": "run-1"}
|
||||
t.process(
|
||||
mux.push(
|
||||
_event(
|
||||
"messages",
|
||||
({"event": "message-start", "message_id": "run-1"}, meta),
|
||||
|
||||
Reference in New Issue
Block a user