mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-24 00:22:25 +02:00
feat(langgraph): add streaming transformer infrastructure and tests
Introduces the StreamingHandler, StreamMux, EventLog, StreamChannel, and StreamTransformer abstractions for ergonomic streaming projections over compiled graphs. Includes ValuesTransformer and MessagesTransformer as built-in native projections, plus support for user-defined custom transformers.
This commit is contained in:
Nick Hollon
@@ -0,0 +1,21 @@
|
||||
"""Streaming infrastructure for LangGraph.
|
||||
|
||||
Provides a ``StreamingHandler`` that wraps a compiled graph and exposes
|
||||
ergonomic streaming projections through a transformer pipeline.
|
||||
"""
|
||||
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream
|
||||
from langgraph.stream.stream_channel import StreamChannel
|
||||
from langgraph.stream.streaming_handler import StreamingHandler
|
||||
|
||||
__all__ = [
|
||||
"AsyncGraphRunStream",
|
||||
"EventLog",
|
||||
"GraphRunStream",
|
||||
"ProtocolEvent",
|
||||
"StreamChannel",
|
||||
"StreamTransformer",
|
||||
"StreamingHandler",
|
||||
]
|
||||
@@ -0,0 +1,24 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
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).
|
||||
"""
|
||||
params: _ProtocolEventParams = {
|
||||
"namespace": list(part["ns"]),
|
||||
"data": part["data"],
|
||||
}
|
||||
if "interrupts" in part:
|
||||
params["interrupts"] = part["interrupts"]
|
||||
return {
|
||||
"type": "event",
|
||||
"method": part["type"],
|
||||
"params": params,
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class EventLog(Generic[T]):
|
||||
"""Append-only buffer with multi-cursor sync and async iteration.
|
||||
|
||||
Supports multiple independent consumers iterating the same log
|
||||
concurrently. Each call to ``__iter__`` or ``__aiter__`` creates a
|
||||
new cursor starting from the beginning.
|
||||
|
||||
A given instance should be used in either sync or async mode — the
|
||||
two notification paths are independent and do not interfere, but
|
||||
mixing them on one instance is not tested.
|
||||
|
||||
Producer API (thread-safe):
|
||||
push(item) — append an item, notify all waiting cursors
|
||||
close() — mark the log as done
|
||||
fail(err) — mark the log as errored
|
||||
|
||||
Consumer API:
|
||||
__iter__() — new sync cursor (blocks via threading.Condition)
|
||||
__aiter__() — new async cursor (awaits via asyncio.Future)
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._items: list[T] = []
|
||||
self._closed = False
|
||||
self._error: BaseException | None = None
|
||||
# Sync notification
|
||||
self._lock = threading.Lock()
|
||||
self._cond = threading.Condition(self._lock)
|
||||
# Async notification — futures created lazily by async cursors
|
||||
self._async_waiters: list[asyncio.Future[None]] = []
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Producer API
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def push(self, item: T) -> None:
|
||||
"""Append *item* and wake all waiting cursors."""
|
||||
with self._lock:
|
||||
if self._closed:
|
||||
raise RuntimeError("Cannot push to a closed EventLog")
|
||||
self._items.append(item)
|
||||
self._cond.notify_all()
|
||||
self._wake_async()
|
||||
|
||||
def close(self) -> None:
|
||||
"""Mark the log as complete — open cursors will finish cleanly."""
|
||||
with self._lock:
|
||||
self._closed = True
|
||||
self._cond.notify_all()
|
||||
self._wake_async()
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
"""Mark the log as errored — open cursors will raise *err*."""
|
||||
with self._lock:
|
||||
self._error = err
|
||||
self._closed = True
|
||||
self._cond.notify_all()
|
||||
self._wake_async()
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Sync iteration
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def __iter__(self) -> Iterator[T]:
|
||||
"""Return a new independent sync cursor over the log."""
|
||||
return self._sync_cursor()
|
||||
|
||||
def _sync_cursor(self) -> Iterator[T]:
|
||||
cursor = 0
|
||||
while True:
|
||||
with self._lock:
|
||||
# Wait until data is available or the log is done.
|
||||
while cursor >= len(self._items) and not self._closed:
|
||||
self._cond.wait()
|
||||
# Yield available items before raising errors, matching
|
||||
# the async cursor's behavior.
|
||||
if cursor < len(self._items):
|
||||
item = self._items[cursor]
|
||||
cursor += 1
|
||||
elif self._error is not None:
|
||||
raise self._error
|
||||
else:
|
||||
# closed and no more items
|
||||
return
|
||||
yield item
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Async iteration
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[T]:
|
||||
"""Return a new independent async cursor over the log."""
|
||||
return self._async_cursor()
|
||||
|
||||
async def _async_cursor(self) -> AsyncIterator[T]:
|
||||
cursor = 0
|
||||
while True:
|
||||
if cursor < len(self._items):
|
||||
yield self._items[cursor]
|
||||
cursor += 1
|
||||
elif self._closed:
|
||||
if self._error is not None:
|
||||
raise self._error
|
||||
return
|
||||
else:
|
||||
# Wait for notification from push/close/fail.
|
||||
loop = asyncio.get_running_loop()
|
||||
fut: asyncio.Future[None] = loop.create_future()
|
||||
self._async_waiters.append(fut)
|
||||
await fut
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _wake_async(self) -> None:
|
||||
"""Resolve all pending async futures (safe from any thread)."""
|
||||
waiters = self._async_waiters
|
||||
if not waiters:
|
||||
return
|
||||
self._async_waiters = []
|
||||
for fut in waiters:
|
||||
if not fut.done():
|
||||
try:
|
||||
fut.get_loop().call_soon_threadsafe(fut.set_result, None)
|
||||
except RuntimeError:
|
||||
# Event loop already closed — nothing to notify.
|
||||
pass
|
||||
@@ -0,0 +1,139 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
from langgraph.stream.stream_channel import StreamChannel
|
||||
|
||||
|
||||
class StreamMux:
|
||||
"""Central event dispatcher for the streaming infrastructure.
|
||||
|
||||
Owns the main ``EventLog[ProtocolEvent]`` and routes events through
|
||||
a transformer pipeline. StreamChannels discovered in transformer
|
||||
projections are auto-wired so that every ``push()`` also injects a
|
||||
``ProtocolEvent`` into the main log.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._events: EventLog[ProtocolEvent] = EventLog()
|
||||
self._transformers: list[StreamTransformer] = []
|
||||
self._channels: list[StreamChannel[Any]] = []
|
||||
self._seq = 0
|
||||
|
||||
def register(self, transformer: StreamTransformer) -> dict[str, Any]:
|
||||
"""Register a transformer and return its projection dict.
|
||||
|
||||
Calls ``transformer.init()``, stores the transformer for event
|
||||
processing, and returns the projection. StreamChannels in the
|
||||
projection are auto-wired.
|
||||
"""
|
||||
projection = transformer.init()
|
||||
if not isinstance(projection, dict):
|
||||
raise TypeError(
|
||||
f"StreamTransformer.init() must return a dict, "
|
||||
f"got {type(projection).__name__}"
|
||||
)
|
||||
self._transformers.append(transformer)
|
||||
self._wire_channels(projection)
|
||||
return projection
|
||||
|
||||
def push(self, event: ProtocolEvent) -> None:
|
||||
"""Route *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).
|
||||
|
||||
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
|
||||
earlier seq numbers than the original event, preserving
|
||||
monotonic ordering in the log.
|
||||
"""
|
||||
keep = True
|
||||
for transformer in self._transformers:
|
||||
if not transformer.process(event):
|
||||
keep = False
|
||||
if keep:
|
||||
self._seq += 1
|
||||
event["seq"] = self._seq
|
||||
self._events.push(event)
|
||||
|
||||
def close(self) -> None:
|
||||
"""Finalize all transformers, close all channels and the main log.
|
||||
|
||||
If any transformer's ``finalize()`` raises, the remaining
|
||||
transformers, channels, and the main log are still closed.
|
||||
The first error is re-raised after cleanup completes.
|
||||
"""
|
||||
first_error: BaseException | None = None
|
||||
for transformer in self._transformers:
|
||||
try:
|
||||
transformer.finalize()
|
||||
except BaseException as e:
|
||||
if first_error is None:
|
||||
first_error = e
|
||||
for ch in self._channels:
|
||||
ch._close()
|
||||
self._events.close()
|
||||
if first_error is not None:
|
||||
raise first_error
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
"""Fail all transformers, channels, and the main log.
|
||||
|
||||
If any transformer's ``fail()`` raises, the remaining
|
||||
transformers, channels, and the main log are still failed.
|
||||
"""
|
||||
for transformer in self._transformers:
|
||||
try:
|
||||
transformer.fail(err)
|
||||
except BaseException:
|
||||
pass
|
||||
for ch in self._channels:
|
||||
ch._fail(err)
|
||||
self._events.fail(err)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# StreamChannel auto-wiring
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _wire_channels(self, projection: dict[str, Any]) -> None:
|
||||
"""Find StreamChannel instances in *projection* and wire them."""
|
||||
for value in projection.values():
|
||||
if isinstance(value, StreamChannel):
|
||||
self._channels.append(value)
|
||||
channel_name = value.name
|
||||
|
||||
def _make_forward(name: str) -> Callable[[Any], None]:
|
||||
def _forward(item: Any) -> None:
|
||||
self._forward(name, item)
|
||||
|
||||
return _forward
|
||||
|
||||
value._wire(_make_forward(channel_name))
|
||||
|
||||
def _forward(self, channel_name: str, item: Any) -> None:
|
||||
"""Inject a ProtocolEvent for a StreamChannel push.
|
||||
|
||||
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.
|
||||
"""
|
||||
self._seq += 1
|
||||
event: ProtocolEvent = {
|
||||
"type": "event",
|
||||
"seq": self._seq,
|
||||
"method": f"custom:{channel_name}",
|
||||
"params": {
|
||||
"namespace": [],
|
||||
"data": item,
|
||||
},
|
||||
}
|
||||
self._events.push(event)
|
||||
@@ -0,0 +1,79 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Literal
|
||||
|
||||
from typing_extensions import NotRequired, TypedDict
|
||||
|
||||
|
||||
class _ProtocolEventParams(TypedDict):
|
||||
"""Parameters for a protocol event."""
|
||||
|
||||
namespace: list[str]
|
||||
data: Any
|
||||
interrupts: NotRequired[tuple[Any, ...]]
|
||||
|
||||
|
||||
class ProtocolEvent(TypedDict):
|
||||
"""A protocol event emitted by the streaming infrastructure.
|
||||
|
||||
Wraps a raw stream part (values, messages, custom, etc.) in a uniform
|
||||
envelope with a monotonic sequence number assigned by the StreamMux.
|
||||
"""
|
||||
|
||||
type: Literal["event"]
|
||||
seq: NotRequired[int]
|
||||
method: str # StreamMode value: "values", "messages", "custom", etc.
|
||||
params: _ProtocolEventParams
|
||||
|
||||
|
||||
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.).
|
||||
|
||||
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 `process`. The `finalize` and
|
||||
`fail` hooks are optional — the default implementations are no-ops.
|
||||
StreamChannel instances are auto-closed/failed by the mux regardless.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def init(self) -> Any:
|
||||
"""Return the projection dict.
|
||||
|
||||
Keys become entries in `run.extensions`. If the transformer has
|
||||
`_native = True`, keys are also set as direct attributes on the
|
||||
run stream.
|
||||
|
||||
StreamChannel instances in the return value are automatically
|
||||
wired by the StreamMux for protocol event auto-forwarding.
|
||||
"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
"""Process a protocol event.
|
||||
|
||||
Called for every event before it is appended to the main event log.
|
||||
Return False to suppress the event from the main log.
|
||||
"""
|
||||
...
|
||||
|
||||
def finalize(self) -> None:
|
||||
"""Called when the run ends normally.
|
||||
|
||||
Override to close EventLogs, resolve promises, or perform other
|
||||
teardown. StreamChannel instances are auto-closed by the mux.
|
||||
"""
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
"""Called when the run ends with an error.
|
||||
|
||||
Override to fail EventLogs, reject promises, or perform other
|
||||
teardown. StreamChannel instances are auto-failed by the mux.
|
||||
"""
|
||||
@@ -0,0 +1,119 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from collections.abc import AsyncIterator, Iterator
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._types import ProtocolEvent
|
||||
from langgraph.stream.transformers import ValuesTransformer
|
||||
|
||||
|
||||
class GraphRunStream:
|
||||
"""Sync run stream with transformer-driven projections.
|
||||
|
||||
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
|
||||
from the mux's main event log.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mux: StreamMux,
|
||||
extensions: dict[str, Any],
|
||||
values_transformer: ValuesTransformer,
|
||||
pump_thread: threading.Thread,
|
||||
) -> None:
|
||||
self._mux = mux
|
||||
self.extensions = extensions
|
||||
self._values_transformer = values_transformer
|
||||
self._pump_thread = pump_thread
|
||||
|
||||
@property
|
||||
def output(self) -> dict[str, Any] | None:
|
||||
"""Block until the run completes and return the final state."""
|
||||
self._pump_thread.join()
|
||||
if self._values_transformer._log._error is not None:
|
||||
raise self._values_transformer._log._error
|
||||
return self._values_transformer._latest
|
||||
|
||||
@property
|
||||
def interrupted(self) -> bool:
|
||||
"""Block until the run completes, then return whether it was interrupted."""
|
||||
self._pump_thread.join()
|
||||
return self._values_transformer._interrupted
|
||||
|
||||
@property
|
||||
def interrupts(self) -> list[Any]:
|
||||
"""Block until the run completes, then return interrupt payloads."""
|
||||
self._pump_thread.join()
|
||||
return self._values_transformer._interrupts
|
||||
|
||||
def __iter__(self) -> Iterator[ProtocolEvent]:
|
||||
"""Iterate all protocol events from the mux's main event log."""
|
||||
return iter(self._mux._events)
|
||||
|
||||
|
||||
class AsyncGraphRunStream:
|
||||
"""Async run stream with transformer-driven projections.
|
||||
|
||||
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
|
||||
from the mux's main event log.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
mux: StreamMux,
|
||||
extensions: dict[str, Any],
|
||||
values_transformer: ValuesTransformer,
|
||||
pump_task: asyncio.Task[None],
|
||||
) -> None:
|
||||
self._mux = mux
|
||||
self.extensions = extensions
|
||||
self._values_transformer = values_transformer
|
||||
self._pump_task = pump_task
|
||||
|
||||
@property
|
||||
def output(self) -> Any:
|
||||
"""Return an awaitable that resolves to the final state.
|
||||
|
||||
Usage::
|
||||
|
||||
output = await run.output
|
||||
"""
|
||||
return self._get_output()
|
||||
|
||||
async def _get_output(self) -> dict[str, Any] | None:
|
||||
try:
|
||||
await self._pump_task
|
||||
except BaseException:
|
||||
pass
|
||||
if self._values_transformer._log._error is not None:
|
||||
raise self._values_transformer._log._error
|
||||
return self._values_transformer._latest
|
||||
|
||||
@property
|
||||
def interrupted(self) -> bool:
|
||||
"""Whether the run was interrupted.
|
||||
|
||||
Only meaningful after the run has completed (after consuming the
|
||||
stream or awaiting ``output``).
|
||||
"""
|
||||
return self._values_transformer._interrupted
|
||||
|
||||
@property
|
||||
def interrupts(self) -> list[Any]:
|
||||
"""Interrupt payloads, populated when interrupted is True."""
|
||||
return self._values_transformer._interrupts
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[ProtocolEvent]:
|
||||
"""Iterate all protocol events from the mux's main event log."""
|
||||
return self._mux._events.__aiter__()
|
||||
@@ -0,0 +1,63 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import AsyncIterator, Callable, Iterator
|
||||
from typing import Generic, TypeVar
|
||||
|
||||
from langgraph.stream._event_log import EventLog
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
class StreamChannel(Generic[T]):
|
||||
"""A named projection channel with optional protocol auto-forwarding.
|
||||
|
||||
Wraps an `EventLog[T]` 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``.
|
||||
|
||||
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>")``.
|
||||
|
||||
Lifecycle (``_close`` / ``_fail``) is managed by the mux — transformers
|
||||
using only StreamChannels don't need ``finalize`` / ``fail`` hooks.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str) -> None:
|
||||
self.name = name
|
||||
self._log: EventLog[T] = EventLog()
|
||||
self._wire_fn: Callable[[T], None] | None = None
|
||||
|
||||
def push(self, item: T) -> None:
|
||||
"""Append *item* to the log and auto-forward if wired."""
|
||||
self._log.push(item)
|
||||
if self._wire_fn is not None:
|
||||
self._wire_fn(item)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Mux lifecycle hooks (not called by transformers directly)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def _wire(self, fn: Callable[[T], None]) -> None:
|
||||
"""Install the auto-forward callback (called by StreamMux)."""
|
||||
self._wire_fn = fn
|
||||
|
||||
def _close(self) -> None:
|
||||
"""Close the underlying log (called by StreamMux on run end)."""
|
||||
self._log.close()
|
||||
|
||||
def _fail(self, err: BaseException) -> None:
|
||||
"""Fail the underlying log (called by StreamMux on run error)."""
|
||||
self._log.fail(err)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Iteration — delegates to the inner EventLog (multi-cursor)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
def __iter__(self) -> Iterator[T]:
|
||||
return iter(self._log)
|
||||
|
||||
def __aiter__(self) -> AsyncIterator[T]:
|
||||
return self._log.__aiter__()
|
||||
@@ -0,0 +1,161 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from collections.abc import Sequence
|
||||
from typing import Any
|
||||
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph.stream._convert import convert_to_protocol_event
|
||||
from langgraph.stream._mux import StreamMux
|
||||
from langgraph.stream._types import StreamTransformer
|
||||
from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream
|
||||
from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer
|
||||
from langgraph.types import All, StreamMode
|
||||
|
||||
# All stream modes to request from the graph.
|
||||
STREAM_V2_MODES: list[StreamMode] = [
|
||||
"values",
|
||||
"updates",
|
||||
"messages",
|
||||
"custom",
|
||||
"checkpoints",
|
||||
"tasks",
|
||||
"debug",
|
||||
]
|
||||
|
||||
|
||||
class StreamingHandler:
|
||||
"""Wraps a compiled graph and provides ergonomic streaming projections.
|
||||
|
||||
Usage::
|
||||
|
||||
handler = StreamingHandler(graph)
|
||||
|
||||
# Sync
|
||||
run = handler.stream(input_data)
|
||||
for state in run.values:
|
||||
print(state)
|
||||
output = run.output
|
||||
|
||||
# Async
|
||||
run = await handler.astream(input_data)
|
||||
async for state in run.values:
|
||||
print(state)
|
||||
output = await run.output
|
||||
"""
|
||||
|
||||
def __init__(self, graph: Any) -> None:
|
||||
self._graph = graph
|
||||
|
||||
def stream(
|
||||
self,
|
||||
input: Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
) -> GraphRunStream:
|
||||
"""Start a sync streaming run.
|
||||
|
||||
Returns a `GraphRunStream` immediately. A background daemon thread
|
||||
pumps events from the graph into the transformer pipeline.
|
||||
"""
|
||||
mux, extensions, native_keys, values_t = self._setup(transformers)
|
||||
|
||||
def pump() -> None:
|
||||
try:
|
||||
for part in self._graph.stream(
|
||||
input,
|
||||
config,
|
||||
stream_mode=STREAM_V2_MODES,
|
||||
subgraphs=True,
|
||||
version="v2",
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
):
|
||||
mux.push(convert_to_protocol_event(part))
|
||||
mux.close()
|
||||
except BaseException as e:
|
||||
mux.fail(e)
|
||||
|
||||
thread = threading.Thread(target=pump, daemon=True)
|
||||
thread.start()
|
||||
|
||||
run = GraphRunStream(mux, extensions, values_t, thread)
|
||||
for key in native_keys:
|
||||
setattr(run, key, extensions[key])
|
||||
return run
|
||||
|
||||
async def astream(
|
||||
self,
|
||||
input: Any,
|
||||
config: RunnableConfig | None = None,
|
||||
*,
|
||||
interrupt_before: All | Sequence[str] | None = None,
|
||||
interrupt_after: All | Sequence[str] | None = None,
|
||||
transformers: list[StreamTransformer] | None = None,
|
||||
) -> AsyncGraphRunStream:
|
||||
"""Start an async streaming run.
|
||||
|
||||
Returns an `AsyncGraphRunStream` immediately. A background asyncio
|
||||
task pumps events from the graph into the transformer pipeline.
|
||||
"""
|
||||
mux, extensions, native_keys, values_t = self._setup(transformers)
|
||||
|
||||
async def pump() -> None:
|
||||
try:
|
||||
async for part in self._graph.astream(
|
||||
input,
|
||||
config,
|
||||
stream_mode=STREAM_V2_MODES,
|
||||
subgraphs=True,
|
||||
version="v2",
|
||||
interrupt_before=interrupt_before,
|
||||
interrupt_after=interrupt_after,
|
||||
):
|
||||
mux.push(convert_to_protocol_event(part))
|
||||
mux.close()
|
||||
except BaseException as e:
|
||||
mux.fail(e)
|
||||
|
||||
task = asyncio.create_task(pump())
|
||||
|
||||
run = AsyncGraphRunStream(mux, extensions, values_t, task)
|
||||
for key in native_keys:
|
||||
setattr(run, key, extensions[key])
|
||||
return run
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _setup(
|
||||
user_transformers: list[StreamTransformer] | None,
|
||||
) -> tuple[StreamMux, dict[str, Any], set[str], ValuesTransformer]:
|
||||
"""Create the mux, register all transformers.
|
||||
|
||||
Returns (mux, extensions, native_keys, values_transformer).
|
||||
"""
|
||||
mux = StreamMux()
|
||||
|
||||
values_t = ValuesTransformer()
|
||||
messages_t = MessagesTransformer()
|
||||
|
||||
all_transformers: list[StreamTransformer] = [values_t, messages_t]
|
||||
if user_transformers:
|
||||
all_transformers.extend(user_transformers)
|
||||
|
||||
extensions: dict[str, Any] = {}
|
||||
native_keys: set[str] = set()
|
||||
|
||||
for t in all_transformers:
|
||||
projection = mux.register(t)
|
||||
extensions.update(projection)
|
||||
if getattr(t, "_native", False):
|
||||
native_keys.update(projection.keys())
|
||||
|
||||
return mux, extensions, native_keys, values_t
|
||||
@@ -0,0 +1,82 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from langgraph.stream._event_log import EventLog
|
||||
from langgraph.stream._types import ProtocolEvent, StreamTransformer
|
||||
|
||||
|
||||
class ValuesTransformer(StreamTransformer):
|
||||
"""Captures values events and projects them into an iterable of state snapshots.
|
||||
|
||||
Native transformer — projection keys are exposed as direct attributes
|
||||
on the run stream (e.g. ``run.values``).
|
||||
"""
|
||||
|
||||
_native = True
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._log: EventLog[dict[str, Any]] = EventLog()
|
||||
self._latest: dict[str, Any] | None = None
|
||||
self._interrupted = False
|
||||
self._interrupts: list[Any] = []
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"values": self._log}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] != "values":
|
||||
return True
|
||||
params = event["params"]
|
||||
# Only capture root namespace events
|
||||
if params["namespace"]:
|
||||
return True
|
||||
self._latest = params["data"]
|
||||
self._log.push(params["data"])
|
||||
interrupts = params.get("interrupts", ())
|
||||
if interrupts:
|
||||
self._interrupted = True
|
||||
self._interrupts.extend(interrupts)
|
||||
return True
|
||||
|
||||
def finalize(self) -> None:
|
||||
self._log.close()
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
self._log.fail(err)
|
||||
|
||||
|
||||
class MessagesTransformer(StreamTransformer):
|
||||
"""Captures messages events and passes through raw (chunk, metadata) tuples.
|
||||
|
||||
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 = True
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._log: EventLog[tuple[Any, dict[str, Any]]] = EventLog()
|
||||
|
||||
def init(self) -> dict[str, Any]:
|
||||
return {"messages": self._log}
|
||||
|
||||
def process(self, event: ProtocolEvent) -> bool:
|
||||
if event["method"] != "messages":
|
||||
return True
|
||||
params = event["params"]
|
||||
# Only capture root namespace events
|
||||
if params["namespace"]:
|
||||
return True
|
||||
self._log.push(params["data"])
|
||||
return True
|
||||
|
||||
def finalize(self) -> None:
|
||||
self._log.close()
|
||||
|
||||
def fail(self, err: BaseException) -> None:
|
||||
self._log.fail(err)
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user