From e01037d90ce39d8216b2b725699b3000b0f24020 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Mon, 20 Apr 2026 12:13:08 -0400 Subject: [PATCH] feat(core, prebuilt): tools-channel streaming and extensible GraphStreamer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- .../langgraph/_internal/_constants.py | 2 +- libs/langgraph/langgraph/config.py | 41 +++ libs/langgraph/langgraph/pregel/_messages.py | 2 +- libs/langgraph/langgraph/pregel/_tools.py | 223 +++++++++++++ libs/langgraph/langgraph/pregel/main.py | 13 + libs/langgraph/langgraph/stream/__init__.py | 6 +- libs/langgraph/langgraph/stream/_types.py | 7 + ...streaming_handler.py => graph_streamer.py} | 151 ++++++--- libs/langgraph/langgraph/stream/run_stream.py | 60 ++-- .../langgraph/stream/transformers.py | 3 + libs/langgraph/langgraph/types.py | 2 + ...ming_handler.py => test_graph_streamer.py} | 263 ++++++++++++--- .../tests/test_stream_messages_transformer.py | 42 +-- .../tests/test_stream_subgraph_transformer.py | 18 +- .../tests/test_tool_stream_handler.py | 290 ++++++++++++++++ libs/prebuilt/langgraph/prebuilt/__init__.py | 4 + .../langgraph/prebuilt/_tool_call_stream.py | 117 +++++++ .../prebuilt/_tool_call_transformer.py | 128 +++++++ .../tests/test_tool_call_transformer.py | 311 ++++++++++++++++++ 19 files changed, 1529 insertions(+), 154 deletions(-) create mode 100644 libs/langgraph/langgraph/pregel/_tools.py rename libs/langgraph/langgraph/stream/{streaming_handler.py => graph_streamer.py} (53%) rename libs/langgraph/tests/{test_streaming_handler.py => test_graph_streamer.py} (89%) create mode 100644 libs/langgraph/tests/test_tool_stream_handler.py create mode 100644 libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py create mode 100644 libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py create mode 100644 libs/prebuilt/tests/test_tool_call_transformer.py diff --git a/libs/langgraph/langgraph/_internal/_constants.py b/libs/langgraph/langgraph/_internal/_constants.py index d28289053..a6b8a4f2a 100644 --- a/libs/langgraph/langgraph/_internal/_constants.py +++ b/libs/langgraph/langgraph/_internal/_constants.py @@ -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") diff --git a/libs/langgraph/langgraph/config.py b/libs/langgraph/langgraph/config.py index c8a321d83..112a7e1ad 100644 --- a/libs/langgraph/langgraph/config.py +++ b/libs/langgraph/langgraph/config.py @@ -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) diff --git a/libs/langgraph/langgraph/pregel/_messages.py b/libs/langgraph/langgraph/pregel/_messages.py index 1c196add6..3416f9854 100644 --- a/libs/langgraph/langgraph/pregel/_messages.py +++ b/libs/langgraph/langgraph/pregel/_messages.py @@ -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. diff --git a/libs/langgraph/langgraph/pregel/_tools.py b/libs/langgraph/langgraph/pregel/_tools.py new file mode 100644 index 000000000..879e36c9b --- /dev/null +++ b/libs/langgraph/langgraph/pregel/_tools.py @@ -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) diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 6d5b7cc48..7f3690702 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -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( diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py index 5a4633c71..5287fda3d 100644 --- a/libs/langgraph/langgraph/stream/__init__.py +++ b/libs/langgraph/langgraph/stream/__init__.py @@ -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", ] diff --git a/libs/langgraph/langgraph/stream/_types.py b/libs/langgraph/langgraph/stream/_types.py index 797be7e76..1a6c5ce22 100644 --- a/libs/langgraph/langgraph/stream/_types.py +++ b/libs/langgraph/langgraph/stream/_types.py @@ -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. diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/graph_streamer.py similarity index 53% rename from libs/langgraph/langgraph/stream/streaming_handler.py rename to libs/langgraph/langgraph/stream/graph_streamer.py index bbc764dfe..752d71901 100644 --- a/libs/langgraph/langgraph/stream/streaming_handler.py +++ b/libs/langgraph/langgraph/stream/graph_streamer.py @@ -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) diff --git a/libs/langgraph/langgraph/stream/run_stream.py b/libs/langgraph/langgraph/stream/run_stream.py index f12c642f9..a8b06a362 100644 --- a/libs/langgraph/langgraph/stream/run_stream.py +++ b/libs/langgraph/langgraph/stream/run_stream.py @@ -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) diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index 8bc53ab1f..16e994cf4 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -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) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index 73954e5eb..a4021271e 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -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] diff --git a/libs/langgraph/tests/test_streaming_handler.py b/libs/langgraph/tests/test_graph_streamer.py similarity index 89% rename from libs/langgraph/tests/test_streaming_handler.py rename to libs/langgraph/tests/test_graph_streamer.py index a70181589..21ae77d4c 100644 --- a/libs/langgraph/tests/test_streaming_handler.py +++ b/libs/langgraph/tests/test_graph_streamer.py @@ -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 diff --git a/libs/langgraph/tests/test_stream_messages_transformer.py b/libs/langgraph/tests/test_stream_messages_transformer.py index 27492ac40..7cf7e353a 100644 --- a/libs/langgraph/tests/test_stream_messages_transformer.py +++ b/libs/langgraph/tests/test_stream_messages_transformer.py @@ -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) diff --git a/libs/langgraph/tests/test_stream_subgraph_transformer.py b/libs/langgraph/tests/test_stream_subgraph_transformer.py index b4d7dadcb..92a9d384b 100644 --- a/libs/langgraph/tests/test_stream_subgraph_transformer.py +++ b/libs/langgraph/tests/test_stream_subgraph_transformer.py @@ -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) diff --git a/libs/langgraph/tests/test_tool_stream_handler.py b/libs/langgraph/tests/test_tool_stream_handler.py new file mode 100644 index 000000000..2573c70e9 --- /dev/null +++ b/libs/langgraph/tests/test_tool_stream_handler.py @@ -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:") diff --git a/libs/prebuilt/langgraph/prebuilt/__init__.py b/libs/prebuilt/langgraph/prebuilt/__init__.py index a93cc0219..13052a66d 100644 --- a/libs/prebuilt/langgraph/prebuilt/__init__.py +++ b/libs/prebuilt/langgraph/prebuilt/__init__.py @@ -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", diff --git a/libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py b/libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py new file mode 100644 index 000000000..2df12e823 --- /dev/null +++ b/libs/prebuilt/langgraph/prebuilt/_tool_call_stream.py @@ -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})" + ) diff --git a/libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py b/libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py new file mode 100644 index 000000000..9db87fd2e --- /dev/null +++ b/libs/prebuilt/langgraph/prebuilt/_tool_call_transformer.py @@ -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() diff --git a/libs/prebuilt/tests/test_tool_call_transformer.py b/libs/prebuilt/tests/test_tool_call_transformer.py new file mode 100644 index 000000000..ac22417fd --- /dev/null +++ b/libs/prebuilt/tests/test_tool_call_transformer.py @@ -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