WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-06 09:47:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(core, prebuilt): tools-channel streaming and extensible GraphStreamer

Browse Source 
Adds a first-class `tools` stream mode with `tool-started` /
`tool-output-delta` / `tool-finished` / `tool-error` events, exposed as
`run.tool_calls` (an EventLog of ToolCallStream handles) through the new
`ToolCallTransformer`. Tool authors call
`langgraph.config.emit_tool_output_delta(chunk)` from inside a tool body
to stream partial output; outside a tool call it is a silent no-op.

`StreamTransformer` now declares `required_stream_modes`, and
`GraphStreamer` computes the request set purely as the union of those
declarations — no hardcoded base set. Built-in transformers
(Values/Messages/Subgraph) migrated to the new scheme; raw channels
like `custom` / `updates` / `checkpoints` / `tasks` / `debug` are
opt-in via a transformer that declares them.

Renames `StreamingHandler` → `GraphStreamer` (and file/test) and makes
it subclassable:

- `builtin_factories: ClassVar[tuple[TransformerFactory, ...]]` — extend
  to bundle default transformers (an `AgentStreamer` appends
  `ToolCallTransformer` here).
- `_make_run_stream` / `_make_async_run_stream` — override to return a
  `GraphRunStream` / `AsyncGraphRunStream` subclass with typed
  accessors over the extra projections.

