mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-30 11:49:38 +02:00
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:
@@ -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.
|
||||
"""
|
||||
|
||||
|
||||
@@ -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"]),
|
||||
|
||||
@@ -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. "
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user