Convert stream module docstrings to Google style

Per repo convention (CLAUDE.md) and general project style: use single
backticks for inline code, Google-style Args/Returns/Raises sections,
and triple-backtick fenced code blocks instead of Sphinx double
backticks, :param: markers, or Usage:: blocks.

No behavior changes — docs only.
This commit is contained in:
Nick Hollon
2026-04-16 14:45:02 -04:00
parent 6fcca359df
commit adda5f0341
9 changed files with 380 additions and 200 deletions
+1 -1
View File
@@ -1,6 +1,6 @@
"""Streaming infrastructure for LangGraph.
Provides a ``StreamingHandler`` that wraps a compiled graph and exposes
Provides a `StreamingHandler` that wraps a compiled graph and exposes
ergonomic streaming projections through a transformer pipeline.
"""
+6 -2
View File
@@ -9,8 +9,12 @@ from langgraph.stream._types import ProtocolEvent, _ProtocolEventParams
def convert_to_protocol_event(part: dict[str, Any]) -> ProtocolEvent:
"""Convert a v2 StreamPart dict to a ProtocolEvent.
Expects a dict with keys ``type``, ``ns``, ``data``, and optionally
``interrupts`` (present on values events).
Args:
part: A stream part dict with keys `type`, `ns`, `data`, and
optionally `interrupts` (present on values events).
Returns:
The equivalent ProtocolEvent.
"""
params: _ProtocolEventParams = {
"namespace": list(part["ns"]),
+61 -26
View File
@@ -11,7 +11,7 @@ T = TypeVar("T")
class BufferOverflowError(RuntimeError):
"""Raised when an EventLog cursor falls off the back of a bounded buffer.
Mirrors the ``restored: false`` signal from the protocol's reconnection
Mirrors the `restored: false` signal from the protocol's reconnection
story (§ 06): consumers that fall behind the retention window get an
explicit error and can decide to rebuild from a snapshot rather than
silently losing events.
@@ -21,31 +21,27 @@ class BufferOverflowError(RuntimeError):
class EventLog(Generic[T]):
"""Append-only buffer that supports multiple independent consumers.
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``.
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`.
All access is single-threaded: sync mode is caller-driven (no
background thread), async mode runs entirely on the event loop.
Producer API:
push(item) append an item, notify all waiting cursors
close() mark the log as done
fail(err) mark the log as errored
- `push(item)`: append an item, notify all waiting cursors.
- `close()`: mark the log as done.
- `fail(err)`: mark the log as errored.
Sync iteration is pull-based: when a cursor catches up it calls
``_request_more`` to drive the graph forward.
`_request_more` to drive the graph forward. Async iteration uses a
shared `asyncio.Event` — cursors await the event when they catch
up, and the producer sets it on each push.
Async iteration uses a shared ``asyncio.Event`` — cursors await
the event when they catch up, and the producer sets it on each push.
Bounded mode
------------
Pass ``maxlen=N`` to cap memory. When the buffer is full, ``push``
Pass `maxlen=N` to cap memory. When the buffer is full, `push`
drops the oldest item to make room. Cursors track an absolute
sequence number; a cursor that falls off the back of the retention
window raises ``BufferOverflowError`` on its next read.
window raises `BufferOverflowError` on its next read.
New cursors start at the current head of the buffer, not at seq 0
— they see whatever is still retained. This matches the protocol's
@@ -54,6 +50,16 @@ class EventLog(Generic[T]):
"""
def __init__(self, maxlen: int | None = None) -> None:
"""Initialize an empty, unbound log.
Args:
maxlen: Optional cap on retained items. When reached, the
oldest item is dropped on each new push. `None` (the
default) leaves the log unbounded.
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()
@@ -78,8 +84,14 @@ class EventLog(Generic[T]):
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.
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")
@@ -92,11 +104,17 @@ class EventLog(Generic[T]):
# ------------------------------------------------------------------
def push(self, item: T) -> None:
"""Append *item* and wake all waiting cursors.
"""Append an item and wake all waiting cursors.
In bounded mode, evicts the oldest item first if the buffer
is full, advancing ``_first_seq`` so cursors can detect that
they've fallen off the back of the retention window.
In bounded mode, evicts the oldest item first if the buffer is
full, advancing `_first_seq` so cursors can detect that they've
fallen off the back of the retention window.
Args:
item: The item to append.
Raises:
RuntimeError: If the log is closed.
"""
if self._closed:
raise RuntimeError("Cannot push to a closed EventLog")
@@ -107,12 +125,21 @@ class EventLog(Generic[T]):
self._notify()
def close(self) -> None:
"""Mark the log as complete — open cursors will finish cleanly."""
"""Mark the log as complete.
Open cursors will finish cleanly once they drain the buffer.
"""
self._closed = True
self._notify()
def fail(self, err: BaseException) -> None:
"""Mark the log as errored — open cursors will raise *err*."""
"""Mark the log as errored.
Open cursors will raise `err` once they drain the buffer.
Args:
err: The exception to surface to consumers.
"""
self._error = err
self._closed = True
self._notify()
@@ -134,7 +161,11 @@ class EventLog(Generic[T]):
# ------------------------------------------------------------------
def __iter__(self) -> Iterator[T]:
"""Return a new independent sync cursor over the log."""
"""Return a new independent sync cursor over the log.
Raises:
TypeError: If the log is unbound or bound to async mode.
"""
if self._is_async is None:
raise TypeError(
"EventLog has not been bound yet. "
@@ -180,7 +211,11 @@ class EventLog(Generic[T]):
# ------------------------------------------------------------------
def __aiter__(self) -> AsyncIterator[T]:
"""Return a new independent async cursor over the log."""
"""Return a new independent async cursor over the log.
Raises:
TypeError: If the log is unbound or bound to sync mode.
"""
if self._is_async is None:
raise TypeError(
"EventLog has not been bound yet. "
+96 -63
View File
@@ -19,13 +19,20 @@ class StreamMux:
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``
auto-wired so that every `push()` also injects a `ProtocolEvent`
into the main log.
Pass ``is_async=True`` when the mux will be consumed via async
iteration (``handler.astream()``). All ``EventLog`` and
``StreamChannel`` instances discovered during ``register()`` are
automatically bound to the matching mode.
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.
Attributes:
extensions: Merged projection dict across all registered
transformers. Treat as read-only — mutations won't be
reflected back in individual transformers' state.
native_keys: Projection keys contributed by transformers with
`_native = True`.
"""
def __init__(
@@ -35,22 +42,30 @@ class StreamMux:
is_async: bool = False,
max_events: int | None = None,
) -> None:
"""Initialize the mux and register *transformers* in order.
"""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,
projections are merged into ``self.extensions``, ``_native``
keys are recorded in ``self.native_keys``, and any ``EventLog``
/ ``StreamChannel`` instances are bound/wired.
post-init `register()`. 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.
*max_events* sets a default capacity for every ``EventLog`` /
``StreamChannel`` the mux binds, including the main event log.
Logs that were constructed with an explicit ``maxlen`` keep
their own setting — the mux only fills in ``None`` defaults.
Unbounded when ``max_events`` is ``None``.
Args:
transformers: Transformers to register, in dispatch order.
`None` or empty gives a mux with no projections.
is_async: True for async dispatch (`apush` / `aclose` /
`afail`), False for the sync path.
max_events: Default capacity for every EventLog and
StreamChannel the mux binds, including the main event
log. Logs constructed with an explicit `maxlen` keep
their own setting — the mux only fills in unset
defaults. `None` leaves the logs unbounded.
Raises ``RuntimeError`` if any transformer requires an async run
under sync mode, and ``ValueError`` on projection-key conflicts.
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.
"""
self._is_async = is_async
self._default_maxlen = max_events
@@ -61,23 +76,18 @@ class StreamMux:
self._logs: list[EventLog[Any]] = []
self._seq = 0
#: Merged projection dict across all registered transformers.
#: Treat as read-only — mutations won't be reflected back in
#: individual transformers' state.
self.extensions: dict[str, Any] = {}
#: Projection keys from transformers with ``_native = True``.
self.native_keys: set[str] = set()
for transformer in transformers or ():
self._register(transformer)
def _register(self, transformer: StreamTransformer) -> None:
"""Register a transformer.
"""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
``self.extensions``.
Calls `transformer.init()`, stores the transformer for event
processing, binds any EventLog or StreamChannel instances in
the projection, and merges the projection into `extensions`.
"""
if transformer_requires_async(transformer) and not self._is_async:
raise RuntimeError(
@@ -105,18 +115,21 @@ class StreamMux:
self.native_keys.update(projection.keys())
def push(self, event: ProtocolEvent) -> None:
"""Route *event* through all transformers, then append to the main log.
"""Route an event through all transformers, then append to the main log.
Each transformer's ``process()`` is called in registration order.
If any transformer returns ``False``, the event is suppressed
from the main log (but transformers that already saw it keep
their side-effects).
Each transformer's `process()` is called in registration order.
If any transformer returns False, the event is suppressed from
the main log, but transformers that already saw it keep their
side effects.
Seq is assigned right before an event enters the main log, not
before the transformer pipeline runs. This ensures that events
auto-forwarded from StreamChannels during ``process()`` get
auto-forwarded from StreamChannels during `process()` get
earlier seq numbers than the original event, preserving
monotonic ordering in the log.
Args:
event: The protocol event to dispatch.
"""
keep = True
for transformer in self._transformers:
@@ -130,11 +143,16 @@ 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.
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.
Raises:
BaseException: The first error raised by a transformer's
`finalize()`, re-raised after cleanup finishes.
"""
first_error: BaseException | None = None
for transformer in self._transformers:
@@ -156,10 +174,14 @@ 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.
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.
Args:
err: The exception that ended the run.
"""
for transformer in self._transformers:
try:
@@ -179,13 +201,17 @@ class StreamMux:
# ------------------------------------------------------------------
async def apush(self, event: ProtocolEvent) -> None:
"""Async counterpart to ``push``. Awaits each transformer's
``aprocess`` in registration order before appending to the main log.
"""Dispatch an event on the async lane.
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.
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.
Args:
event: The protocol event to dispatch.
"""
keep = True
for transformer in self._transformers:
@@ -197,15 +223,19 @@ class StreamMux:
self._events.push(event)
async def aclose(self) -> None:
"""Async counterpart to ``close``.
"""Finalize on the async lane.
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.
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.
If any scheduled task raised under ``on_error="raise"``, or any
transformer's ``afinalize`` raises, the exception propagates.
The caller (the pump) handles it by routing into ``afail``.
If any scheduled task raised under `on_error="raise"`, or any
transformer's `afinalize` raises, the exception propagates.
The caller (the pump) handles it by routing into `afail`.
Raises:
BaseException: The first scheduled-task or `afinalize`
error, re-raised after cleanup.
"""
pending = self._collect_scheduled_tasks()
if pending:
@@ -240,11 +270,14 @@ class StreamMux:
raise first_error
async def afail(self, err: BaseException) -> None:
"""Async counterpart to ``fail``.
"""Fail on the async lane.
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.
them to completion, then runs each transformer's `afail` hook
and auto-fails logs, channels, and the main event log.
Args:
err: The exception that ended the run.
"""
pending = self._collect_scheduled_tasks()
for task in pending:
@@ -267,7 +300,7 @@ class StreamMux:
self._events.fail(err)
def _collect_scheduled_tasks(self) -> list[asyncio.Task[Any]]:
"""Snapshot of all in-flight tasks scheduled via transformers."""
"""Return a snapshot of in-flight tasks scheduled via transformers."""
return [
task
for transformer in self._transformers
@@ -280,7 +313,7 @@ class StreamMux:
# ------------------------------------------------------------------
def _bind_and_wire(self, projection: dict[str, Any]) -> None:
"""Bind and wire EventLog / StreamChannel instances in *projection*."""
"""Bind and wire EventLog / StreamChannel instances in a projection."""
for value in projection.values():
if isinstance(value, StreamChannel):
self._apply_default_maxlen(value._log)
@@ -301,7 +334,7 @@ class StreamMux:
self._logs.append(value)
def _apply_default_maxlen(self, log: EventLog[Any]) -> None:
"""Fill in the mux's default maxlen if the log hasn't set its own."""
"""Fill in the mux's default maxlen when the log hasn't set its own."""
if log._maxlen is None and self._default_maxlen is not None:
log._maxlen = self._default_maxlen
@@ -310,9 +343,9 @@ class StreamMux:
Forwarded events bypass the transformer pipeline to avoid
infinite recursion (a transformer that pushes to a channel
during ``process()`` would re-trigger itself). These events
are visible in the main event log but are not passed through
transformers' ``process()`` methods.
during `process()` would re-trigger itself). These events are
visible in the main event log but are not passed through
transformers' `process()` methods.
"""
self._seq += 1
event: ProtocolEvent = {
+88 -49
View File
@@ -38,41 +38,43 @@ 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 (EventLogs, 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
appearing in `run.extensions`).
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.
`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/failed by the mux, so most transformers don't need
``finalize`` or ``fail`` at all.
auto-closed / auto-failed by the mux, so most transformers don't
need `finalize` or `fail` at all.
Async lane
----------
Transformers that need async work pick the async lane by:
1. Overriding ``aprocess`` (and optionally ``afinalize``/``afail``), or
2. Calling ``self.schedule(coro)`` from inside a sync ``process``, or
3. Setting ``requires_async = True`` explicitly.
1. Overriding `aprocess` (and optionally `afinalize` / `afail`), or
2. Calling `self.schedule(coro)` from inside a sync `process`, or
3. Setting `requires_async = True` explicitly.
The mux detects these cases at registration and raises if they're
used under sync ``stream()`` — they only work under ``astream()``.
used under sync `stream()` — they only work under `astream()`.
Use ``aprocess`` when the pump must *wait* for async work before
the next transformer sees the event (e.g. PII redaction that
mutates ``event`` in place). Use ``schedule()`` for decoupled async
work whose result lands on an independent projection (e.g. async
moderation scoring, cost lookup, external tracing).
Use `aprocess` when the pump must wait for async work before the
next transformer sees the event (e.g. PII redaction that mutates
`event` in place). Use `schedule()` for decoupled async work whose
result lands on an independent projection (e.g. async moderation
scoring, cost lookup, external tracing).
Attributes:
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
`process`). The mux also auto-detects the async lane when
`aprocess`, `afinalize`, or `afail` is overridden.
"""
#: 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 ``process``). The mux also auto-detects the
#: async lane when ``aprocess``/``afinalize``/``afail`` is overridden.
requires_async: ClassVar[bool] = False
@abstractmethod
@@ -89,30 +91,40 @@ class StreamTransformer(ABC):
...
def process(self, event: ProtocolEvent) -> bool:
"""Sync event handler. Override for the sync lane.
"""Handle an event on the sync lane.
Called for every event before it is appended to the main event log.
Return False to suppress the event from the main log.
Called for every event before it is appended to the main event
log. Subclasses must override either `process` or `aprocess`.
The default raises so a missing override fails loudly rather
than silently passing every event through.
Subclasses must override either ``process`` or ``aprocess``. The
default raises so a missing override fails loudly rather than
silently passing every event through.
Args:
event: The protocol event to observe.
Returns:
True to keep the event in the main log, False to suppress it.
"""
raise NotImplementedError(
f"{type(self).__name__} must override process() or aprocess()"
)
async def aprocess(self, event: ProtocolEvent) -> bool:
"""Async event handler. Override for the async lane.
"""Handle an event on the async lane.
The mux awaits this before dispatching to the next transformer,
so a slow ``aprocess`` serializes the pipeline. Use this only
when a later transformer — or a consumer reading the event
so a slow `aprocess` serializes the pipeline. Use it only when
a later transformer — or a consumer reading the event
synchronously — must see the result of the async work (e.g.
PII redaction that mutates ``event`` in place).
PII redaction that mutates `event` in place).
The default delegates to ``process``, so purely-sync transformers
run unchanged under ``astream()``.
The default delegates to `process`, so purely-sync transformers
run unchanged under `astream()`.
Args:
event: The protocol event to observe.
Returns:
True to keep the event in the main log, False to suppress it.
"""
return self.process(event)
@@ -127,10 +139,10 @@ class StreamTransformer(ABC):
"""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 EventLogs can be closed here
without a last-task-wins race.
The default delegates to ``finalize``.
The default delegates to `finalize`.
"""
self.finalize()
@@ -139,15 +151,21 @@ class StreamTransformer(ABC):
Override to fail EventLogs, reject promises, or perform other
teardown. StreamChannel instances are auto-failed by the mux.
Args:
err: The exception that ended the run.
"""
async def afail(self, err: BaseException) -> None:
"""Called when the run ends with an error (async lane).
The mux cancels and awaits every task started via ``schedule()``
The mux cancels and awaits every task started via `schedule()`
before calling this, so cleanup doesn't race with in-flight work.
The default delegates to ``fail``.
The default delegates to `fail`.
Args:
err: The exception that ended the run.
"""
self.fail(err)
@@ -164,19 +182,31 @@ class StreamTransformer(ABC):
"""Schedule a coroutine tied to this transformer's lifecycle.
The mux holds the task reference, awaits all scheduled tasks
during ``aclose()`` before calling ``afinalize()``, and cancels
them on ``afail()``. Authors don't need to track tasks or
during `aclose()` before calling `afinalize()`, and cancels
them on `afail()`. Authors don't need to track tasks or
implement the last-task-closes-the-log dance.
``on_error="log"`` (default): exceptions are caught and logged;
a single failed call doesn't tear down the run.
Requires a running event loop — call only under `astream()`.
Set `requires_async = True` on the class so registration under
sync `stream()` fails fast with a clear message.
``on_error="raise"``: exceptions propagate when the mux joins
pendings, converting the close path into the fail path.
Args:
coro: The coroutine to run. Its lifecycle is owned by the
mux from this point on.
on_error: `"log"` (default) catches and logs any exception
the coroutine raises, so a single failure doesn't tear
down the run. `"raise"` lets the exception propagate
when the mux joins pendings, converting the close path
into the fail path.
Requires a running event loop — call only under ``astream()``.
Set ``requires_async = True`` on the class so registration under
sync ``stream()`` fails fast with a clear message.
Returns:
The asyncio Task. Authors rarely need to await it directly
— consumers read results from whatever projection the
coroutine pushes into.
Raises:
RuntimeError: If called without a running event loop (i.e.
under sync `stream()` rather than `astream()`).
"""
try:
asyncio.get_running_loop()
@@ -205,7 +235,10 @@ class StreamTransformer(ABC):
_logger.exception("Scheduled StreamTransformer task failed")
def _scheduled_task_set(self) -> set[asyncio.Task[Any]]:
"""Lazily-allocated task set. Avoids requiring super().__init__()."""
"""Return the lazily-allocated task set.
Avoids requiring subclasses to call `super().__init__()`.
"""
tasks: set[asyncio.Task[Any]] | None = getattr(
self, "_stream_scheduled_tasks", None
)
@@ -216,11 +249,17 @@ class StreamTransformer(ABC):
def transformer_requires_async(transformer: StreamTransformer) -> bool:
"""True if the transformer needs a running event loop.
"""Return True if the transformer needs a running event loop.
A transformer requires async if it explicitly opts in
(``requires_async = True``) or overrides any of the async-lane methods
(``aprocess``, ``afinalize``, ``afail``).
(`requires_async = True`) or overrides any of the async-lane methods
(`aprocess`, `afinalize`, `afail`).
Args:
transformer: The transformer to inspect.
Returns:
True if the transformer cannot run under sync `stream()`.
"""
if transformer.requires_async:
return True
+47 -22
View File
@@ -15,16 +15,16 @@ from langgraph.stream.transformers import ValuesTransformer
class GraphRunStream:
"""Sync run stream with caller-driven pumping.
The caller's iteration on any projection (``values``, ``messages``,
raw events, or ``output``) drives the graph forward. No background
thread is used — this matches v1's model where the caller's ``for``
The caller's iteration on any projection (`values`, `messages`,
raw events, or `output`) drives the graph forward. No background
thread is used — this matches v1's model where the caller's `for`
loop is the pump.
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``).
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`).
Iterating the run stream directly yields raw ``ProtocolEvent`` objects
Iterating the run stream directly yields raw `ProtocolEvent` objects
from the mux's main event log.
"""
@@ -34,6 +34,14 @@ class GraphRunStream:
mux: StreamMux,
values_transformer: ValuesTransformer,
) -> None:
"""Initialize the run stream.
Args:
graph_iter: Pull-based iterator over the graph's stream.
mux: The StreamMux owning projections and the main log.
values_transformer: The built-in values transformer
providing `output` / `interrupted` / `interrupts`.
"""
self._graph_iter = graph_iter
self._mux = mux
self.extensions = mux.extensions
@@ -47,7 +55,11 @@ class GraphRunStream:
self._wire_request_more(mux)
def _wire_request_more(self, mux: StreamMux) -> None:
"""Set _request_more on all sync EventLogs so iteration drives the graph."""
"""Install `_request_more` on every sync EventLog.
Sync iteration is caller-driven, so a cursor that catches up to
the buffer's tail needs a way to ask the graph for more events.
"""
mux._events._request_more = self._pump_next
for value in mux.extensions.values():
if isinstance(value, EventLog):
@@ -56,9 +68,11 @@ class GraphRunStream:
value._log._request_more = self._pump_next
def _pump_next(self) -> bool:
"""Pull one event from the graph and push through the mux.
"""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.
Returns:
True if an event was pulled, False if the graph is exhausted
or has raised.
"""
if self._exhausted:
return False
@@ -110,14 +124,14 @@ class AsyncGraphRunStream:
A background asyncio task pumps events from the graph into the
transformer pipeline. This is the standard async pattern — the task
runs on the same event loop and async consumers can iterate multiple
projections concurrently.
runs on the same event loop and async consumers can iterate
multiple projections concurrently.
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``).
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`).
Async-iterating the run stream yields raw ``ProtocolEvent`` objects
Async-iterating the run stream yields raw `ProtocolEvent` objects
from the mux's main event log.
"""
@@ -127,6 +141,14 @@ class AsyncGraphRunStream:
values_transformer: ValuesTransformer,
pump_task: asyncio.Task[None],
) -> None:
"""Initialize the async run stream.
Args:
mux: The StreamMux owning projections and the main log.
values_transformer: The built-in values transformer
providing `output` / `interrupted` / `interrupts`.
pump_task: Background task pumping graph events into the mux.
"""
self._mux = mux
self.extensions = mux.extensions
self._values_transformer = values_transformer
@@ -139,9 +161,10 @@ class AsyncGraphRunStream:
def output(self) -> Any:
"""Return an awaitable that resolves to the final state.
Usage::
Example:
```python
output = await run.output
```
"""
return self._get_output()
@@ -158,9 +181,10 @@ class AsyncGraphRunStream:
def interrupted(self) -> Any:
"""Return an awaitable that resolves to whether the run was interrupted.
Usage::
Example:
```python
interrupted = await run.interrupted
```
"""
return self._get_interrupted()
@@ -175,9 +199,10 @@ class AsyncGraphRunStream:
def interrupts(self) -> Any:
"""Return an awaitable that resolves to interrupt payloads.
Usage::
Example:
```python
interrupts = await run.interrupts
```
"""
return self._get_interrupts()
@@ -12,34 +12,50 @@ class StreamChannel(Generic[T]):
"""A named projection channel 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``.
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.
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>")``.
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>")`.
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.
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.
Lifecycle (``_close`` / ``_fail``) is managed by the mux — transformers
using only StreamChannels don't need ``finalize`` / ``fail`` hooks.
Lifecycle (`_close` / `_fail`) is managed by the mux — transformers
using only StreamChannels don't need `finalize` or `fail` hooks.
"""
def __init__(self, name: str, *, maxlen: int | None = None) -> None:
"""Initialize the channel with an empty inner log.
Args:
name: The protocol channel name used for auto-forwarded
events (`custom:<name>` on the wire).
maxlen: Optional retention cap on the inner EventLog. See
`EventLog.__init__` for semantics.
"""
self.name = name
self._log: EventLog[T] = EventLog(maxlen=maxlen)
self._wire_fn: Callable[[T], None] | None = None
def _bind(self, *, is_async: bool) -> None:
"""Bind the underlying event log to sync or async mode."""
"""Bind the underlying event log to sync or async mode.
Args:
is_async: True for async iteration, False for sync.
"""
self._log._bind(is_async=is_async)
def push(self, item: T) -> None:
"""Append *item* to the log and auto-forward if wired."""
"""Append an item to the log and auto-forward if wired.
Args:
item: The item to push.
"""
self._log.push(item)
if self._wire_fn is not None:
self._wire_fn(item)
@@ -26,10 +26,10 @@ STREAM_V2_MODES: list[StreamMode] = [
class StreamingHandler:
"""Wraps a compiled graph and provides ergonomic streaming projections.
Usage::
"""Wrap a compiled graph with ergonomic streaming projections.
Example:
```python
handler = StreamingHandler(graph)
# Sync
@@ -43,9 +43,15 @@ class StreamingHandler:
async for state in run.values:
print(state)
output = await run.output
```
"""
def __init__(self, graph: Any) -> None:
"""Initialize the handler.
Args:
graph: A compiled LangGraph graph to stream from.
"""
self._graph = graph
def stream(
@@ -60,17 +66,27 @@ class StreamingHandler:
) -> GraphRunStream:
"""Start a sync streaming run.
Returns a `GraphRunStream` immediately. The caller's iteration on
any projection drives the graph forward — no background thread is
used. This matches v1's model where the caller's ``for`` loop is
the pump.
Returns a GraphRunStream immediately. The caller's iteration on
any projection drives the graph forward — no background thread
is used. This matches v1's model where the caller's `for` loop
is the pump.
*max_events* caps the retention of every ``EventLog`` /
``StreamChannel`` the mux binds (main event log plus each
transformer's projection logs) to the given number of items,
dropping the oldest when full. Transformers that constructed
their own logs with an explicit ``maxlen`` keep their setting.
Unbounded when ``None``.
Args:
input: Graph input.
config: Optional runnable config forwarded to the graph.
interrupt_before: Nodes to interrupt before, if any.
interrupt_after: Nodes to interrupt after, if any.
transformers: User transformers appended after the built-in
`ValuesTransformer` and `MessagesTransformer`.
max_events: Caps the retention of every EventLog and
StreamChannel the mux binds (main event log plus each
transformer's projection logs), dropping the oldest
when full. Transformers that constructed their own
logs with an explicit `maxlen` keep their setting.
Unbounded when `None`.
Returns:
A GraphRunStream the caller can iterate to drive the run.
"""
values_t = ValuesTransformer()
mux = StreamMux(
@@ -105,11 +121,23 @@ class StreamingHandler:
) -> AsyncGraphRunStream:
"""Start an async streaming run.
Returns an `AsyncGraphRunStream` immediately. A background asyncio
Returns an AsyncGraphRunStream immediately. A background asyncio
task pumps events from the graph into the transformer pipeline.
*max_events* caps retention of every ``EventLog`` / ``StreamChannel``
the mux binds — see ``stream()`` for the full semantics.
Args:
input: Graph input.
config: Optional runnable config forwarded to the graph.
interrupt_before: Nodes to interrupt before, if any.
interrupt_after: Nodes to interrupt after, if any.
transformers: User transformers appended after the built-in
`ValuesTransformer` and `MessagesTransformer`.
max_events: Caps retention of every EventLog and
StreamChannel the mux binds — see `stream()` for the
full semantics.
Returns:
An AsyncGraphRunStream whose projections can be awaited
concurrently while the background pump runs.
"""
values_t = ValuesTransformer()
mux = StreamMux(
@@ -7,10 +7,10 @@ from langgraph.stream._types import ProtocolEvent, StreamTransformer
class ValuesTransformer(StreamTransformer):
"""Captures values events and projects them into an iterable of state snapshots.
"""Capture values events as an iterable of state snapshots.
Native transformer — projection keys are exposed as direct attributes
on the run stream (e.g. ``run.values``).
Native transformer — projection keys are exposed as direct
attributes on the run stream (e.g. `run.values`).
"""
_native = True
@@ -41,14 +41,14 @@ class ValuesTransformer(StreamTransformer):
class MessagesTransformer(StreamTransformer):
"""Captures messages events and passes through raw (chunk, metadata) tuples.
"""Pass through raw (chunk, metadata) tuples from messages events.
This is the same shape as today's ``stream_mode="messages"`` output.
This is the same shape as today's `stream_mode="messages"` output.
A follow-on PR will replace this with a richer transformer that
produces ChatModelStream objects using the protocol handler.
Native transformer — projection keys are exposed as direct attributes
on the run stream (e.g. ``run.messages``).
Native transformer — projection keys are exposed as direct
attributes on the run stream (e.g. `run.messages`).
"""
_native = True