diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py index 9f0c161c5..294e81f28 100644 --- a/libs/langgraph/langgraph/stream/__init__.py +++ b/libs/langgraph/langgraph/stream/__init__.py @@ -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. """ diff --git a/libs/langgraph/langgraph/stream/_convert.py b/libs/langgraph/langgraph/stream/_convert.py index bb2afb14d..1274fa7e7 100644 --- a/libs/langgraph/langgraph/stream/_convert.py +++ b/libs/langgraph/langgraph/stream/_convert.py @@ -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"]), diff --git a/libs/langgraph/langgraph/stream/_event_log.py b/libs/langgraph/langgraph/stream/_event_log.py index 6008d53af..01c22a442 100644 --- a/libs/langgraph/langgraph/stream/_event_log.py +++ b/libs/langgraph/langgraph/stream/_event_log.py @@ -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. " diff --git a/libs/langgraph/langgraph/stream/_mux.py b/libs/langgraph/langgraph/stream/_mux.py index 8677e863c..0481cd3eb 100644 --- a/libs/langgraph/langgraph/stream/_mux.py +++ b/libs/langgraph/langgraph/stream/_mux.py @@ -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 = { diff --git a/libs/langgraph/langgraph/stream/_types.py b/libs/langgraph/langgraph/stream/_types.py index c13c46582..60e3b5994 100644 --- a/libs/langgraph/langgraph/stream/_types.py +++ b/libs/langgraph/langgraph/stream/_types.py @@ -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 diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index c278360ff..c17d5db76 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -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() diff --git a/libs/langgraph/langgraph/stream/stream_channel.py b/libs/langgraph/langgraph/stream/stream_channel.py index c5b079e29..a863c556f 100644 --- a/libs/langgraph/langgraph/stream/stream_channel.py +++ b/libs/langgraph/langgraph/stream/stream_channel.py @@ -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:")``. + 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:")`. - 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:` 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) diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/streaming_handler.py index db88ff26b..373d74c72 100644 --- a/libs/langgraph/langgraph/stream/streaming_handler.py +++ b/libs/langgraph/langgraph/stream/streaming_handler.py @@ -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( diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index 59a470230..1c6805a39 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -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