Drops the now-redundant `values_transformer` arg from `GraphRunStream` /
`AsyncGraphRunStream` constructors — `output` / `interrupted` /
`interrupts` resolve the transformer lazily off the mux via a shared
`BaseRunStream._values_transformer` property.
This commit is contained in:
Nick Hollon
2026-04-20 12:13:08 -04:00
parent bebcd20815
commit e01037d90c
19 changed files with 1529 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/_internal/_constants.py
+1 -1
View File
@@ -68,7 +68,7 @@ CONFIG_KEY_RESUME_MAP = sys.intern("__pregel_resume_map")
# holds a mapping of task ns -> resume value for resuming tasks
CONFIG_KEY_STREAM_MESSAGES_V2 = sys.intern("__pregel_stream_messages_v2")
# when True, attach StreamMessagesHandlerV2 so content-block (v2) events
# flow through stream_mode="messages"; set by StreamingHandler only.
# flow through stream_mode="messages"; set by GraphStreamer only.
# --- Other constants ---
PUSH = sys.intern("__pregel_push")
libs/langgraph/langgraph/config.py
+41
View File
@@ -1,5 +1,7 @@
import asyncio
import sys
from collections.abc import Callable
from contextvars import ContextVar
from typing import Any
from langchain_core.runnables import RunnableConfig
@@ -9,6 +11,18 @@ from langgraph.store.base import BaseStore
from langgraph._internal._constants import CONF, CONFIG_KEY_RUNTIME
from langgraph.types import StreamWriter
_tool_call_writer: ContextVar[Callable[[Any], None] | None] = ContextVar(
"langgraph_tool_call_writer", default=None
)
"""ContextVar holding the writer for the currently-executing tool call.
Set by `StreamToolCallHandler.on_tool_start` and reset on end/error.
Defined here (rather than alongside the handler in `pregel/_tools.py`)
so `emit_tool_output_delta` can import it without triggering the
pregel import chain — user tool code does
`from langgraph.config import emit_tool_output_delta` at import time.
"""
def _no_op_stream_writer(c: Any) -> None:
pass
@@ -194,3 +208,30 @@ def get_stream_writer() -> StreamWriter:
"""
runtime = get_config()[CONF][CONFIG_KEY_RUNTIME]
return runtime.stream_writer
def emit_tool_output_delta(delta: Any) -> None:
"""Emit a `tool-output-delta` event onto the `tools` stream mode.
Must be called from inside a tool's execution scope (sync or async).
While a tool is running, `StreamToolCallHandler.on_tool_start` sets a
writer closure on a ContextVar keyed to that call's `tool_call_id`
and namespace; this helper reads the ContextVar and forwards `delta`
through it.
When called outside any tool call, or when the graph was not
streamed with `"tools"` in `stream_mode`, this is a silent no-op —
tool authors can leave `emit_tool_output_delta` calls in place
without gating them on stream mode.
Args:
delta: The partial output chunk to stream. Shape is up to the
caller — strings are the common case, but any JSON-
serializable value is accepted and surfaced as-is on the
`tools` channel's `tool-output-delta` payload under
`"delta"`.
"""
writer = _tool_call_writer.get()
if writer is None:
return
writer(delta)
libs/langgraph/langgraph/pregel/_messages.py
+1 -1
View File
@@ -273,7 +273,7 @@ class StreamMessagesHandlerV2(StreamMessagesHandler, _V2StreamingCallbackHandler
which forwards protocol events onto the messages stream channel.
Pregel attaches this class instead of the v1 handler only when
`StreamingHandler` opts in via the internal
`GraphStreamer` opts in via the internal
`CONFIG_KEY_STREAM_MESSAGES_V2` config key; direct
`graph.stream(stream_mode="messages")` callers keep the v1
AIMessageChunk shape.
libs/langgraph/langgraph/pregel/_tools.py
+223
View File
@@ -0,0 +1,223 @@
from __future__ import annotations
from collections.abc import AsyncIterator, Callable, Iterator
from contextvars import Token
from typing import Any, TypeVar, cast
from uuid import UUID
from langchain_core.callbacks import BaseCallbackHandler
from langgraph._internal._constants import NS_SEP
from langgraph.config import _tool_call_writer
from langgraph.pregel.protocol import StreamChunk
try:
from langchain_core.tracers._streaming import _StreamingCallbackHandler
except ImportError:
_StreamingCallbackHandler = object # type: ignore[assignment,misc]
T = TypeVar("T")
ToolCallWriter = Callable[[Any], None]
"""A closure bound to a single tool call that emits `tool-output-delta` events."""
class StreamToolCallHandler(BaseCallbackHandler, _StreamingCallbackHandler):
"""Callback handler that emits tool-call lifecycle events on the stream.
Fires on LangChain's `on_tool_*` callbacks and pushes to the `tools`
stream mode. Emits `tool-started` / `tool-output-delta` /
`tool-finished` / `tool-error` payloads keyed by `tool_call_id`.
While a tool is executing, this handler sets `_tool_call_writer` to a
closure bound to that call's namespace and `tool_call_id`. The
`emit_tool_output_delta` helper in `langgraph.config` reads that
ContextVar so tool bodies can stream partial output without threading
the writer through their own signature.
Attached by `Pregel.stream` / `astream` when `"tools"` is in
`stream_modes`. `run_inline = True` keeps event ordering
deterministic.
"""
run_inline = True
def __init__(self, stream: Callable[[StreamChunk], None]) -> None:
"""Initialize the handler.
Args:
stream: Callable that accepts a `StreamChunk` tuple
`(namespace, mode, payload)` and enqueues it.
"""
self.stream = stream
# run_id → (namespace, tool_call_id, ContextVar token)
# `on_tool_end` does not receive `tool_call_id` in kwargs, so
# we correlate by `run_id` which is present on every callback.
self._run_to_call: dict[
UUID, tuple[tuple[str, ...], str, Token[ToolCallWriter | None]]
] = {}
@staticmethod
def _containing_ns_from_metadata(
metadata: dict[str, Any] | None,
) -> tuple[str, ...]:
"""Return the namespace of the subgraph that contains this tool call.
`langgraph_checkpoint_ns` on a tool's callback metadata ends with
the `node_name:task_id` segment of the node that invoked the
tool. Dropping that segment gives the subgraph's own namespace,
which matches what other `tools` / `lifecycle` / `messages`
emitters use.
"""
if not metadata:
return ()
nskey = metadata.get("langgraph_checkpoint_ns")
if not nskey:
return ()
return tuple(cast(str, nskey).split(NS_SEP))[:-1]
def _start(
self,
serialized: dict[str, Any] | None,
input_str: str,
*,
run_id: UUID,
metadata: dict[str, Any] | None,
inputs: dict[str, Any] | None,
kwargs: dict[str, Any],
) -> None:
tool_call_id = cast("str | None", kwargs.get("tool_call_id")) or str(run_id)
tool_name = (
(serialized or {}).get("name")
or cast("str | None", kwargs.get("name"))
or ""
)
ns = self._containing_ns_from_metadata(metadata)
def writer(delta: Any) -> None:
self.stream(
(
ns,
"tools",
{
"event": "tool-output-delta",
"tool_call_id": tool_call_id,
"delta": delta,
},
)
)
token = _tool_call_writer.set(writer)
self._run_to_call[run_id] = (ns, tool_call_id, token)
payload: dict[str, Any] = {
"event": "tool-started",
"tool_call_id": tool_call_id,
"tool_name": tool_name,
}
if inputs is not None:
payload["input"] = inputs
self.stream((ns, "tools", payload))
def _end(self, output: Any, *, run_id: UUID) -> None:
info = self._run_to_call.pop(run_id, None)
if info is None:
return
ns, tool_call_id, token = info
self._reset_writer(token)
self.stream(
(
ns,
"tools",
{
"event": "tool-finished",
"tool_call_id": tool_call_id,
"output": output,
},
)
)
def _error(self, error: BaseException, *, run_id: UUID) -> None:
info = self._run_to_call.pop(run_id, None)
if info is None:
return
ns, tool_call_id, token = info
self._reset_writer(token)
self.stream(
(
ns,
"tools",
{
"event": "tool-error",
"tool_call_id": tool_call_id,
"message": str(error),
},
)
)
def tap_output_aiter(
self, run_id: UUID, output: AsyncIterator[T]
) -> AsyncIterator[T]:
"""Pass-through — required by the `_StreamingCallbackHandler` protocol."""
return output
def tap_output_iter(self, run_id: UUID, output: Iterator[T]) -> Iterator[T]:
"""Pass-through — sync counterpart to `tap_output_aiter`."""
return output
@staticmethod
def _reset_writer(token: Token[ToolCallWriter | None]) -> None:
# Token is invalid if `on_tool_end` runs in a different context
# than `on_tool_start` (e.g. langchain may hand off to a thread
# worker without copying the context). Swallow that case; the
# ContextVar lifetime is bounded by the enclosing task anyway.
try:
_tool_call_writer.reset(token)
except ValueError:
pass
# ------------------------------------------------------------------
# Sync callbacks
# ------------------------------------------------------------------
def on_tool_start(
self,
serialized: dict[str, Any],
input_str: str,
*,
run_id: UUID,
parent_run_id: UUID | None = None,
tags: list[str] | None = None,
metadata: dict[str, Any] | None = None,
inputs: dict[str, Any] | None = None,
**kwargs: Any,
) -> Any:
self._start(
serialized,
input_str,
run_id=run_id,
metadata=metadata,
inputs=inputs,
kwargs=kwargs,
)
def on_tool_end(
self,
output: Any,
*,
run_id: UUID,
parent_run_id: UUID | None = None,
**kwargs: Any,
) -> Any:
self._end(output, run_id=run_id)
def on_tool_error(
self,
error: BaseException,
*,
run_id: UUID,
parent_run_id: UUID | None = None,
**kwargs: Any,
) -> Any:
self._error(error, run_id=run_id)
libs/langgraph/langgraph/pregel/main.py
+13
View File
@@ -142,6 +142,7 @@ from langgraph.pregel._messages import (
from langgraph.pregel._read import DEFAULT_BOUND, PregelNode
from langgraph.pregel._retry import RetryPolicy
from langgraph.pregel._runner import PregelRunner
from langgraph.pregel._tools import StreamToolCallHandler
from langgraph.pregel._utils import get_new_channel_versions
from langgraph.pregel._validate import validate_graph, validate_keys
from langgraph.pregel._write import ChannelWrite, ChannelWriteEntry
@@ -2653,6 +2654,12 @@ class Pregel(
)
)
# set up tools stream mode
if "tools" in stream_modes:
run_manager.inheritable_handlers.append(
StreamToolCallHandler(stream.put)
)
# set up custom stream mode
if "custom" in stream_modes:
@@ -3044,6 +3051,12 @@ class Pregel(
)
)
# set up tools stream mode
if "tools" in stream_modes:
run_manager.inheritable_handlers.append(
StreamToolCallHandler(stream_put)
)
# set up custom stream mode
def stream_writer(c: Any) -> None:
aioloop.call_soon_threadsafe(
libs/langgraph/langgraph/stream/__init__.py
+3 -3
View File
@@ -1,14 +1,14 @@
"""Streaming infrastructure for LangGraph.
Provides a `StreamingHandler` that wraps a compiled graph and exposes
Provides a `GraphStreamer` 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.graph_streamer import GraphStreamer
from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream
from langgraph.stream.stream_channel import StreamChannel
from langgraph.stream.streaming_handler import StreamingHandler
__all__ = [
"AsyncGraphRunStream",
@@ -17,5 +17,5 @@ __all__ = [
"ProtocolEvent",
"StreamChannel",
"StreamTransformer",
"StreamingHandler",
"GraphStreamer",
]
libs/langgraph/langgraph/stream/_types.py
+7
View File
@@ -94,10 +94,17 @@ class StreamTransformer(ABC):
example, transformers that call `schedule()` from a sync
`process`). The mux also auto-detects the async lane when
`aprocess`, `afinalize`, or `afail` is overridden.
required_stream_modes: Stream modes the graph must emit for
this transformer to have anything to process. Computed as
the union across all registered transformers to determine
which modes a `GraphStreamer` run requests from the
graph. Empty tuple means the transformer consumes only
synthetic events (or is purely passive).
"""
requires_async: ClassVar[bool] = False
scope_exact: ClassVar[bool] = True
required_stream_modes: ClassVar[tuple[str, ...]] = ()
def __init__(self, scope: tuple[str, ...] = ()) -> None:
"""Initialize the transformer with its mux's scope.
libs/langgraph/langgraph/stream/streaming_handler.py → libs/langgraph/langgraph/stream/graph_streamer.py
+106 -45
View File
@@ -1,7 +1,7 @@
from __future__ import annotations
from collections.abc import Sequence
from typing import Any
from collections.abc import AsyncIterator, Iterator, Sequence
from typing import Any, ClassVar, cast
from langchain_core.runnables import RunnableConfig
@@ -33,7 +33,7 @@ def _coerce_factories(
for item in transformers or ():
if isinstance(item, StreamTransformer):
raise TypeError(
"StreamingHandler.transformers takes factories, not "
"GraphStreamer.transformers takes factories, not "
"pre-built instances. Pass the transformer class "
"(e.g. `MyTransformer`) or a callable taking `scope` "
"(e.g. `lambda scope: MyTransformer(scope, foo=...)`), "
@@ -41,20 +41,13 @@ def _coerce_factories(
)
if not callable(item):
raise TypeError(
f"StreamingHandler.transformers entries must be callable; "
f"GraphStreamer.transformers entries must be callable; "
f"got {type(item).__name__}."
)
coerced.append(item)
return coerced
_BUILTIN_FACTORIES: list[TransformerFactory] = [
ValuesTransformer,
MessagesTransformer,
SubgraphTransformer,
]
def _merge_v2_messages_flag(
config: RunnableConfig | None,
) -> RunnableConfig:
@@ -72,49 +65,117 @@ def _merge_v2_messages_flag(
return merged
# All stream modes to request from the graph.
STREAM_V2_MODES: list[StreamMode] = [
"values",
"updates",
"messages",
"custom",
"checkpoints",
"tasks",
"debug",
"lifecycle",
]
def _collect_stream_modes(mux: StreamMux) -> list[StreamMode]:
"""Return the union of `required_stream_modes` across registered transformers.
Transformers declare the stream modes they need to function, and
`GraphStreamer` asks the graph for exactly that union — no
hardcoded default set. If zero transformers are registered (or none
declares a given mode), the graph does not stream events for that
mode; consumers that want raw `custom` / `updates` / `checkpoints`
/ `tasks` / `debug` visibility must register a transformer that
declares those modes as required.
"""
modes: set[str] = set()
for transformer in mux._transformers:
modes.update(transformer.required_stream_modes)
return cast("list[StreamMode]", list(modes))
class StreamingHandler:
class GraphStreamer:
"""Wrap a compiled graph with ergonomic streaming projections.
Example:
```python
handler = StreamingHandler(graph)
streamer = GraphStreamer(graph)
# Sync
run = handler.stream(input_data)
run = streamer.stream(input_data)
for state in run.values:
print(state)
output = run.output
# Async — terminal accessors are methods so a missing `await`
# fails loudly instead of silently yielding a coroutine.
run = await handler.astream(input_data)
run = await streamer.astream(input_data)
async for state in run.values:
print(state)
output = await run.output()
```
Subclassing hooks:
- `builtin_factories`: tuple of transformer factories registered on
every run. Override to append domain-specific transformers
without the caller threading them through `transformers=`. An
`AgentStreamer` subclass, for example, can append
`ToolCallTransformer` so every agent run exposes `run.tool_calls`
by default.
- `_make_run_stream(...)` / `_make_async_run_stream(...)`: factory
hooks that return the run stream instance. Override to return a
subclass of `GraphRunStream` / `AsyncGraphRunStream` with typed
accessors over the projections contributed by the extra
transformers.
See the end of this module for a minimal subclassing sketch.
"""
builtin_factories: ClassVar[tuple[TransformerFactory, ...]] = (
ValuesTransformer,
MessagesTransformer,
SubgraphTransformer,
)
"""Factories registered on every run before caller-supplied transformers.
Subclasses append to this tuple to bundle domain-specific
transformers — for example, an `AgentStreamer` appends
`ToolCallTransformer` so every agent run exposes `run.tool_calls`
without the caller opting in.
"""
def __init__(self, graph: Pregel) -> None:
"""Initialize the handler.
"""Initialize the streamer.
Args:
graph: A compiled LangGraph graph to stream from.
"""
self._graph = graph
def _build_factories(
self,
transformers: list[TransformerFactory] | None,
) -> list[TransformerFactory]:
"""Return `builtin_factories` plus the caller's transformers.
Subclasses rarely override this directly — extend
`builtin_factories` instead. Override only when you need to
inspect or re-order the caller's transformers (e.g. to inject
a transformer *after* user-supplied ones).
"""
return [*self.builtin_factories, *_coerce_factories(transformers)]
def _make_run_stream(
self,
graph_iter: Iterator[Any],
mux: StreamMux,
) -> GraphRunStream:
"""Construct the sync run stream returned from `stream()`.
Override in a subclass to return a `GraphRunStream` subclass
with typed accessors over the extra projections contributed by
the subclass's `builtin_factories` (e.g. `AgentRunStream` with
a typed `.tool_calls`).
"""
return GraphRunStream(graph_iter, mux)
def _make_async_run_stream(
self,
graph_aiter: AsyncIterator[Any],
mux: StreamMux,
) -> AsyncGraphRunStream:
"""Async counterpart to `_make_run_stream`."""
return AsyncGraphRunStream(graph_aiter, mux)
def stream(
self,
input: Any,
@@ -136,24 +197,25 @@ class StreamingHandler:
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`.
transformers: User transformers appended after
`self.builtin_factories`.
Returns:
A GraphRunStream the caller can iterate to drive the run.
A `GraphRunStream` the caller can iterate to drive the run.
Subclasses may narrow this to a `GraphRunStream` subtype via
`_make_run_stream`.
"""
mux = StreamMux(
factories=_BUILTIN_FACTORIES + _coerce_factories(transformers),
factories=self._build_factories(transformers),
is_async=False,
)
values_t = mux.transformer_by_key("values")
assert isinstance(values_t, ValuesTransformer)
stream_modes = _collect_stream_modes(mux)
graph_iter = iter(
self._graph.stream(
input,
_merge_v2_messages_flag(config),
stream_mode=STREAM_V2_MODES,
stream_mode=stream_modes,
subgraphs=True,
version="v2",
interrupt_before=interrupt_before,
@@ -161,7 +223,7 @@ class StreamingHandler:
)
)
return GraphRunStream(graph_iter, mux, values_t)
return self._make_run_stream(graph_iter, mux)
async def astream(
self,
@@ -184,29 +246,28 @@ class StreamingHandler:
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`.
transformers: User transformers appended after
`self.builtin_factories`.
Returns:
An AsyncGraphRunStream whose projections can be awaited
concurrently; each subscribed cursor drives the pump when
its buffer is empty.
An `AsyncGraphRunStream` whose projections can be awaited
concurrently; subclasses may narrow this via
`_make_async_run_stream`.
"""
mux = StreamMux(
factories=_BUILTIN_FACTORIES + _coerce_factories(transformers),
factories=self._build_factories(transformers),
is_async=True,
)
values_t = mux.transformer_by_key("values")
assert isinstance(values_t, ValuesTransformer)
stream_modes = _collect_stream_modes(mux)
graph_aiter = self._graph.astream(
input,
_merge_v2_messages_flag(config),
stream_mode=STREAM_V2_MODES,
stream_mode=stream_modes,
subgraphs=True,
version="v2",
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
).__aiter__()
return AsyncGraphRunStream(graph_aiter, mux, values_t)
return self._make_async_run_stream(graph_aiter, mux)
libs/langgraph/langgraph/stream/run_stream.py
+40 -20
View File
@@ -47,6 +47,28 @@ class BaseRunStream:
for key in mux.native_keys:
setattr(self, key, mux.extensions[key])
@property
def _values_transformer(self) -> ValuesTransformer:
"""Look up the `ValuesTransformer` backing `output` / `interrupted`.
Resolved lazily off the mux so subclasses don't have to thread
it through their constructors. Raises if no `ValuesTransformer`
is registered — `output` / `interrupted` / `interrupts` have
nothing to return in that case, so failing loudly is better
than returning `None` silently.
"""
from langgraph.stream.transformers import ValuesTransformer
vt = self._mux.transformer_by_key("values")
if not isinstance(vt, ValuesTransformer):
raise RuntimeError(
"No ValuesTransformer is registered on this mux — "
"`output`, `interrupted`, and `interrupts` require one. "
"Add it to your GraphStreamer subclass's "
"`builtin_factories` or pass it via `transformers=`."
)
return vt
def __iter__(self) -> Iterator[ProtocolEvent]:
"""Sync iteration of protocol events on this mux's main log.
@@ -128,19 +150,18 @@ class GraphRunStream(BaseRunStream):
self,
graph_iter: Iterator[Any],
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`.
Must have a `ValuesTransformer` registered under the
`"values"` key — `output` / `interrupted` / `interrupts`
read from it lazily.
"""
super().__init__(mux)
self._graph_iter = graph_iter
self._values_transformer = values_transformer
self._exhausted = False
mux.bind_pump(self._pump_next)
@@ -196,10 +217,10 @@ class GraphRunStream(BaseRunStream):
def output(self) -> dict[str, Any] | None:
"""Drive the run to completion and return the final state."""
_drive_until_done(self._pump_next)
err = self._values_transformer.error
if err is not None:
raise err
return self._values_transformer._latest
vt = self._values_transformer
if vt.error is not None:
raise vt.error
return vt._latest
@property
def interrupted(self) -> bool:
@@ -209,10 +230,10 @@ class GraphRunStream(BaseRunStream):
BaseException: If the run ended with an error.
"""
_drive_until_done(self._pump_next)
err = self._values_transformer.error
if err is not None:
raise err
return self._values_transformer._interrupted
vt = self._values_transformer
if vt.error is not None:
raise vt.error
return vt._interrupted
@property
def interrupts(self) -> list[Any]:
@@ -222,10 +243,10 @@ class GraphRunStream(BaseRunStream):
BaseException: If the run ended with an error.
"""
_drive_until_done(self._pump_next)
err = self._values_transformer.error
if err is not None:
raise err
return self._values_transformer._interrupts
vt = self._values_transformer
if vt.error is not None:
raise vt.error
return vt._interrupts
class AsyncGraphRunStream(BaseRunStream):
@@ -256,19 +277,18 @@ class AsyncGraphRunStream(BaseRunStream):
self,
graph_aiter: AsyncIterator[Any],
mux: StreamMux,
values_transformer: ValuesTransformer,
) -> None:
"""Initialize the async run stream.
Args:
graph_aiter: Async 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`.
Must have a `ValuesTransformer` registered under the
`"values"` key — `output` / `interrupted` / `interrupts`
read from it lazily.
"""
super().__init__(mux)
self._graph_aiter = graph_aiter
self._values_transformer = values_transformer
self._exhausted = False
self._pump_lock = asyncio.Lock()
mux.bind_apump(self._apump_next)
libs/langgraph/langgraph/stream/transformers.py
+3
View File
@@ -49,6 +49,7 @@ class ValuesTransformer(StreamTransformer):
"""
_native = True
required_stream_modes = ("values",)
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
@@ -121,6 +122,7 @@ class MessagesTransformer(StreamTransformer):
"""
_native = True
required_stream_modes = ("messages",)
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
@@ -340,6 +342,7 @@ class SubgraphTransformer(StreamTransformer):
_native = True
scope_exact = False
required_stream_modes = ("lifecycle",)
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
libs/langgraph/langgraph/types.py
+2
View File
@@ -124,6 +124,7 @@ StreamMode = Literal[
"messages",
"custom",
"lifecycle",
"tools",
]
"""How the stream method should emit outputs.
@@ -137,6 +138,7 @@ StreamMode = Literal[
- `"tasks"`: Emit events when tasks start and finish, including their results and errors.
- `"debug"`: Emit `"checkpoints"` and `"tasks"` events for debugging purposes.
- `"lifecycle"`: Emit subgraph lifecycle events (`started`, `running`, `completed`, `failed`, `interrupted`) with payloads matching `LifecycleData`.
- `"tools"`: Emit tool-call lifecycle events (`tool-started`, `tool-output-delta`, `tool-finished`, `tool-error`) keyed by `tool_call_id`.
"""
StreamWriter = Callable[[Any], None]
libs/langgraph/tests/test_streaming_handler.py → libs/langgraph/tests/test_graph_streamer.py
+209 -54
View File
@@ -1,4 +1,4 @@
"""Tests for the StreamingHandler and its supporting infrastructure."""
"""Tests for the GraphStreamer and its supporting infrastructure."""
from __future__ import annotations
@@ -6,6 +6,7 @@ import asyncio
import operator
import sys
import time
from collections.abc import AsyncIterator, Iterator
from typing import Annotated, Any
import pytest
@@ -15,9 +16,11 @@ from typing_extensions import TypedDict
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from langgraph.stream import (
AsyncGraphRunStream,
EventLog,
GraphRunStream,
GraphStreamer,
StreamChannel,
StreamingHandler,
StreamTransformer,
)
from langgraph.stream._convert import convert_to_protocol_event
@@ -132,6 +135,24 @@ def _build_custom_stream_graph():
return builder.compile()
class _CustomPassthroughTransformer(StreamTransformer):
"""Opts a run into the `custom` stream mode without building a projection.
`GraphStreamer` requests only the modes that registered
transformers declare via `required_stream_modes`. Custom events are
raw user emissions from `StreamWriter`, so tests that want them
visible on the main event log register this pass-through transformer.
"""
required_stream_modes = ("custom",)
def init(self) -> dict[str, Any]:
return {}
def process(self, event: ProtocolEvent) -> bool:
return True
# ---------------------------------------------------------------------------
# EventLog unit tests
# ---------------------------------------------------------------------------
@@ -404,14 +425,14 @@ class TestStreamChannel:
# ---------------------------------------------------------------------------
# StreamingHandler sync tests
# GraphStreamer sync tests
# ---------------------------------------------------------------------------
class TestStreamingHandlerSync:
class TestGraphStreamerSync:
def test_values_projection(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
snapshots = list(run.values)
# Should have at least the initial + per-node snapshots.
@@ -423,7 +444,7 @@ class TestStreamingHandlerSync:
def test_output(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
output = run.output
assert output is not None
@@ -432,7 +453,7 @@ class TestStreamingHandlerSync:
def test_raw_event_iteration(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
events = list(run)
assert len(events) > 0
@@ -444,7 +465,7 @@ class TestStreamingHandlerSync:
def test_extensions_has_native_keys(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
# Drain events so the run completes.
_ = run.output
@@ -457,7 +478,7 @@ class TestStreamingHandlerSync:
def test_extensions_is_read_only(self) -> None:
"""`run.extensions` must reject mutations so users can't corrupt mux state."""
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
with pytest.raises(TypeError):
run.extensions["new_key"] = object() # type: ignore[index]
@@ -466,16 +487,34 @@ class TestStreamingHandlerSync:
def test_custom_stream_events(self) -> None:
graph = _build_custom_stream_graph()
handler = StreamingHandler(graph)
run = handler.stream({"value": "x", "items": []})
handler = GraphStreamer(graph)
run = handler.stream(
{"value": "x", "items": []},
transformers=[_CustomPassthroughTransformer],
)
custom_events = [e for e in run if e["method"] == "custom"]
assert len(custom_events) == 2
assert custom_events[0]["params"]["data"] == {"step": "start"}
assert custom_events[1]["params"]["data"] == {"step": "end"}
def test_custom_events_suppressed_without_transformer(self) -> None:
"""Without a transformer declaring `"custom"`, no custom events flow.
`GraphStreamer` asks the graph only for the modes that
registered transformers require. Built-ins cover
`values` / `messages` / `lifecycle`; consumers that want raw
custom events surface them by registering a transformer whose
`required_stream_modes` includes `"custom"`.
"""
graph = _build_custom_stream_graph()
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
custom_events = [e for e in run if e["method"] == "custom"]
assert custom_events == []
def test_interleave_values_and_messages(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
tagged = list(run.interleave("values", "messages"))
@@ -489,7 +528,7 @@ class TestStreamingHandlerSync:
def test_abort_marks_exhausted_and_closes_mux(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
values_iter = iter(run.values)
# Consume one item so the pump advances.
@@ -503,7 +542,7 @@ class TestStreamingHandlerSync:
def test_context_manager_calls_abort_on_exit(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
with handler.stream({"value": "x", "items": []}) as run:
values_iter = iter(run.values)
_ = next(values_iter)
@@ -511,30 +550,30 @@ class TestStreamingHandlerSync:
def test_interleave_unknown_projection(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
with pytest.raises(KeyError):
list(run.interleave("values", "does_not_exist"))
class TestStreamingHandlerSyncErrors:
class TestGraphStreamerSyncErrors:
def test_error_propagation_output(self) -> None:
graph = _build_error_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
_ = run.output
def test_error_propagation_values(self) -> None:
graph = _build_error_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
list(run.values)
def test_error_propagation_raw_events(self) -> None:
graph = _build_error_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
list(run)
@@ -542,7 +581,7 @@ class TestStreamingHandlerSyncErrors:
def test_error_propagation_interrupted(self) -> None:
"""`run.interrupted` should raise on a failed run, not silently return False."""
graph = _build_error_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
_ = run.interrupted
@@ -550,16 +589,16 @@ class TestStreamingHandlerSyncErrors:
def test_error_propagation_interrupts(self) -> None:
"""`run.interrupts` should raise on a failed run."""
graph = _build_error_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
_ = run.interrupts
class TestStreamingHandlerSyncInterrupt:
class TestGraphStreamerSyncInterrupt:
def test_interrupted(self) -> None:
graph = _build_interrupt_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream(
{"value": "x", "items": []},
{"configurable": {"thread_id": "t1"}},
@@ -570,16 +609,16 @@ class TestStreamingHandlerSyncInterrupt:
# ---------------------------------------------------------------------------
# StreamingHandler async tests
# GraphStreamer async tests
# ---------------------------------------------------------------------------
class TestStreamingHandlerAsync:
class TestGraphStreamerAsync:
@pytest.mark.anyio
@NEEDS_CONTEXTVARS
async def test_values_projection(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
snapshots = [s async for s in run.values]
assert len(snapshots) >= 1
@@ -591,7 +630,7 @@ class TestStreamingHandlerAsync:
@NEEDS_CONTEXTVARS
async def test_output(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
output = await run.output()
assert output is not None
@@ -602,7 +641,7 @@ class TestStreamingHandlerAsync:
@NEEDS_CONTEXTVARS
async def test_raw_event_iteration(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
events = [e async for e in run]
assert len(events) > 0
@@ -613,7 +652,7 @@ class TestStreamingHandlerAsync:
@NEEDS_CONTEXTVARS
async def test_abort_marks_exhausted_and_closes_mux(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
values_iter = aiter(run.values)
_ = await anext(values_iter)
@@ -628,7 +667,7 @@ class TestStreamingHandlerAsync:
@NEEDS_CONTEXTVARS
async def test_async_context_manager_calls_abort_on_exit(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
async with run:
values_iter = aiter(run.values)
@@ -639,7 +678,7 @@ class TestStreamingHandlerAsync:
@NEEDS_CONTEXTVARS
async def test_extensions_has_native_keys(self) -> None:
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
_ = await run.output()
assert "values" in run.extensions
@@ -648,12 +687,12 @@ class TestStreamingHandlerAsync:
assert run.messages is run.extensions["messages"]
class TestStreamingHandlerAsyncErrors:
class TestGraphStreamerAsyncErrors:
@pytest.mark.anyio
@NEEDS_CONTEXTVARS
async def test_error_propagation_output(self) -> None:
graph = _build_error_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
await run.output()
@@ -662,7 +701,7 @@ class TestStreamingHandlerAsyncErrors:
@NEEDS_CONTEXTVARS
async def test_error_propagation_values(self) -> None:
graph = _build_error_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
async for _ in run.values:
@@ -672,7 +711,7 @@ class TestStreamingHandlerAsyncErrors:
@NEEDS_CONTEXTVARS
async def test_error_propagation_raw_events(self) -> None:
graph = _build_error_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
async for _ in run:
@@ -683,7 +722,7 @@ class TestStreamingHandlerAsyncErrors:
async def test_error_propagation_interrupted(self) -> None:
"""`await run.interrupted()` should raise on a failed async run."""
graph = _build_error_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
await run.interrupted()
@@ -693,18 +732,18 @@ class TestStreamingHandlerAsyncErrors:
async def test_error_propagation_interrupts(self) -> None:
"""`await run.interrupts()` should raise on a failed async run."""
graph = _build_error_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
with pytest.raises(ValueError, match="boom"):
await run.interrupts()
class TestStreamingHandlerAsyncInterrupt:
class TestGraphStreamerAsyncInterrupt:
@pytest.mark.anyio
@NEEDS_CONTEXTVARS
async def test_interrupted(self) -> None:
graph = _build_interrupt_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream(
{"value": "x", "items": []},
{"configurable": {"thread_id": "t2"}},
@@ -714,13 +753,16 @@ class TestStreamingHandlerAsyncInterrupt:
assert len(await run.interrupts()) > 0
class TestStreamingHandlerAsyncCustom:
class TestGraphStreamerAsyncCustom:
@pytest.mark.anyio
@NEEDS_CONTEXTVARS
async def test_custom_stream_events(self) -> None:
graph = _build_custom_stream_graph()
handler = StreamingHandler(graph)
run = await handler.astream({"value": "x", "items": []})
handler = GraphStreamer(graph)
run = await handler.astream(
{"value": "x", "items": []},
transformers=[_CustomPassthroughTransformer],
)
events = [e async for e in run]
custom_events = [e for e in events if e["method"] == "custom"]
assert len(custom_events) == 2
@@ -1105,7 +1147,7 @@ class TestCustomTransformer:
return True
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
counter_t = CounterTransformer()
run = handler.stream(
{"value": "x", "items": []},
@@ -1139,7 +1181,7 @@ class TestCustomTransformer:
return True
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
foo_t = FooTransformer()
run = handler.stream(
{"value": "x", "items": []},
@@ -1172,7 +1214,7 @@ class TestCustomTransformer:
return True
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream(
{"value": "x", "items": []},
transformers=[lambda _scope: EmitterTransformer()],
@@ -1232,7 +1274,7 @@ class TestCustomTransformer:
return True
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
with pytest.raises(
ValueError,
match=r"conflict.*'values'.*ValuesTransformer",
@@ -1334,7 +1376,7 @@ class TestEventLogAutoLifecycle:
return True
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
t = MinimalTransformer()
run = handler.stream(
{"value": "x", "items": []},
@@ -1638,7 +1680,7 @@ class TestAsyncTransformerLane:
self._log.close()
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
scorer = Scorer()
run = await handler.astream(
{"value": "x", "items": []},
@@ -1665,7 +1707,7 @@ class TestMemoryBounds:
"""With a single sync consumer, the pump produces exactly one event
per cursor advance, so the buffer never holds more than one."""
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
events_iter = iter(run)
max_buffered = 0
@@ -1683,7 +1725,7 @@ class TestMemoryBounds:
"""Projections without a subscriber drop pushes silently —
their buffers stay empty regardless of run length."""
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
# Subscribe to main events only; leave values and messages unsubscribed.
list(run)
@@ -1699,7 +1741,7 @@ class TestMemoryBounds:
process() without populating the log, so the values log buffer
stays empty even across a full run."""
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
_ = run.output
values_log = run.extensions["values"]
@@ -1709,7 +1751,7 @@ class TestMemoryBounds:
def test_drained_subscriber_buffer_returns_to_empty(self) -> None:
"""After fully draining a subscribed log, the internal deque is empty."""
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "x", "items": []})
values_log = run.extensions["values"]
list(run.values)
@@ -1720,7 +1762,7 @@ class TestMemoryBounds:
async def test_async_single_consumer_buffer_stays_at_most_one(self) -> None:
"""Same drain-on-consume guarantee for the async lane."""
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
max_buffered = 0
count = 0
@@ -1735,7 +1777,7 @@ class TestMemoryBounds:
async def test_async_unsubscribed_projections_never_accumulate(self) -> None:
"""Projections with no async subscriber stay empty under astream."""
graph = _build_simple_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "x", "items": []})
_ = await run.output()
values_log = run.extensions["values"]
@@ -1787,3 +1829,116 @@ class TestDrainOnConsume:
log.close()
assert list(a) == [0, 1, 2]
assert list(b) == [0, 1, 2]
# ---------------------------------------------------------------------------
# Subclassing hooks
# ---------------------------------------------------------------------------
class _CountingTransformer(StreamTransformer):
"""Record how many events the transformer saw — nothing more."""
required_stream_modes = ("values",)
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._log: EventLog[int] = EventLog()
self.count = 0
def init(self) -> dict[str, Any]:
return {"counts": self._log}
def process(self, event: ProtocolEvent) -> bool:
if event["method"] == "values":
self.count += 1
return True
class _CustomRunStream(GraphRunStream):
"""Trivial subclass used to prove the `_make_run_stream` hook fires."""
marker = "custom-sync"
class _CustomAsyncRunStream(AsyncGraphRunStream):
marker = "custom-async"
class _CustomStreamer(GraphStreamer):
"""`GraphStreamer` subclass that injects a transformer and narrows run types."""
builtin_factories = (
*GraphStreamer.builtin_factories,
_CountingTransformer,
)
def _make_run_stream(
self,
graph_iter: Iterator[Any],
mux: StreamMux,
) -> _CustomRunStream:
return _CustomRunStream(graph_iter, mux)
def _make_async_run_stream(
self,
graph_aiter: AsyncIterator[Any],
mux: StreamMux,
) -> _CustomAsyncRunStream:
return _CustomAsyncRunStream(graph_aiter, mux)
class TestGraphStreamerSubclassing:
def test_subclass_injects_transformer_without_user_opt_in(self) -> None:
graph = _build_simple_graph()
streamer = _CustomStreamer(graph)
run = streamer.stream({"value": "", "items": []})
# Subclass-injected transformer is present in every run's mux
# without the caller passing `transformers=`.
assert isinstance(run, _CustomRunStream)
assert run.marker == "custom-sync"
counting = run._mux.transformer_by_key("counts")
assert isinstance(counting, _CountingTransformer)
run.output # drive to completion
assert counting.count >= 1
@pytest.mark.anyio
async def test_async_subclass_returns_custom_async_stream(self) -> None:
graph = _build_simple_graph()
streamer = _CustomStreamer(graph)
run = await streamer.astream({"value": "", "items": []})
assert isinstance(run, _CustomAsyncRunStream)
assert run.marker == "custom-async"
counting = run._mux.transformer_by_key("counts")
assert isinstance(counting, _CountingTransformer)
await run.output()
assert counting.count >= 1
def test_user_transformers_appended_after_builtin_factories(self) -> None:
"""User transformers run after `builtin_factories` in registration order."""
class _UserTransformer(StreamTransformer):
required_stream_modes = ()
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
def init(self) -> dict[str, Any]:
return {"user_flag": EventLog()}
def process(self, event: ProtocolEvent) -> bool:
return True
graph = _build_simple_graph()
streamer = _CustomStreamer(graph)
run = streamer.stream(
{"value": "", "items": []}, transformers=[_UserTransformer]
)
# Both the subclass-injected and user-supplied projections are
# available.
assert "counts" in run.extensions
assert "user_flag" in run.extensions
run.output # drain
libs/langgraph/tests/test_stream_messages_transformer.py
+21 -21
View File
@@ -23,8 +23,8 @@ from langgraph.constants import END, START
from langgraph.graph import MessagesState, StateGraph
from langgraph.stream._event_log import EventLog
from langgraph.stream._mux import StreamMux
from langgraph.stream.graph_streamer import GraphStreamer
from langgraph.stream.run_stream import GraphRunStream
from langgraph.stream.streaming_handler import StreamingHandler
from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer
TS = int(time.time() * 1000)
@@ -426,7 +426,7 @@ class TestWireRequestMore:
mux = StreamMux([values_t, messages_t], is_async=False)
assert messages_t._pump_fn is None
run = GraphRunStream(iter([]), mux, values_t)
run = GraphRunStream(iter([]), mux)
# After wire, the transformer's pump callback is set.
assert messages_t._pump_fn is not None
# And calling it invokes GraphRunStream._pump_next (drains an empty
@@ -439,7 +439,7 @@ class TestWireRequestMore:
messages_t = MessagesTransformer()
mux = StreamMux([values_t, messages_t], is_async=False)
GraphRunStream(iter([]), mux, values_t)
GraphRunStream(iter([]), mux)
log: EventLog[ChatModelStream] = mux.extensions["messages"]
log._subscribed = True
@@ -507,7 +507,7 @@ class TestViaMux:
# ---------------------------------------------------------------------------
# End-to-end: full graph → StreamingHandler → run.messages
# End-to-end: full graph → GraphStreamer → run.messages
# ---------------------------------------------------------------------------
@@ -541,7 +541,7 @@ class TestEndToEnd:
.compile()
)
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"messages": "hi"})
streams = list(run.messages)
@@ -565,7 +565,7 @@ class TestEndToEnd:
.compile()
)
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"messages": "go"})
# Pull the stream handle out, then iterate its text deltas.
@@ -587,7 +587,7 @@ class TestEndToEnd:
.compile()
)
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"messages": "hi"})
streams = list(run.messages)
@@ -611,7 +611,7 @@ class TestEndToEnd:
.compile()
)
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"messages": "hi"})
streams = []
@@ -649,7 +649,7 @@ class TestEndToEnd:
.compile()
)
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"messages": "hi"})
async def consume_nested() -> list[str]:
@@ -664,11 +664,11 @@ class TestEndToEnd:
class TestEndToEndV2Invoke:
"""Nodes call `model.invoke()`; `StreamingHandler` routes through v2.
"""Nodes call `model.invoke()`; `GraphStreamer` routes through v2.
Exercises the auto-routing path added in
`feat(core): route invoke through v2 event path for
_V2StreamingCallbackHandler`: `StreamingHandler` injects
_V2StreamingCallbackHandler`: `GraphStreamer` injects
`CONFIG_KEY_STREAM_MESSAGES_V2` into the config, pregel attaches
`StreamMessagesHandlerV2`, `BaseChatModel._should_stream_v2` sees the
v2 marker and drives the protocol event generator, and
@@ -690,7 +690,7 @@ class TestEndToEndV2Invoke:
.compile()
)
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"messages": "hi"})
streams = list(run.messages)
@@ -717,7 +717,7 @@ class TestEndToEndV2Invoke:
.compile()
)
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"messages": "go"})
(stream,) = list(run.messages)
@@ -751,7 +751,7 @@ class TestEndToEndV2Invoke:
.compile()
)
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"messages": "hi"})
(stream,) = list(run.messages)
@@ -779,7 +779,7 @@ class TestEndToEndV2Invoke:
.compile()
)
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"messages": "hi"})
streams = list(run.messages)
@@ -810,7 +810,7 @@ class TestEndToEndV2Invoke:
.compile()
)
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"messages": "hi"})
streams = list(run.messages)
@@ -823,7 +823,7 @@ class TestEndToEndV2Invoke:
@pytest.mark.anyio
async def test_ainvoke_with_v2_marker_populates_messages(self) -> None:
"""Async mirror: `model.ainvoke()` + `StreamingHandler.astream()`."""
"""Async mirror: `model.ainvoke()` + `GraphStreamer.astream()`."""
model = GenericFakeChatModel(messages=iter(["async invoke"]))
async def call_model(state: MessagesState) -> dict[str, Any]:
@@ -837,7 +837,7 @@ class TestEndToEndV2Invoke:
.compile()
)
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"messages": "hi"})
streams = []
@@ -852,8 +852,8 @@ class TestEndToEndV2Invoke:
class TestDirectMessagesModeStaysV1:
"""Regression guard: direct `graph.stream(stream_mode="messages")`
(no `StreamingHandler`) must keep the v1 `(AIMessageChunk, metadata)`
tuple shape. The v2 flag is only injected by `StreamingHandler`.
(no `GraphStreamer`) must keep the v1 `(AIMessageChunk, metadata)`
tuple shape. The v2 flag is only injected by `GraphStreamer`.
"""
def test_direct_graph_stream_messages_yields_ai_message_chunks(self) -> None:
@@ -878,7 +878,7 @@ class TestDirectMessagesModeStaysV1:
payload, _metadata = part
assert isinstance(payload, AIMessageChunk), (
"direct graph.stream(stream_mode='messages') leaked v2 "
"event dicts — StreamingHandler flag bled through."
"event dicts — GraphStreamer flag bled through."
)
assembled = "".join(
p[0].content for p in parts if isinstance(p[0].content, str)
libs/langgraph/tests/test_stream_subgraph_transformer.py
+9 -9
View File
@@ -13,7 +13,7 @@ from typing_extensions import TypedDict
from langgraph.constants import END, START
from langgraph.errors import GraphInterrupt
from langgraph.graph import StateGraph
from langgraph.stream import StreamingHandler
from langgraph.stream import GraphStreamer
from langgraph.stream._event_log import EventLog
from langgraph.stream._mux import StreamMux
from langgraph.stream._types import ProtocolEvent
@@ -245,7 +245,7 @@ class TestSubgraphTransformerUnit:
# ---------------------------------------------------------------------------
# End-to-end tests via StreamingHandler on real graphs
# End-to-end tests via GraphStreamer on real graphs
# ---------------------------------------------------------------------------
@@ -286,7 +286,7 @@ class TestSubgraphTransformerEndToEnd:
builder.add_edge("n", END)
graph = builder.compile()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "", "items": []})
collected: list[SubgraphRunStream] = []
@@ -298,7 +298,7 @@ class TestSubgraphTransformerEndToEnd:
def test_nested_graph_yields_one_child(self) -> None:
graph = _build_nested_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "", "items": []})
collected: list[SubgraphRunStream] = []
@@ -327,7 +327,7 @@ class TestSubgraphTransformerEndToEnd:
outer_builder.add_edge("sub", END)
graph = outer_builder.compile()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "", "items": []})
collected: list[SubgraphRunStream] = []
@@ -357,7 +357,7 @@ class TestSubgraphTransformerAsyncEndToEnd:
outer_builder.add_edge("sub", END)
graph = outer_builder.compile()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = await handler.astream({"value": "", "items": []})
collected: list[SubgraphRunStream] = []
@@ -374,7 +374,7 @@ class TestSubgraphTriggerCallId:
def test_trigger_call_id_populated_end_to_end(self) -> None:
graph = _build_nested_graph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "", "items": []})
collected: list[SubgraphRunStream] = list(run.subgraphs)
@@ -413,7 +413,7 @@ class TestSubgraphInterrupt:
def test_interrupt_in_subgraph_marks_handle_interrupted(self) -> None:
graph = self._build_interrupt_subgraph()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream(
{"value": "", "items": []},
config={"configurable": {"thread_id": "t1"}},
@@ -451,7 +451,7 @@ class TestSubgraphNameCollision:
outer_builder.add_edge("sub", END)
graph = outer_builder.compile()
handler = StreamingHandler(graph)
handler = GraphStreamer(graph)
run = handler.stream({"value": "", "items": []})
collected: list[SubgraphRunStream] = list(run.subgraphs)
libs/langgraph/tests/test_tool_stream_handler.py
+290
View File
@@ -0,0 +1,290 @@
"""Tests for StreamToolCallHandler and emit_tool_output_delta.
These tests exercise the langgraph-core piece in isolation — the prebuilt
`ToolCallTransformer` has its own test file. Here we feed real graphs
through `Pregel.stream(stream_mode=["tools", ...])` and inspect the raw
`(ns, mode, payload)` tuples on the `tools` channel.
"""
from __future__ import annotations
from typing import Annotated, Any
import pytest
from langchain_core.messages import AIMessage
from langchain_core.tools import tool
from langgraph.prebuilt import ToolNode
from typing_extensions import TypedDict
from langgraph.config import emit_tool_output_delta
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
class _State(TypedDict):
messages: Annotated[list, add_messages]
def _caller_sync(tool_name: str, tool_args: dict[str, Any], tc_id: str = "tc1"):
def caller(state: _State) -> dict:
return {
"messages": [
AIMessage(
content="",
tool_calls=[{"name": tool_name, "args": tool_args, "id": tc_id}],
)
]
}
return caller
def _caller_async(tool_name: str, tool_args: dict[str, Any], tc_id: str = "tc1"):
async def caller(state: _State) -> dict:
return {
"messages": [
AIMessage(
content="",
tool_calls=[{"name": tool_name, "args": tool_args, "id": tc_id}],
)
]
}
return caller
def _build_graph(caller, tools) -> Any:
sg = StateGraph(_State)
sg.add_node("caller", caller)
sg.add_node("tools", ToolNode(tools))
sg.add_edge(START, "caller")
sg.add_edge("caller", "tools")
sg.add_edge("tools", END)
return sg.compile()
def _tool_events(stream) -> list[tuple[tuple[str, ...], dict]]:
"""Collect `(ns, payload)` for every `tools`-mode chunk."""
out: list[tuple[tuple[str, ...], dict]] = []
for ns, mode, payload in stream:
if mode == "tools":
out.append((tuple(ns), payload))
return out
class TestSyncGraphSyncTool:
def test_started_finished_cycle(self) -> None:
@tool
def echo(text: str) -> str:
"""echo."""
return f"echoed:{text}"
graph = _build_graph(_caller_sync("echo", {"text": "hi"}), [echo])
events = _tool_events(
graph.stream(
{"messages": []},
stream_mode=["tools"],
subgraphs=True,
)
)
assert [p["event"] for _, p in events] == [
"tool-started",
"tool-finished",
]
assert events[0][1]["tool_call_id"] == "tc1"
assert events[0][1]["tool_name"] == "echo"
assert events[0][1]["input"] == {"text": "hi"}
# ToolNode wraps the return in a ToolMessage.
assert events[1][1]["tool_call_id"] == "tc1"
def test_emit_tool_output_delta_produces_delta_events(self) -> None:
@tool
def streaming_echo(text: str) -> str:
"""stream chunks."""
for chunk in ("a", "b", "c"):
emit_tool_output_delta(chunk)
return text
graph = _build_graph(
_caller_sync("streaming_echo", {"text": "x"}), [streaming_echo]
)
events = _tool_events(
graph.stream(
{"messages": []},
stream_mode=["tools"],
subgraphs=True,
)
)
deltas = [p["delta"] for _, p in events if p["event"] == "tool-output-delta"]
assert deltas == ["a", "b", "c"]
# The deltas must be bracketed by started and finished.
ordered = [p["event"] for _, p in events]
assert ordered[0] == "tool-started"
assert ordered[-1] == "tool-finished"
def test_tool_error_event(self) -> None:
@tool
def boom() -> str:
"""raises."""
raise ValueError("nope")
graph = _build_graph(_caller_sync("boom", {}), [boom])
events: list[tuple[tuple[str, ...], dict]] = []
with pytest.raises(ValueError, match="nope"):
for ns, mode, payload in graph.stream(
{"messages": []},
stream_mode=["tools"],
subgraphs=True,
):
if mode == "tools":
events.append((tuple(ns), payload))
kinds = [p["event"] for _, p in events]
assert kinds == ["tool-started", "tool-error"]
assert events[1][1]["message"] == "nope"
def test_emit_outside_tool_is_noop(self) -> None:
# Called at import time (outside any tool body) — must not raise.
emit_tool_output_delta("ignored")
emit_tool_output_delta({"any": "payload"})
def test_no_events_without_tools_mode(self) -> None:
@tool
def echo(text: str) -> str:
"""echo."""
return text
graph = _build_graph(_caller_sync("echo", {"text": "hi"}), [echo])
# No "tools" in stream_mode — handler is not attached and zero
# `tools`-method events fire.
chunks = list(
graph.stream(
{"messages": []},
stream_mode=["values"],
subgraphs=True,
)
)
assert all(
not (isinstance(c, tuple) and len(c) == 3 and c[1] == "tools")
for c in chunks
)
class TestAsyncGraphAsyncTool:
@pytest.mark.anyio
async def test_async_tool_produces_events(self) -> None:
@tool
async def aecho(text: str) -> str:
"""async echo."""
emit_tool_output_delta(text)
return f"got:{text}"
graph = _build_graph(_caller_async("aecho", {"text": "hi"}), [aecho])
events: list[tuple[tuple[str, ...], dict]] = []
async for ns, mode, payload in graph.astream(
{"messages": []},
stream_mode=["tools"],
subgraphs=True,
):
if mode == "tools":
events.append((tuple(ns), payload))
kinds = [p["event"] for _, p in events]
assert kinds == ["tool-started", "tool-output-delta", "tool-finished"]
assert events[1][1]["delta"] == "hi"
class TestConcurrentToolCalls:
def test_parallel_tool_calls_do_not_bleed(self) -> None:
@tool
def streamer(marker: str) -> str:
"""emits marker twice."""
emit_tool_output_delta(f"{marker}-1")
emit_tool_output_delta(f"{marker}-2")
return marker
def caller(state: _State) -> dict:
return {
"messages": [
AIMessage(
content="",
tool_calls=[
{"name": "streamer", "args": {"marker": "A"}, "id": "a"},
{"name": "streamer", "args": {"marker": "B"}, "id": "b"},
],
)
]
}
graph = _build_graph(caller, [streamer])
events = _tool_events(
graph.stream(
{"messages": []},
stream_mode=["tools"],
subgraphs=True,
)
)
# Group deltas by tool_call_id.
by_id: dict[str, list[str]] = {}
for _, p in events:
if p["event"] == "tool-output-delta":
by_id.setdefault(p["tool_call_id"], []).append(p["delta"])
assert by_id["a"] == ["A-1", "A-2"]
assert by_id["b"] == ["B-1", "B-2"]
class TestSubgraphNamespacePropagation:
def test_tool_inside_subgraph_emits_with_subgraph_ns(self) -> None:
@tool
def inner_tool(text: str) -> str:
"""inner tool."""
return text
def sub_caller(state: _State) -> dict:
return {
"messages": [
AIMessage(
content="",
tool_calls=[
{
"name": "inner_tool",
"args": {"text": "x"},
"id": "tc1",
}
],
)
]
}
inner = StateGraph(_State)
inner.add_node("sub_caller", sub_caller)
inner.add_node("sub_tools", ToolNode([inner_tool]))
inner.add_edge(START, "sub_caller")
inner.add_edge("sub_caller", "sub_tools")
inner.add_edge("sub_tools", END)
inner_graph = inner.compile()
outer = StateGraph(_State)
outer.add_node("sub", inner_graph)
outer.add_edge(START, "sub")
outer.add_edge("sub", END)
graph = outer.compile()
events = _tool_events(
graph.stream(
{"messages": []},
stream_mode=["tools"],
subgraphs=True,
)
)
# All `tools` events should carry a non-empty namespace rooted
# at the `sub` node.
assert events, "expected at least one tools event"
for ns, _ in events:
assert ns # non-empty
assert ns[0].startswith("sub:")
libs/prebuilt/langgraph/prebuilt/__init__.py
+4
View File
@@ -1,5 +1,7 @@
"""langgraph.prebuilt exposes a higher-level API for creating and executing agents and tools."""
from langgraph.prebuilt._tool_call_stream import ToolCallStream
from langgraph.prebuilt._tool_call_transformer import ToolCallTransformer
from langgraph.prebuilt.chat_agent_executor import create_react_agent
from langgraph.prebuilt.tool_node import (
InjectedState,
@@ -13,6 +15,8 @@ from langgraph.prebuilt.tool_validator import ValidationNode
__all__ = [
"create_react_agent",
"ToolNode",
"ToolCallStream",
"ToolCallTransformer",
"tools_condition",
"ValidationNode",
"InjectedState",
libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py
+117
View File
@@ -0,0 +1,117 @@
"""In-process handle for a single tool call's streaming execution.
Mirrors the shape of `ChatModelStream` from langchain-core but simpler —
a tool has one output channel, no content-block multiplexing. Populated
by `ToolCallTransformer` as `tool-started` / `tool-output-delta` /
`tool-finished` / `tool-error` events flow in on the `tools` channel.
"""
from __future__ import annotations
from collections.abc import AsyncIterator, Iterator
from typing import Any
from langgraph.stream._event_log import EventLog
class ToolCallStream:
"""Scoped view of a single tool call's lifecycle.
Yielded on `run.tool_calls` once per `tool-started` event. Fields
are populated as events arrive:
- `tool_call_id`, `tool_name`, `input`: stable from the start event.
- `output_deltas`: an `EventLog` of delta chunks. Iterate (sync or
async) to consume partial output in arrival order.
- `output`: terminal payload from `tool-finished`, or `None` if the
call failed or is still in flight.
- `error`: terminal error string from `tool-error`, or `None` if the
call succeeded or is still in flight.
- `completed`: True once a terminal event (`tool-finished` or
`tool-error`) has been observed.
`ToolCallStream` is not meant to be constructed by end users — it's
produced by `ToolCallTransformer` as events flow through the mux.
"""
def __init__(
self,
tool_call_id: str,
tool_name: str,
input: dict[str, Any] | None = None,
) -> None:
"""Initialize a fresh handle for a tool call.
Args:
tool_call_id: The `tool_call_id` from the AIMessage.
tool_name: The tool's name.
input: The tool's input arguments (as reported by
`on_tool_start`), or `None` if none were captured.
"""
self.tool_call_id = tool_call_id
self.tool_name = tool_name
self.input = input
self._output_deltas: EventLog[Any] = EventLog()
self.output: Any = None
self.error: str | None = None
self.completed = False
@property
def output_deltas(self) -> EventLog[Any]:
"""The EventLog of streamed `tool-output-delta` payloads.
Iterate (sync or async depending on how the run was started)
to consume partial output in arrival order. The log closes when
the tool finishes or errors.
"""
return self._output_deltas
def _bind(self, *, is_async: bool) -> None:
"""Bind the deltas log to sync or async iteration.
Called by `ToolCallTransformer` when constructing this handle so
the log matches the enclosing mux's mode.
"""
self._output_deltas._bind(is_async=is_async)
def _push_delta(self, delta: Any) -> None:
self._output_deltas.push(delta)
def _finish(self, output: Any) -> None:
self.output = output
self.completed = True
self._output_deltas.close()
def _fail(self, message: str) -> None:
self.error = message
self.completed = True
self._output_deltas.close()
def __iter__(self) -> Iterator[Any]:
"""Iterate delta chunks synchronously.
Equivalent to `iter(self.output_deltas)`. Raises `TypeError` if
the underlying log is bound to async mode.
"""
return iter(self._output_deltas)
def __aiter__(self) -> AsyncIterator[Any]:
"""Iterate delta chunks asynchronously.
Equivalent to `aiter(self.output_deltas)`. Raises `TypeError`
if the underlying log is bound to sync mode.
"""
return self._output_deltas.__aiter__()
def __repr__(self) -> str:
status = (
"completed"
if self.completed and self.error is None
else "failed"
if self.completed
else "running"
)
return (
f"ToolCallStream(tool_call_id={self.tool_call_id!r}, "
f"tool_name={self.tool_name!r}, status={status})"
)
libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py
+128
View File
@@ -0,0 +1,128 @@
"""Transformer that projects `tools` channel events into `ToolCallStream`s."""
from __future__ import annotations
from collections.abc import Awaitable, Callable
from typing import Any
from langgraph.stream._event_log import EventLog
from langgraph.stream._types import ProtocolEvent, StreamTransformer
from langgraph.prebuilt._tool_call_stream import ToolCallStream
class ToolCallTransformer(StreamTransformer):
"""Project `tools` channel events into `ToolCallStream` handles.
Each `tool-started` event spawns a `ToolCallStream`, pushed onto
`run.tool_calls`. Subsequent `tool-output-delta` events append to
that stream's deltas log; `tool-finished` and `tool-error` close it.
Native transformer — the `tool_calls` projection is exposed as a
direct attribute on the run stream.
`EventLog[ToolCallStream]` is used (not `StreamChannel`) because the
live handles are not serializable and should not be auto-forwarded
onto the main event log. Wire consumers subscribe to the `tools`
channel instead, where the raw protocol events flow through
untouched by this transformer (`process` returns `True`).
Registered explicitly by users via
`GraphStreamer(graph).stream(transformers=[ToolCallTransformer])`
— not a default built-in, so the `tools` channel is user-opt-in.
"""
_native = True
required_stream_modes = ("tools",)
def __init__(self, scope: tuple[str, ...] = ()) -> None:
super().__init__(scope)
self._log: EventLog[ToolCallStream] = EventLog()
self._active: dict[str, ToolCallStream] = {}
self._is_async = False
self._pump_fn: Callable[[], bool] | None = None
self._apump_fn: Callable[[], Awaitable[bool]] | None = None
def init(self) -> dict[str, Any]:
return {"tool_calls": self._log}
def _bind_pump(self, fn: Callable[[], bool]) -> None:
"""Wire the sync pull callback onto this transformer.
Called by `StreamMux.bind_pump`. Stored so each new
`ToolCallStream` created by `process` can wire its deltas log
for pump-driven iteration.
"""
self._pump_fn = fn
self._is_async = False
def _bind_apump(self, fn: Callable[[], Awaitable[bool]]) -> None:
"""Async counterpart to `_bind_pump`."""
self._apump_fn = fn
self._is_async = True
def _new_stream(
self,
tool_call_id: str,
tool_name: str,
tool_input: dict[str, Any] | None,
) -> ToolCallStream:
stream = ToolCallStream(tool_call_id, tool_name, tool_input)
stream._bind(is_async=self._is_async)
if self._apump_fn is not None:
stream._output_deltas._arequest_more = self._apump_fn
if self._pump_fn is not None:
stream._output_deltas._request_more = self._pump_fn
return stream
def process(self, event: ProtocolEvent) -> bool:
# Namespace filtering is handled by the mux via `scope_exact`.
if event["method"] != "tools":
return True
data = event["params"]["data"]
tool_call_id = data.get("tool_call_id")
if tool_call_id is None:
return True
event_type = data.get("event")
stream: ToolCallStream | None
if event_type == "tool-started":
stream = self._new_stream(
tool_call_id,
data.get("tool_name", ""),
data.get("input"),
)
self._active[tool_call_id] = stream
self._log.push(stream)
elif event_type == "tool-output-delta":
stream = self._active.get(tool_call_id)
if stream is not None:
stream._push_delta(data.get("delta"))
elif event_type == "tool-finished":
stream = self._active.pop(tool_call_id, None)
if stream is not None:
stream._finish(data.get("output"))
elif event_type == "tool-error":
stream = self._active.pop(tool_call_id, None)
if stream is not None:
stream._fail(data.get("message", ""))
# Pass-through — wire consumers subscribe to the `tools` channel
# directly and reconstruct handles client-side.
return True
def finalize(self) -> None:
"""Close any still-active tool streams left open at run end."""
for stream in self._active.values():
if not stream.completed:
stream._finish(None)
self._active.clear()
def fail(self, err: BaseException) -> None:
"""Fail any still-active tool streams when the run errors."""
message = str(err)
for stream in self._active.values():
if not stream.completed:
stream._fail(message)
self._active.clear()
libs/prebuilt/tests/test_tool_call_transformer.py
+311
View File
@@ -0,0 +1,311 @@
"""Tests for ToolCallTransformer and the ToolCallStream projection."""
from __future__ import annotations
import time
from typing import Annotated, Any
import pytest
from langchain_core.messages import AIMessage
from langchain_core.tools import tool
from langgraph.config import emit_tool_output_delta
from langgraph.constants import END, START
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langgraph.stream import GraphStreamer
from langgraph.stream._event_log import EventLog
from langgraph.stream._mux import StreamMux
from langgraph.stream._types import ProtocolEvent
from langgraph.stream.transformers import (
MessagesTransformer,
SubgraphTransformer,
ValuesTransformer,
)
from typing_extensions import TypedDict
from langgraph.prebuilt import ToolCallStream, ToolCallTransformer, ToolNode
TS = int(time.time() * 1000)
def _tool_event(
event: str,
tool_call_id: str,
*,
tool_name: str = "",
input: dict[str, Any] | None = None,
delta: Any = None,
output: Any = None,
message: str = "",
namespace: list[str] | None = None,
) -> ProtocolEvent:
data: dict[str, Any] = {"event": event, "tool_call_id": tool_call_id}
if event == "tool-started":
data["tool_name"] = tool_name
if input is not None:
data["input"] = input
elif event == "tool-output-delta":
data["delta"] = delta
elif event == "tool-finished":
data["output"] = output
elif event == "tool-error":
data["message"] = message
return {
"type": "event",
"method": "tools",
"params": {
"namespace": namespace or [],
"timestamp": TS,
"data": data,
},
}
def _subscribe(log: EventLog) -> None:
log._subscribed = True
def _mux() -> tuple[StreamMux, ToolCallTransformer]:
mux = StreamMux(
factories=[
ValuesTransformer,
MessagesTransformer,
SubgraphTransformer,
ToolCallTransformer,
],
is_async=False,
)
transformer = mux.transformer_by_key("tool_calls")
assert isinstance(transformer, ToolCallTransformer)
_subscribe(transformer._log)
return mux, transformer
class TestToolCallTransformerUnit:
def test_required_stream_modes_declares_tools(self) -> None:
assert ToolCallTransformer.required_stream_modes == ("tools",)
def test_tool_started_yields_handle(self) -> None:
mux, transformer = _mux()
mux.push(
_tool_event(
"tool-started",
"tc1",
tool_name="echo",
input={"text": "hi"},
)
)
handles = list(transformer._log._items)
assert len(handles) == 1
h = handles[0]
assert isinstance(h, ToolCallStream)
assert h.tool_call_id == "tc1"
assert h.tool_name == "echo"
assert h.input == {"text": "hi"}
assert h.completed is False
def test_delta_accumulates_on_active_stream(self) -> None:
mux, transformer = _mux()
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
_subscribe(transformer._active["tc1"]._output_deltas)
mux.push(_tool_event("tool-output-delta", "tc1", delta="a"))
mux.push(_tool_event("tool-output-delta", "tc1", delta="b"))
stream = transformer._active["tc1"]
assert list(stream._output_deltas._items) == ["a", "b"]
def test_finish_closes_stream(self) -> None:
mux, transformer = _mux()
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
stream = transformer._active["tc1"]
mux.push(_tool_event("tool-finished", "tc1", output="done"))
assert stream.completed is True
assert stream.output == "done"
assert stream.error is None
assert "tc1" not in transformer._active
def test_error_closes_stream(self) -> None:
mux, transformer = _mux()
mux.push(_tool_event("tool-started", "tc1", tool_name="boom"))
stream = transformer._active["tc1"]
mux.push(_tool_event("tool-error", "tc1", message="nope"))
assert stream.completed is True
assert stream.output is None
assert stream.error == "nope"
assert "tc1" not in transformer._active
def test_concurrent_tool_calls_do_not_bleed(self) -> None:
mux, transformer = _mux()
mux.push(_tool_event("tool-started", "a", tool_name="t"))
mux.push(_tool_event("tool-started", "b", tool_name="t"))
for tc in ("a", "b"):
_subscribe(transformer._active[tc]._output_deltas)
mux.push(_tool_event("tool-output-delta", "a", delta="A1"))
mux.push(_tool_event("tool-output-delta", "b", delta="B1"))
mux.push(_tool_event("tool-output-delta", "a", delta="A2"))
assert list(transformer._active["a"]._output_deltas._items) == ["A1", "A2"]
assert list(transformer._active["b"]._output_deltas._items) == ["B1"]
def test_tools_event_passes_through_main_log(self) -> None:
mux, transformer = _mux()
_subscribe(mux._events)
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
kept = [e for e in mux._events._items if e["method"] == "tools"]
assert len(kept) == 1
# ---------------------------------------------------------------------------
# End-to-end tests with a real graph
# ---------------------------------------------------------------------------
class _State(TypedDict):
messages: Annotated[list, add_messages]
def _build_graph(caller, tools):
sg = StateGraph(_State)
sg.add_node("caller", caller)
sg.add_node("tools", ToolNode(tools))
sg.add_edge(START, "caller")
sg.add_edge("caller", "tools")
sg.add_edge("tools", END)
return sg.compile()
class TestToolCallTransformerEndToEnd:
def test_sync_streaming_tool_populates_tool_calls(self) -> None:
@tool
def streamer(text: str) -> str:
"""streams chunks."""
for chunk in ("one", "two"):
emit_tool_output_delta(chunk)
return text
def caller(state: _State) -> dict:
return {
"messages": [
AIMessage(
content="",
tool_calls=[
{"name": "streamer", "args": {"text": "x"}, "id": "tc1"}
],
)
]
}
graph = _build_graph(caller, [streamer])
handler = GraphStreamer(graph)
run = handler.stream({"messages": []}, transformers=[ToolCallTransformer])
tool_calls: list[ToolCallStream] = []
for tc in run.tool_calls:
tool_calls.append(tc)
deltas = list(tc.output_deltas)
assert deltas == ["one", "two"]
assert len(tool_calls) == 1
tc = tool_calls[0]
assert tc.tool_call_id == "tc1"
assert tc.tool_name == "streamer"
assert tc.completed is True
assert tc.error is None
def test_stream_modes_union_includes_tools(self) -> None:
@tool
def echo(text: str) -> str:
"""echo."""
return text
def caller(state: _State) -> dict:
return {
"messages": [
AIMessage(
content="",
tool_calls=[
{"name": "echo", "args": {"text": "x"}, "id": "tc1"}
],
)
]
}
graph = _build_graph(caller, [echo])
handler = GraphStreamer(graph)
# Without ToolCallTransformer, no tool_calls projection is
# exposed and no `tools` events flow through (required_stream_modes
# omits it).
run_no_tc = handler.stream({"messages": []})
assert "tool_calls" not in run_no_tc._mux.extensions # type: ignore[attr-defined]
# With ToolCallTransformer, the projection is present.
run = handler.stream({"messages": []}, transformers=[ToolCallTransformer])
assert "tool_calls" in run._mux.extensions # type: ignore[attr-defined]
# Drain so the run closes cleanly.
list(run.tool_calls)
@pytest.mark.anyio
async def test_async_streaming_tool_populates_tool_calls(self) -> None:
@tool
async def astreamer(text: str) -> str:
"""async streams."""
emit_tool_output_delta(text)
emit_tool_output_delta(text + "!")
return text
async def caller(state: _State) -> dict:
return {
"messages": [
AIMessage(
content="",
tool_calls=[
{"name": "astreamer", "args": {"text": "hi"}, "id": "tc1"}
],
)
]
}
graph = _build_graph(caller, [astreamer])
handler = GraphStreamer(graph)
run = await handler.astream(
{"messages": []}, transformers=[ToolCallTransformer]
)
collected: list[ToolCallStream] = []
async for tc in run.tool_calls:
collected.append(tc)
deltas = [d async for d in tc.output_deltas]
assert deltas == ["hi", "hi!"]
assert len(collected) == 1
assert collected[0].completed is True
assert collected[0].error is None
def test_tool_error_populates_error_field(self) -> None:
@tool
def boom() -> str:
"""raises."""
raise ValueError("nope")
def caller(state: _State) -> dict:
return {
"messages": [
AIMessage(
content="",
tool_calls=[{"name": "boom", "args": {}, "id": "tc1"}],
)
]
}
graph = _build_graph(caller, [boom])
handler = GraphStreamer(graph)
run = handler.stream({"messages": []}, transformers=[ToolCallTransformer])
collected: list[ToolCallStream] = []
with pytest.raises(ValueError, match="nope"):
for tc in run.tool_calls:
collected.append(tc)
# Drain deltas so the error field is populated before we
# inspect it below.
list(tc.output_deltas)
assert len(collected) == 1
assert collected[0].error == "nope"
assert collected[0].output is None
assert collected[0].completed is True