WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-17 21:25:46 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor(langgraph,prebuilt): merge EventLog into StreamChannel with optional name (#7637)

Browse Source
This commit is contained in:
Nick Hollon
2026-04-28 18:43:51 -04:00
committed by GitHub
parent f4388df77f
commit 5af4c5addf
13 changed files with 435 additions and 544 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/stream/__init__.py
-2
View File
@@ -5,7 +5,6 @@ Compile a graph with `transformers=[...]` and call `graph.stream_v2()` /
graph's raw events into ergonomic per-channel streams.
"""
from langgraph.stream._event_log import EventLog
from langgraph.stream._types import ProtocolEvent, StreamTransformer
from langgraph.stream.run_stream import (
AsyncGraphRunStream,
@@ -24,7 +23,6 @@ from langgraph.stream.transformers import (
__all__ = [
"AsyncGraphRunStream",
"AsyncSubgraphRunStream",
"EventLog",
"GraphRunStream",
"LifecyclePayload",
"LifecycleTransformer",
libs/langgraph/langgraph/stream/_event_log.py
-306
View File
@@ -1,306 +0,0 @@
from __future__ import annotations
import asyncio
from collections import deque
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
from typing import Generic, TypeVar
T = TypeVar("T")
class EventLog(Generic[T]):
"""Single-consumer drainable queue for streaming events.
Items are popped off the front as the consumer advances — there is
no retention beyond what's currently queued. A log accepts exactly
one subscriber; a second `__iter__` / `__aiter__` call raises. Use
`tee(n)` / `atee(n)` for fan-out.
Starts unbound — neither `__iter__` nor `__aiter__` is available
until the StreamMux calls `_bind(is_async)`. After binding, only
the matching iteration protocol works; the other raises `TypeError`.
Pump wiring (set by the run stream, not by `_bind`):
- `_request_more`: sync pump callable, returns True if a new
event was produced.
- `_arequest_more`: async pump coroutine factory, same contract.
Memory is bounded by caller pace: both sync and async use caller-
driven pumps, so each cursor advance produces at most one event.
The only shape where a log can accumulate meaningfully is
concurrent async consumers at unequal rates — a slow consumer's
log grows while fast consumers drive the shared pump. That's the
documented tradeoff for concurrent consumption; consume at similar
rates or use a single consumer if memory matters.
Lazy-subscribe: `push` is a no-op when no subscriber has registered.
Transformers still execute `process()` (so scalar state like
`ValuesTransformer._latest` stays current); only the log append is
skipped.
"""
def __init__(self, maxlen: int | None = None) -> None:
"""Initialize an empty, unbound log.
Args:
maxlen: Accepted for forward compatibility; currently unused.
The caller-driven pump bounds memory naturally for
single-consumer use.
Raises:
ValueError: If `maxlen` is not a positive integer or `None`.
"""
if maxlen is not None and maxlen <= 0:
raise ValueError("EventLog maxlen must be a positive int or None")
self._items: deque[T] = deque()
self._maxlen: int | None = maxlen
self._closed = False
self._error: BaseException | None = None
# Binding state — None means unbound.
self._is_async: bool | None = None
# Flipped on first __iter__ / __aiter__. Pre-subscription
# pushes are silent no-ops.
self._subscribed = False
# Pump wiring set by the run stream after bind.
self._request_more: Callable[[], bool] | None = None
self._arequest_more: Callable[[], Awaitable[bool]] | None = None
# ------------------------------------------------------------------
# Binding
# ------------------------------------------------------------------
def _bind(self, *, is_async: bool) -> None:
"""Bind this log to sync or async mode.
Called by the StreamMux after transformer registration. Must be
called exactly once before any iteration.
Args:
is_async: True to enable async iteration, False for sync.
Raises:
RuntimeError: If the log has already been bound.
"""
if self._is_async is not None:
raise RuntimeError("EventLog is already bound")
self._is_async = is_async
# ------------------------------------------------------------------
# Producer API
# ------------------------------------------------------------------
def push(self, item: T) -> None:
"""Append an item. No-op when no subscriber is registered.
Non-blocking in both sync and async — matches v1's
`put_nowait` producer shape. Memory is bounded by caller pace
via the caller-driven pump.
Raises:
RuntimeError: If the log is closed (and subscribed).
"""
if not self._subscribed:
return
if self._closed:
raise RuntimeError("Cannot push to a closed EventLog")
self._items.append(item)
def close(self) -> None:
"""Mark the log as complete."""
self._closed = True
def fail(self, err: BaseException) -> None:
"""Mark the log as errored.
Args:
err: The exception to surface to the subscriber.
"""
self._error = err
self._closed = True
# ------------------------------------------------------------------
# Sync iteration (caller-driven pump)
# ------------------------------------------------------------------
def __iter__(self) -> Iterator[T]:
"""Subscribe and return a sync cursor. Can be called only once.
Raises:
TypeError: If the log is unbound or bound to async mode.
RuntimeError: If the log already has a subscriber.
"""
if self._is_async is None:
raise TypeError(
"EventLog has not been bound yet. "
"Register the transformer with a StreamMux first."
)
if self._is_async:
raise TypeError(
"This EventLog is bound to async mode — use 'async for' instead."
)
if self._subscribed:
raise RuntimeError(
"EventLog already has a subscriber; use .tee(n) for fan-out."
)
self._subscribed = True
return self._sync_cursor()
def _sync_cursor(self) -> Iterator[T]:
while True:
if self._items:
yield self._items.popleft()
elif self._closed:
if self._error is not None:
raise self._error
return
elif self._request_more is not None:
if not self._request_more():
if not self._items and not self._closed:
return
else:
return
# ------------------------------------------------------------------
# Async iteration (caller-driven pump)
# ------------------------------------------------------------------
def __aiter__(self) -> AsyncIterator[T]:
"""Subscribe and return an async cursor. Can be called only once.
Raises:
TypeError: If the log is unbound or bound to sync mode.
RuntimeError: If the log already has a subscriber.
"""
if self._is_async is None:
raise TypeError(
"EventLog has not been bound yet. "
"Register the transformer with a StreamMux first."
)
if not self._is_async:
raise TypeError("This EventLog is bound to sync mode — use 'for' instead.")
if self._subscribed:
raise RuntimeError(
"EventLog already has a subscriber; use .atee(n) for fan-out."
)
self._subscribed = True
return self._async_cursor()
async def _async_cursor(self) -> AsyncIterator[T]:
while True:
if self._items:
yield self._items.popleft()
elif self._closed:
if self._error is not None:
raise self._error
return
elif self._arequest_more is not None:
if not await self._arequest_more():
if not self._items and not self._closed:
return
else:
return
# ------------------------------------------------------------------
# Fan-out via tee
# ------------------------------------------------------------------
def tee(self, n: int = 2) -> tuple[Iterator[T], ...]:
"""Subscribe and return `n` independent sync iterators.
Each branch has its own buffer; items pulled from the
underlying cursor are copied into every branch. Branches are
naturally bounded by caller pace since the sync pump is
caller-driven.
Args:
n: Number of branches to create. Must be >= 1.
Returns:
A tuple of `n` iterators over the same underlying stream.
Raises:
TypeError: If the log is unbound or bound to async mode.
RuntimeError: If the log already has a subscriber.
ValueError: If `n` < 1.
"""
if n < 1:
raise ValueError("tee() requires n >= 1")
source = self.__iter__()
buffers: list[deque[T]] = [deque() for _ in range(n)]
exhausted = [False]
def branch(i: int) -> Iterator[T]:
buf = buffers[i]
while True:
if buf:
yield buf.popleft()
elif exhausted[0]:
return
else:
try:
item = next(source)
except StopIteration:
exhausted[0] = True
return
for b in buffers:
b.append(item)
return tuple(branch(i) for i in range(n))
def atee(self, n: int = 2) -> tuple[AsyncIterator[T], ...]:
"""Subscribe and return `n` independent async iterators.
Caller-driven fan-out: each branch's `__anext__` either pops
from its own buffer or, under a shared `asyncio.Lock`, pulls
one item from the underlying cursor and distributes it to
every branch's buffer.
Args:
n: Number of branches to create. Must be >= 1.
Returns:
A tuple of `n` async iterators over the same underlying
stream.
Raises:
TypeError: If the log is unbound or bound to sync mode.
RuntimeError: If the log already has a subscriber.
ValueError: If `n` < 1.
"""
if n < 1:
raise ValueError("atee() requires n >= 1")
source = self.__aiter__()
buffers: list[deque[T]] = [deque() for _ in range(n)]
exhausted = [False]
error: list[BaseException | None] = [None]
lock = asyncio.Lock()
async def branch(i: int) -> AsyncIterator[T]:
buf = buffers[i]
while True:
if buf:
yield buf.popleft()
continue
if exhausted[0]:
if error[0] is not None:
raise error[0]
return
async with lock:
if buf or exhausted[0]:
continue
try:
item = await source.__anext__()
except StopAsyncIteration:
exhausted[0] = True
continue
except Exception as e:
error[0] = e
exhausted[0] = True
continue
for b in buffers:
b.append(item)
return tuple(branch(i) for i in range(n))
libs/langgraph/langgraph/stream/_mux.py
+51 -70
View File
@@ -5,7 +5,6 @@ import time
from collections.abc import Awaitable, Callable
from typing import Any
from langgraph.stream._event_log import EventLog
from langgraph.stream._types import (
ProtocolEvent,
StreamTransformer,
@@ -28,14 +27,15 @@ class StreamMux:
"""Central event dispatcher for the streaming infrastructure.
Owns the main event log and routes events through a transformer
pipeline. StreamChannels discovered in transformer projections are
auto-wired so that every `push()` also injects a `ProtocolEvent`
into the main log.
pipeline. StreamChannels with a name discovered in transformer
projections are auto-wired so that every `push()` also injects a
`ProtocolEvent` into the main log. StreamChannels without a name
are local-only.
Pass `is_async=True` when the mux will be consumed via async
iteration (`handler.astream()`). All EventLog and StreamChannel
instances discovered during registration are automatically bound
to the matching mode.
iteration (`handler.astream()`). All StreamChannel instances
discovered during registration are automatically bound to the
matching mode.
Attributes:
extensions: Merged projection dict across all registered
@@ -60,7 +60,7 @@ class StreamMux:
`factories` (callables producing fresh instances per mux). Each
transformer's `init()` is called, projections are merged into
`extensions`, `_native` keys are recorded in `native_keys`, and
any EventLog / StreamChannel instances are bound and wired.
any StreamChannel instances are bound and (if named) wired.
Args:
transformers: Already-built transformer instances. Registered
@@ -89,11 +89,10 @@ class StreamMux:
self.is_async = is_async
self.scope: tuple[str, ...] = scope
self._assign_seq = _assign_seq
self._events: EventLog[ProtocolEvent] = EventLog()
self._events: StreamChannel[ProtocolEvent] = StreamChannel()
self._events._bind(is_async=is_async)
self._transformers: list[StreamTransformer] = []
self._channels: list[StreamChannel[Any]] = []
self._logs: list[EventLog[Any]] = []
self._seq = 0
self.extensions: dict[str, Any] = {}
@@ -135,18 +134,15 @@ class StreamMux:
Records the pump on the mux so child mini-muxes built by
`_make_child` can inherit it. Propagates to:
- the main event log (`self._events`)
- every projection EventLog and StreamChannel in `extensions`
- every projection StreamChannel in `extensions`
- any registered transformer that exposes `_bind_pump` (e.g.
`MessagesTransformer` so `ChatModelStream` instances drive the
shared pump from their 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 ch in self._channels:
ch._request_more = fn
for transformer in self._transformers:
bind = getattr(transformer, "_bind_pump", None)
if bind is not None:
@@ -156,11 +152,8 @@ class StreamMux:
"""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 ch in self._channels:
ch._arequest_more = fn
for transformer in self._transformers:
abind = getattr(transformer, "_bind_apump", None)
if abind is not None:
@@ -204,8 +197,8 @@ class StreamMux:
"""Register a single transformer.
Calls `transformer.init()`, stores the transformer for event
processing, binds any EventLog or StreamChannel instances in
the projection, and merges the projection into `extensions`.
processing, binds any StreamChannel instances in the projection,
and merges the projection into `extensions`.
"""
if transformer_requires_async(transformer) and not self.is_async:
raise RuntimeError(
@@ -274,12 +267,12 @@ class StreamMux:
def close(self) -> None:
"""Finalize all transformers, close all projections and the main log.
EventLogs and StreamChannels discovered in transformer
projections are auto-closed after `finalize()` runs —
transformers don't need to close them manually. If any
transformer's `finalize()` raises, the remaining transformers,
projections, and the main log are still closed; the first error
is re-raised after cleanup completes.
StreamChannels discovered in transformer projections are
auto-closed after `finalize()` runs — transformers don't need
to close them manually. If any transformer's `finalize()` raises,
the remaining transformers, projections, and the main log are
still closed; the first error is re-raised after cleanup
completes.
Raises:
BaseException: The first error raised by a transformer's
@@ -292,12 +285,9 @@ class StreamMux:
except BaseException as e:
if first_error is None:
first_error = e
for log in self._logs:
if not log._closed:
log.close()
for ch in self._channels:
if not ch._log._closed:
ch._close()
if not ch._closed:
ch.close()
self._events.close()
if first_error is not None:
raise first_error
@@ -305,11 +295,10 @@ class StreamMux:
def fail(self, err: BaseException) -> None:
"""Fail all transformers, projections, and the main log.
EventLogs and StreamChannels discovered in transformer
projections are auto-failed — transformers don't need to fail
them manually. If any transformer's `fail()` raises, the
remaining transformers, projections, and the main log are still
failed.
StreamChannels discovered in transformer projections are
auto-failed — transformers don't need to fail them manually.
If any transformer's `fail()` raises, the remaining
transformers, projections, and the main log are still failed.
Args:
err: The exception that ended the run.
@@ -319,12 +308,9 @@ class StreamMux:
transformer.fail(err)
except BaseException:
pass
for log in self._logs:
if not log._closed:
log.fail(err)
for ch in self._channels:
if not ch._log._closed:
ch._fail(err)
if not ch._closed:
ch.fail(err)
self._events.fail(err)
# ------------------------------------------------------------------
@@ -345,7 +331,7 @@ class StreamMux:
`put_nowait` shape. The root mux assigns `seq`; child muxes do
not, so forwarded subgraph events can be shared without copying.
Memory is bounded by caller pace via the caller-driven pump; see
`EventLog` for the full tradeoff story.
`StreamChannel` for the full tradeoff story.
Args:
event: The protocol event to dispatch.
@@ -365,7 +351,7 @@ class StreamMux:
Awaits every task started via `StreamTransformer.schedule()`
across all transformers, then calls `afinalize()` on each,
then auto-closes logs, channels, and the main event log.
then auto-closes channels and the main event log.
If any scheduled task raised under `on_error="raise"`, or any
transformer's `afinalize` raises, the exception propagates.
@@ -397,12 +383,9 @@ class StreamMux:
except BaseException as e:
if first_error is None:
first_error = e
for log in self._logs:
if not log._closed:
log.close()
for ch in self._channels:
if not ch._log._closed:
ch._close()
if not ch._closed:
ch.close()
self._events.close()
if first_error is not None:
raise first_error
@@ -412,7 +395,7 @@ class StreamMux:
Cancels every scheduled task across all transformers, awaits
them to completion, then runs each transformer's `afail` hook
and auto-fails logs, channels, and the main event log.
and auto-fails channels and the main event log.
Args:
err: The exception that ended the run.
@@ -428,12 +411,9 @@ class StreamMux:
await transformer.afail(err)
except BaseException:
pass
for log in self._logs:
if not log._closed:
log.fail(err)
for ch in self._channels:
if not ch._log._closed:
ch._fail(err)
if not ch._closed:
ch.fail(err)
if not self._events._closed:
self._events.fail(err)
@@ -453,32 +433,33 @@ class StreamMux:
def _bind_and_wire(
self, projection: dict[str, Any], *, native: bool = False
) -> None:
"""Bind and wire EventLog / StreamChannel instances in a projection.
"""Bind and optionally wire StreamChannel instances in a projection.
All StreamChannels are bound and tracked. Channels with a name
are additionally wired for protocol auto-forwarding.
Args:
projection: The projection dict returned by a transformer's
`init()`.
native: True when the owning transformer is `_native`.
Channels owned by a native transformer use the channel
name directly as the protocol method; user-defined
channels are prefixed with `custom:`.
Named channels owned by a native transformer use the
channel name directly as the protocol method;
user-defined channels are prefixed with `custom:`.
"""
for value in projection.values():
if isinstance(value, StreamChannel):
value._bind(is_async=self.is_async)
self._channels.append(value)
method = value.name if native else f"custom:{value.name}"
if value.name is not None:
method = value.name if native else f"custom:{value.name}"
def _make_forward(method_name: str) -> Callable[[Any], None]:
def _forward(item: Any) -> None:
self._forward(method_name, item)
def _make_forward(method_name: str) -> Callable[[Any], None]:
def _forward(item: Any) -> None:
self._forward(method_name, item)
return _forward
return _forward
value._wire(_make_forward(method))
elif isinstance(value, EventLog):
value._bind(is_async=self.is_async)
self._logs.append(value)
value._wire(_make_forward(method))
def _forward(self, method: str, item: Any) -> None:
"""Inject a ProtocolEvent for a StreamChannel push.
libs/langgraph/langgraph/stream/_types.py
+11 -10
View File
@@ -45,8 +45,7 @@ class StreamTransformer(ABC):
"""Extension point for custom stream projections.
Transformers observe protocol events flowing through the StreamMux and
build typed derived projections (EventLogs, StreamChannels, promises,
etc.).
build typed derived projections (StreamChannels, promises, etc.).
Set `_native = True` on a transformer to have its projection keys
exposed as direct attributes on the run stream (in addition to
@@ -55,9 +54,9 @@ class StreamTransformer(ABC):
Subclasses must implement `init` and override at least one of
`process` / `aprocess`. The `finalize` / `afinalize` and `fail` /
`afail` hooks are optional — the default implementations are no-ops.
EventLog and StreamChannel instances in the projection dict are
auto-closed / auto-failed by the mux, so most transformers don't
need `finalize` or `fail` at all.
StreamChannel instances in the projection dict are auto-closed /
auto-failed by the mux, so most transformers don't need `finalize`
or `fail` at all.
Transformers that need async work pick the async lane by:
@@ -170,15 +169,16 @@ class StreamTransformer(ABC):
def finalize(self) -> None:
"""Called when the run ends normally (sync lane).
Override to close EventLogs, resolve promises, or perform other
teardown. StreamChannel instances are auto-closed by the mux.
Override to close StreamChannels, resolve promises, or perform
other teardown. StreamChannel instances in the projection dict
are auto-closed by the mux.
"""
async def afinalize(self) -> None:
"""Called when the run ends normally (async lane).
By the time this runs, the mux has already awaited every task
started via `schedule()`, so EventLogs can be closed here
started via `schedule()`, so StreamChannels can be closed here
without a last-task-wins race.
The default delegates to `finalize`.
@@ -188,8 +188,9 @@ class StreamTransformer(ABC):
def fail(self, err: BaseException) -> None:
"""Called when the run ends with an error (sync lane).
Override to fail EventLogs, reject promises, or perform other
teardown. StreamChannel instances are auto-failed by the mux.
Override to fail StreamChannels, reject promises, or perform
other teardown. StreamChannel instances in the projection dict
are auto-failed by the mux.
Args:
err: The exception that ended the run.
libs/langgraph/langgraph/stream/stream_channel.py
+277 -63
View File
@@ -1,113 +1,327 @@
from __future__ import annotations
from collections.abc import AsyncIterator, Callable, Iterator
import asyncio
from collections import deque
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
from typing import Generic, TypeVar
from langgraph.stream._event_log import EventLog
T = TypeVar("T")
class StreamChannel(Generic[T]):
"""A named projection channel with optional protocol auto-forwarding.
"""Single-consumer drainable queue for streaming events, with optional
protocol auto-forwarding.
Wraps an event log and declares a protocol channel name. When the
StreamMux detects a StreamChannel in a transformer's `init()`
return value, it automatically wires every `push()` to inject a
`ProtocolEvent` into the main event stream using the channel's
name as the method.
When constructed with a `name`, the StreamMux auto-wires every
`push()` to also inject a `ProtocolEvent` into the main event stream
using the channel's name as the method. When constructed without a
name, the channel is local-only — items are only visible to
in-process consumers that iterate the channel directly.
Auto-forwarded events bypass the transformer pipeline — other
transformers' `process()` / `aprocess()` methods do not see
`custom:<name>` events produced by a channel push. This prevents a
transformer that pushes to its own channel during `process()` from
re-triggering itself, but it also means filter- or tap-style
transformers cannot observe channel output from peer transformers.
Consumers that need that should iterate the main event stream.
Items are popped off the front as the consumer advances — there is
no retention beyond what's currently queued. A channel accepts
exactly one subscriber; a second `__iter__` / `__aiter__` call
raises. Use `tee(n)` / `atee(n)` for fan-out.
In-process consumers iterate the channel directly (`for item in ch`
or `async for item in ch`). Remote SDK clients subscribe via
`session.subscribe("custom:<channelName>")`.
Starts unbound — neither `__iter__` nor `__aiter__` is available
until the StreamMux calls `_bind(is_async)`. After binding, only
the matching iteration protocol works; the other raises `TypeError`.
Like EventLog, a StreamChannel starts unbound. The mux calls
`_bind(is_async)` during registration so the correct iteration
protocol is available by the time user code sees it.
Pump wiring (set by the run stream, not by `_bind`):
- `_request_more`: sync pump callable, returns True if a new
event was produced.
- `_arequest_more`: async pump coroutine factory, same contract.
Lifecycle (`_close` / `_fail`) is managed by the mux — transformers
using only StreamChannels don't need `finalize` or `fail` hooks.
Memory is bounded by caller pace: both sync and async use caller-
driven pumps, so each cursor advance produces at most one event.
Lazy-subscribe: `push` appends to the local buffer only when a
subscriber has registered. Auto-forward via `_wire_fn` always fires
regardless of subscription state.
Lifecycle (`close` / `fail`) is managed by the mux — transformers
don't need to close their channels manually.
"""
def __init__(self, name: str, *, maxlen: int | None = None) -> None:
"""Initialize the channel with an empty inner log.
def __init__(self, name: str | None = None, *, maxlen: int | None = None) -> None:
"""Initialize the channel.
Args:
name: The protocol channel name used for auto-forwarded
events. Surfaced on the wire as `custom:<name>` for
user-defined transformers, or as `<name>` for channels
owned by a native transformer (`_native = True`). The
StreamMux derives the prefix from the owning
transformer at registration time.
maxlen: Optional retention cap on the inner EventLog. See
`EventLog.__init__` for semantics.
name: Optional protocol channel name. When set, the
StreamMux wires every `push()` to also inject a
`ProtocolEvent` into the main event stream. Surfaced
on the wire as `custom:<name>` for user-defined
transformers, or as `<name>` for channels owned by a
native transformer (`_native = True`). When `None`,
the channel is local-only.
maxlen: Accepted for forward compatibility; currently
unused. The caller-driven pump bounds memory naturally
for single-consumer use.
Raises:
ValueError: If `maxlen` is not a positive integer or `None`.
"""
if maxlen is not None and maxlen <= 0:
raise ValueError("StreamChannel maxlen must be a positive int or None")
self.name = name
self._log: EventLog[T] = EventLog(maxlen=maxlen)
self._items: deque[T] = deque()
self._maxlen: int | None = maxlen
self._closed = False
self._error: BaseException | None = None
self._is_async: bool | None = None
self._subscribed = False
self._request_more: Callable[[], bool] | None = None
self._arequest_more: Callable[[], Awaitable[bool]] | None = None
self._wire_fn: Callable[[T], None] | None = None
# ------------------------------------------------------------------
# Binding
# ------------------------------------------------------------------
def _bind(self, *, is_async: bool) -> None:
"""Bind the underlying event log to sync or async mode.
"""Bind this channel to sync or async mode.
Called by the StreamMux after transformer registration. Must be
called exactly once before any iteration.
Args:
is_async: True for async iteration, False for sync.
"""
self._log._bind(is_async=is_async)
is_async: True to enable async iteration, False for sync.
def push(self, item: T) -> None:
"""Append an item to the log and auto-forward if wired.
Args:
item: The item to push.
Raises:
RuntimeError: If the channel has already been bound.
"""
self._log.push(item)
if self._wire_fn is not None:
self._wire_fn(item)
if self._is_async is not None:
raise RuntimeError("StreamChannel is already bound")
self._is_async = is_async
# ------------------------------------------------------------------
# Mux lifecycle hooks (not called by transformers directly)
# Mux wiring (not called by transformers directly)
# ------------------------------------------------------------------
def _wire(self, fn: Callable[[T], None]) -> None:
"""Install the auto-forward callback (called by StreamMux)."""
self._wire_fn = fn
def _close(self) -> None:
"""Close the underlying log (called by StreamMux on run end)."""
self._log.close()
# ------------------------------------------------------------------
# Producer API
# ------------------------------------------------------------------
def _fail(self, err: BaseException) -> None:
"""Fail the underlying log (called by StreamMux on run error)."""
self._log.fail(err)
def push(self, item: T) -> None:
"""Append an item. Auto-forwards if wired.
The local buffer append is a no-op when no subscriber is
registered, but auto-forwarding always fires so wired events
reach the main event log regardless of subscription state.
Raises:
RuntimeError: If the channel is closed (and subscribed).
"""
if self._subscribed:
if self._closed:
raise RuntimeError("Cannot push to a closed StreamChannel")
self._items.append(item)
if self._wire_fn is not None:
self._wire_fn(item)
def close(self) -> None:
"""Mark the channel as complete."""
self._closed = True
def fail(self, err: BaseException) -> None:
"""Mark the channel as errored.
Args:
err: The exception to surface to the subscriber.
"""
self._error = err
self._closed = True
# ------------------------------------------------------------------
# Iteration — delegates to the inner event log (multi-cursor)
# Sync iteration (caller-driven pump)
# ------------------------------------------------------------------
def __iter__(self) -> Iterator[T]:
return iter(self._log)
"""Subscribe and return a sync cursor. Can be called only once.
Raises:
TypeError: If the channel is unbound or bound to async mode.
RuntimeError: If the channel already has a subscriber.
"""
if self._is_async is None:
raise TypeError(
"StreamChannel has not been bound yet. "
"Register the transformer with a StreamMux first."
)
if self._is_async:
raise TypeError(
"This StreamChannel is bound to async mode — use 'async for' instead."
)
if self._subscribed:
raise RuntimeError(
"StreamChannel already has a subscriber; use .tee(n) for fan-out."
)
self._subscribed = True
return self._sync_cursor()
def _sync_cursor(self) -> Iterator[T]:
while True:
if self._items:
yield self._items.popleft()
elif self._closed:
if self._error is not None:
raise self._error
return
elif self._request_more is not None:
if not self._request_more():
if not self._items and not self._closed:
return
else:
return
# ------------------------------------------------------------------
# Async iteration (caller-driven pump)
# ------------------------------------------------------------------
def __aiter__(self) -> AsyncIterator[T]:
return self._log.__aiter__()
"""Subscribe and return an async cursor. Can be called only once.
Raises:
TypeError: If the channel is unbound or bound to sync mode.
RuntimeError: If the channel already has a subscriber.
"""
if self._is_async is None:
raise TypeError(
"StreamChannel has not been bound yet. "
"Register the transformer with a StreamMux first."
)
if not self._is_async:
raise TypeError(
"This StreamChannel is bound to sync mode — use 'for' instead."
)
if self._subscribed:
raise RuntimeError(
"StreamChannel already has a subscriber; use .atee(n) for fan-out."
)
self._subscribed = True
return self._async_cursor()
async def _async_cursor(self) -> AsyncIterator[T]:
while True:
if self._items:
yield self._items.popleft()
elif self._closed:
if self._error is not None:
raise self._error
return
elif self._arequest_more is not None:
if not await self._arequest_more():
if not self._items and not self._closed:
return
else:
return
# ------------------------------------------------------------------
# Fan-out via tee
# ------------------------------------------------------------------
def tee(self, n: int = 2) -> tuple[Iterator[T], ...]:
"""Fan out the channel into `n` independent sync iterators.
"""Subscribe and return `n` independent sync iterators.
Delegates to the underlying EventLog's `tee()`.
Each branch has its own buffer; items pulled from the
underlying cursor are copied into every branch. Branches are
naturally bounded by caller pace since the sync pump is
caller-driven.
Args:
n: Number of branches to create. Must be >= 1.
Returns:
A tuple of `n` iterators over the same underlying stream.
Raises:
TypeError: If the channel is unbound or bound to async mode.
RuntimeError: If the channel already has a subscriber.
ValueError: If `n` < 1.
"""
return self._log.tee(n)
if n < 1:
raise ValueError("tee() requires n >= 1")
source = self.__iter__()
buffers: list[deque[T]] = [deque() for _ in range(n)]
exhausted = [False]
def branch(i: int) -> Iterator[T]:
buf = buffers[i]
while True:
if buf:
yield buf.popleft()
elif exhausted[0]:
return
else:
try:
item = next(source)
except StopIteration:
exhausted[0] = True
return
for b in buffers:
b.append(item)
return tuple(branch(i) for i in range(n))
def atee(self, n: int = 2) -> tuple[AsyncIterator[T], ...]:
"""Fan out the channel into `n` independent async iterators.
"""Subscribe and return `n` independent async iterators.
Delegates to the underlying EventLog's `atee()`.
Caller-driven fan-out: each branch's `__anext__` either pops
from its own buffer or, under a shared `asyncio.Lock`, pulls
one item from the underlying cursor and distributes it to
every branch's buffer.
Args:
n: Number of branches to create. Must be >= 1.
Returns:
A tuple of `n` async iterators over the same underlying
stream.
Raises:
TypeError: If the channel is unbound or bound to sync mode.
RuntimeError: If the channel already has a subscriber.
ValueError: If `n` < 1.
"""
return self._log.atee(n)
if n < 1:
raise ValueError("atee() requires n >= 1")
source = self.__aiter__()
buffers: list[deque[T]] = [deque() for _ in range(n)]
exhausted = [False]
error: list[BaseException | None] = [None]
lock = asyncio.Lock()
async def branch(i: int) -> AsyncIterator[T]:
buf = buffers[i]
while True:
if buf:
yield buf.popleft()
continue
if exhausted[0]:
if error[0] is not None:
raise error[0]
return
async with lock:
if buf or exhausted[0]:
continue
try:
item = await source.__anext__()
except StopAsyncIteration:
exhausted[0] = True
continue
except Exception as e:
error[0] = e
exhausted[0] = True
continue
for b in buffers:
b.append(item)
return tuple(branch(i) for i in range(n))
libs/langgraph/langgraph/stream/transformers.py
+5 -4
View File
@@ -13,7 +13,6 @@ from langchain_protocol.protocol import MessagesData
from typing_extensions import NotRequired, TypedDict
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 AsyncSubgraphRunStream, SubgraphRunStream
from langgraph.stream.stream_channel import StreamChannel
@@ -50,7 +49,7 @@ class ValuesTransformer(StreamTransformer):
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._log: EventLog[dict[str, Any]] = EventLog()
self._log: StreamChannel[dict[str, Any]] = StreamChannel()
self._latest: dict[str, Any] | None = None
self._interrupted = False
self._interrupts: list[Any] = []
@@ -131,7 +130,7 @@ class MessagesTransformer(StreamTransformer):
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._log: EventLog[ChatModelStream] = EventLog()
self._log: StreamChannel[ChatModelStream] = StreamChannel()
# Correlate protocol events back to a ChatModelStream by run_id
# (attached to the event's metadata by StreamMessagesHandler).
self._by_run: dict[str, ChatModelStream] = {}
@@ -519,7 +518,9 @@ class SubgraphTransformer(_TasksLifecycleBase):
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._log: EventLog[SubgraphRunStream | AsyncSubgraphRunStream] = EventLog()
self._log: StreamChannel[SubgraphRunStream | AsyncSubgraphRunStream] = (
StreamChannel()
)
self._handles: dict[
tuple[str, ...], SubgraphRunStream | AsyncSubgraphRunStream
] = {}
libs/langgraph/tests/test_pregel_stream_v2.py
+54 -55
View File
@@ -15,7 +15,6 @@ from typing_extensions import TypedDict
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from langgraph.stream import (
EventLog,
StreamChannel,
StreamTransformer,
)
@@ -141,13 +140,13 @@ class _CustomPassthroughTransformer(StreamTransformer):
# ---------------------------------------------------------------------------
# EventLog unit tests
# StreamChannel (local, unnamed) unit tests
# ---------------------------------------------------------------------------
class TestEventLog:
class TestStreamChannelLocal:
def test_sync_iteration(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
it = iter(log)
log.push(1)
@@ -157,7 +156,7 @@ class TestEventLog:
assert list(it) == [1, 2, 3]
def test_drain_on_consume(self) -> None:
log: EventLog[str] = EventLog()
log: StreamChannel[str] = StreamChannel()
log._bind(is_async=False)
it = iter(log)
log.push("a")
@@ -167,7 +166,7 @@ class TestEventLog:
assert list(log._items) == []
def test_second_subscribe_raises(self) -> None:
log: EventLog[str] = EventLog()
log: StreamChannel[str] = StreamChannel()
log._bind(is_async=False)
log.close()
_ = iter(log)
@@ -176,7 +175,7 @@ class TestEventLog:
def test_pre_subscription_push_is_noop(self) -> None:
# Lazy-subscribe: pushes before subscription are dropped silently.
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
log.push(1)
log.push(2)
@@ -186,7 +185,7 @@ class TestEventLog:
assert list(it) == [3]
def test_fail_propagation(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
it = iter(log)
log.push(1)
@@ -195,7 +194,7 @@ class TestEventLog:
list(it)
def test_sync_cursor_yields_items_before_error(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
it = iter(log)
log.push(1)
@@ -209,60 +208,60 @@ class TestEventLog:
assert items == [1, 2, 3]
def test_push_after_close_raises(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
it = iter(log)
log.push(1)
log.close()
with pytest.raises(RuntimeError, match="Cannot push to a closed EventLog"):
with pytest.raises(RuntimeError, match="Cannot push to a closed StreamChannel"):
log.push(2)
_ = list(it)
def test_push_after_fail_raises(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
it = iter(log)
log.fail(ValueError("err"))
with pytest.raises(RuntimeError, match="Cannot push to a closed EventLog"):
with pytest.raises(RuntimeError, match="Cannot push to a closed StreamChannel"):
log.push(1)
with pytest.raises(ValueError, match="err"):
list(it)
def test_empty_log_sync(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
log.close()
assert list(log) == []
def test_empty_log_fail_sync(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
log.fail(ValueError("empty fail"))
with pytest.raises(ValueError, match="empty fail"):
list(log)
def test_unbound_iter_raises(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log.close()
with pytest.raises(TypeError, match="has not been bound"):
list(log)
def test_sync_bound_aiter_raises(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
log.close()
with pytest.raises(TypeError, match="bound to sync mode"):
log.__aiter__()
def test_double_bind_raises(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
with pytest.raises(RuntimeError, match="already bound"):
log._bind(is_async=True)
@pytest.mark.anyio
async def test_async_iteration(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=True)
cursor = aiter(log)
for i in range(3):
@@ -272,7 +271,7 @@ class TestEventLog:
@pytest.mark.anyio
async def test_async_second_subscribe_raises(self) -> None:
log: EventLog[str] = EventLog()
log: StreamChannel[str] = StreamChannel()
log._bind(is_async=True)
log.close()
_ = log.__aiter__()
@@ -281,7 +280,7 @@ class TestEventLog:
@pytest.mark.anyio
async def test_async_fail(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=True)
cursor = aiter(log)
log.push(1)
@@ -292,7 +291,7 @@ class TestEventLog:
@pytest.mark.anyio
async def test_async_cursor_yields_items_before_error(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=True)
cursor = aiter(log)
log.push(1)
@@ -307,14 +306,14 @@ class TestEventLog:
@pytest.mark.anyio
async def test_empty_log_async(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=True)
log.close()
assert [item async for item in log] == []
@pytest.mark.anyio
async def test_empty_log_fail_async(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=True)
log.fail(ValueError("empty fail"))
with pytest.raises(ValueError, match="empty fail"):
@@ -323,7 +322,7 @@ class TestEventLog:
@pytest.mark.anyio
async def test_async_bound_iter_raises(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=True)
log.close()
with pytest.raises(TypeError, match="bound to async mode"):
@@ -331,18 +330,18 @@ class TestEventLog:
# ---------------------------------------------------------------------------
# StreamChannel unit tests
# StreamChannel (named, wired) unit tests
# ---------------------------------------------------------------------------
class TestStreamChannel:
class TestStreamChannelNamed:
def test_push_and_iterate(self) -> None:
ch: StreamChannel[str] = StreamChannel("test")
ch._bind(is_async=False)
it = iter(ch)
ch.push("a")
ch.push("b")
ch._close()
ch.close()
assert list(it) == ["a", "b"]
def test_wire_callback(self) -> None:
@@ -353,7 +352,7 @@ class TestStreamChannel:
it = iter(ch)
ch.push("x")
ch.push("y")
ch._close()
ch.close()
assert forwarded == ["x", "y"]
assert list(it) == ["x", "y"]
@@ -362,7 +361,7 @@ class TestStreamChannel:
ch._bind(is_async=False)
it = iter(ch)
ch.push("a")
ch._fail(ValueError("channel error"))
ch.fail(ValueError("channel error"))
items: list[str] = []
with pytest.raises(ValueError, match="channel error"):
for item in it:
@@ -375,7 +374,7 @@ class TestStreamChannel:
assert ch._wire_fn is None
it = iter(ch)
ch.push(42)
ch._close()
ch.close()
assert list(it) == [42]
@pytest.mark.anyio
@@ -383,9 +382,9 @@ class TestStreamChannel:
ch: StreamChannel[str] = StreamChannel("test")
ch._bind(is_async=True)
cursor = ch.__aiter__()
ch._log.push("x")
ch._log.push("y")
ch._close()
ch.push("x")
ch.push("y")
ch.close()
assert [item async for item in cursor] == ["x", "y"]
@@ -922,7 +921,7 @@ class TestStreamMuxResilience:
mux = StreamMux([t])
with pytest.raises(RuntimeError, match="finalize broke"):
mux.close()
assert t._channel._log._closed
assert t._channel._closed
# ---------------------------------------------------------------------------
@@ -963,7 +962,7 @@ class TestCustomTransformer:
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._log: EventLog[str] = EventLog()
self._log: StreamChannel[str] = StreamChannel()
def init(self) -> dict[str, Any]:
return {"foo": self._log}
@@ -1046,7 +1045,7 @@ class TestCustomTransformer:
class ConflictTransformer(StreamTransformer):
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._log: EventLog[str] = EventLog()
self._log: StreamChannel[str] = StreamChannel()
def init(self) -> dict[str, Any]:
return {"values": self._log}
@@ -1061,16 +1060,16 @@ class TestCustomTransformer:
# ---------------------------------------------------------------------------
# EventLog auto-lifecycle via StreamMux
# StreamChannel auto-lifecycle via StreamMux
# ---------------------------------------------------------------------------
class TestEventLogAutoLifecycle:
def test_mux_auto_closes_event_logs(self) -> None:
class TestStreamChannelAutoLifecycle:
def test_mux_auto_closes_channels(self) -> None:
class SimpleTransformer(StreamTransformer):
def __init__(self) -> None:
super().__init__()
self._log: EventLog[str] = EventLog()
self._log: StreamChannel[str] = StreamChannel()
def init(self) -> dict[str, Any]:
return {"items": self._log}
@@ -1085,11 +1084,11 @@ class TestEventLogAutoLifecycle:
mux.close()
assert len(list(it)) == 1
def test_mux_auto_fails_event_logs(self) -> None:
def test_mux_auto_fails_channels(self) -> None:
class SimpleTransformer(StreamTransformer):
def __init__(self) -> None:
super().__init__()
self._log: EventLog[str] = EventLog()
self._log: StreamChannel[str] = StreamChannel()
def init(self) -> dict[str, Any]:
return {"items": self._log}
@@ -1110,7 +1109,7 @@ class TestEventLogAutoLifecycle:
class ManualCloseTransformer(StreamTransformer):
def __init__(self) -> None:
super().__init__()
self._log: EventLog[str] = EventLog()
self._log: StreamChannel[str] = StreamChannel()
def init(self) -> dict[str, Any]:
return {"items": self._log}
@@ -1128,7 +1127,7 @@ class TestEventLogAutoLifecycle:
class MinimalTransformer(StreamTransformer):
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._log: EventLog[str] = EventLog()
self._log: StreamChannel[str] = StreamChannel()
def init(self) -> dict[str, Any]:
return {"minimal": self._log}
@@ -1216,7 +1215,7 @@ class TestAsyncTransformerLane:
requires_async = True
def __init__(self) -> None:
self._log: EventLog[str] = EventLog()
self._log: StreamChannel[str] = StreamChannel()
def init(self) -> dict[str, Any]:
return {"out": self._log}
@@ -1273,7 +1272,7 @@ class TestAsyncTransformerLane:
requires_async = True
def __init__(self) -> None:
self._log: EventLog[str] = EventLog()
self._log: StreamChannel[str] = StreamChannel()
def init(self) -> dict[str, Any]:
return {"out": self._log}
@@ -1353,7 +1352,7 @@ class TestAsyncTransformerLane:
requires_async = True
def __init__(self) -> None:
self._log: EventLog[str] = EventLog()
self._log: StreamChannel[str] = StreamChannel()
def init(self) -> dict[str, Any]:
return {"seen": self._log}
@@ -1381,7 +1380,7 @@ class TestAsyncTransformerLane:
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._log: EventLog[int] = EventLog()
self._log: StreamChannel[int] = StreamChannel()
def init(self) -> dict[str, Any]:
return {"scores": self._log}
@@ -1469,20 +1468,20 @@ class TestMemoryBounds:
# ---------------------------------------------------------------------------
# DrainOnConsume: EventLog capacity semantics
# DrainOnConsume: StreamChannel capacity semantics
# ---------------------------------------------------------------------------
class TestDrainOnConsume:
def test_invalid_maxlen_raises(self) -> None:
with pytest.raises(ValueError, match="positive int or None"):
EventLog(maxlen=0)
StreamChannel(maxlen=0)
with pytest.raises(ValueError, match="positive int or None"):
EventLog(maxlen=-3)
StreamChannel(maxlen=-3)
def test_push_unbounded_by_design(self) -> None:
"""Push is non-blocking; the caller-driven pump bounds memory via iteration pace."""
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
it = iter(log)
for i in range(100):
@@ -1491,7 +1490,7 @@ class TestDrainOnConsume:
assert list(it) == list(range(100))
def test_tee_fans_out_sync(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=False)
a, b = log.tee(2)
for i in range(3):
@@ -1502,7 +1501,7 @@ class TestDrainOnConsume:
@pytest.mark.anyio
async def test_atee_fans_out(self) -> None:
log: EventLog[int] = EventLog()
log: StreamChannel[int] = StreamChannel()
log._bind(is_async=True)
a, b = log.atee(2)
for i in range(3):
libs/langgraph/tests/test_stream_lifecycle_transformer.py
+7 -7
View File
@@ -80,23 +80,23 @@ def _tasks_result(
def _arm(mux: StreamMux) -> None:
"""Force projection logs to accept pushes (skip lazy-subscribe gate).
"""Force projection channels to accept pushes (skip lazy-subscribe gate).
`EventLog.push` is a no-op until a subscriber attaches. Tests that
inspect `_items` directly need the gate flipped before any event
is dispatched.
`StreamChannel.push` only appends to the local buffer when a
subscriber is attached. Tests that inspect `_items` directly need
the gate flipped before any event is dispatched.
"""
mux._events._subscribed = True
for transformer in mux._transformers:
if isinstance(transformer, LifecycleTransformer):
transformer._channel._log._subscribed = True
transformer._channel._subscribed = True
def _drain_lifecycle(mux: StreamMux) -> list[LifecyclePayload]:
"""Snapshot the lifecycle channel's underlying log."""
"""Snapshot the lifecycle channel's buffer."""
transformer = mux.transformer_by_key("lifecycle")
assert isinstance(transformer, LifecycleTransformer)
return list(transformer._channel._log._items)
return list(transformer._channel._items)
def _build_lifecycle_mux(*, scope: tuple[str, ...] = ()) -> StreamMux:
libs/langgraph/tests/test_stream_messages_transformer.py
+13 -9
View File
@@ -18,9 +18,9 @@ from typing_extensions import TypedDict
from langgraph.constants import END, START
from langgraph.graph import MessagesState, StateGraph
from langgraph.stream._event_log import EventLog
from langgraph.stream._mux import StreamMux
from langgraph.stream.run_stream import GraphRunStream
from langgraph.stream.stream_channel import StreamChannel
from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer
TS = int(time.time() * 1000)
@@ -90,9 +90,11 @@ def _whole_msg(
}
def _make_sync_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStream]]:
def _make_sync_transformer() -> tuple[
MessagesTransformer, StreamChannel[ChatModelStream]
]:
t = MessagesTransformer()
log: EventLog[ChatModelStream] = t.init()["messages"]
log: StreamChannel[ChatModelStream] = t.init()["messages"]
log._bind(is_async=False)
# Subscribe up front so pushes during process() are retained.
log._subscribed = True
@@ -100,9 +102,11 @@ def _make_sync_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStr
return t, log
def _make_async_transformer() -> tuple[MessagesTransformer, EventLog[ChatModelStream]]:
def _make_async_transformer() -> tuple[
MessagesTransformer, StreamChannel[ChatModelStream]
]:
t = MessagesTransformer()
log: EventLog[ChatModelStream] = t.init()["messages"]
log: StreamChannel[ChatModelStream] = t.init()["messages"]
log._bind(is_async=True)
log._subscribed = True
return t, log
@@ -410,7 +414,7 @@ class TestWireRequestMore:
mux = StreamMux([values_t, messages_t], is_async=False)
GraphRunStream(iter([]), mux, values_t)
log: EventLog[ChatModelStream] = mux.extensions["messages"]
log: StreamChannel[ChatModelStream] = mux.extensions["messages"]
log._subscribed = True
for evt in _lifecycle():
messages_t.process(_proto_event(evt))
@@ -427,12 +431,12 @@ class TestWireRequestMore:
class TestViaMux:
def _make_mux(
self,
) -> tuple[MessagesTransformer, StreamMux, EventLog[ChatModelStream]]:
) -> tuple[MessagesTransformer, StreamMux, StreamChannel[ChatModelStream]]:
t = MessagesTransformer()
v = ValuesTransformer()
mux = StreamMux([v, t], is_async=False)
t._bind_pump(lambda: False)
log: EventLog[ChatModelStream] = mux.extensions["messages"]
log: StreamChannel[ChatModelStream] = mux.extensions["messages"]
log._subscribed = True
return t, mux, log
@@ -456,7 +460,7 @@ class TestViaMux:
t = MessagesTransformer()
v = ValuesTransformer()
mux = StreamMux([v, t], is_async=True)
log: EventLog[ChatModelStream] = mux.extensions["messages"]
log: StreamChannel[ChatModelStream] = mux.extensions["messages"]
log._subscribed = True
for evt in _lifecycle(text="async mux"):
libs/langgraph/tests/test_stream_subgraph_transformer.py
+2 -4
View File
@@ -124,10 +124,8 @@ def _arm(mux: StreamMux) -> None:
"""
mux._events._subscribed = True
for value in mux.extensions.values():
if hasattr(value, "_subscribed"): # EventLog
if hasattr(value, "_subscribed"):
value._subscribed = True
elif hasattr(value, "_log"): # StreamChannel
value._log._subscribed = True
def _arm_recursive(mux: StreamMux) -> None:
@@ -177,7 +175,7 @@ def _event_items(mux: StreamMux) -> list[ProtocolEvent]:
def _lifecycle_payloads(mux: StreamMux) -> list[dict[str, Any]]:
lifecycle_t = mux.transformer_by_key("lifecycle")
assert isinstance(lifecycle_t, LifecycleTransformer)
return list(lifecycle_t._channel._log._items)
return list(lifecycle_t._channel._items)
# ---------------------------------------------------------------------------
libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py
+5 -5
View File
@@ -11,7 +11,7 @@ from __future__ import annotations
from collections.abc import AsyncIterator, Iterator
from typing import Any
from langgraph.stream._event_log import EventLog
from langgraph.stream.stream_channel import StreamChannel
class ToolCallStream:
@@ -21,7 +21,7 @@ class ToolCallStream:
are populated as events arrive:
- `tool_call_id`, `tool_name`, `input`: stable from the start event.
- `output_deltas`: an `EventLog` of delta chunks. Iterate (sync or
- `output_deltas`: a `StreamChannel` of delta chunks. Iterate (sync or
async) to consume partial output in arrival order.
- `output`: terminal payload from `tool-finished`, or `None` if the
call failed or is still in flight.
@@ -51,14 +51,14 @@ class ToolCallStream:
self.tool_call_id = tool_call_id
self.tool_name = tool_name
self.input = input
self._output_deltas: EventLog[Any] = EventLog()
self._output_deltas: StreamChannel[Any] = StreamChannel()
self.output: Any = None
self.error: str | None = None
self.completed = False
@property
def output_deltas(self) -> EventLog[Any]:
"""The EventLog of streamed `tool-output-delta` payloads.
def output_deltas(self) -> StreamChannel[Any]:
"""The channel of streamed `tool-output-delta` payloads.
Iterate (sync or async depending on how the run was started)
to consume partial output in arrival order. The log closes when
libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py
+8 -7
View File
@@ -5,8 +5,8 @@ from __future__ import annotations
from collections.abc import Awaitable, Callable
from typing import Any
from langgraph.stream._event_log import EventLog
from langgraph.stream._types import ProtocolEvent, StreamTransformer
from langgraph.stream.stream_channel import StreamChannel
from langgraph.prebuilt._tool_call_stream import ToolCallStream
@@ -21,11 +21,12 @@ class ToolCallTransformer(StreamTransformer):
Native transformer — the `tool_calls` projection is exposed as a
direct attribute on the run stream.
`EventLog[ToolCallStream]` is used (not `StreamChannel`) because the
live handles are not serializable and should not be auto-forwarded
onto the main event log. Wire consumers subscribe to the `tools`
channel instead, where the raw protocol events flow through
untouched by this transformer (`process` returns `True`).
A nameless `StreamChannel[ToolCallStream]` is used (no protocol
auto-forwarding) because the live handles are not serializable and
should not be injected into the main event log. Wire consumers
subscribe to the `tools` channel instead, where the raw protocol
events flow through untouched by this transformer (`process`
returns `True`).
Registered explicitly by users at compile time via
`builder.compile(transformers=[ToolCallTransformer])` — not a
@@ -37,7 +38,7 @@ class ToolCallTransformer(StreamTransformer):
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._log: EventLog[ToolCallStream] = EventLog()
self._log: StreamChannel[ToolCallStream] = StreamChannel()
self._active: dict[str, ToolCallStream] = {}
self._is_async = False
self._pump_fn: Callable[[], bool] | None = None
libs/prebuilt/tests/test_tool_call_transformer.py
+2 -2
View File
@@ -11,9 +11,9 @@ from langchain_core.tools import tool
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
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 (
MessagesTransformer,
ValuesTransformer,
@@ -63,7 +63,7 @@ def _tool_event(
}
def _subscribe(log: EventLog) -> None:
def _subscribe(log: StreamChannel) -> None:
log._subscribed = True