From 40055e92cc8e0123169c885f1c0c0701c0eb4561 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Mon, 20 Apr 2026 16:08:19 -0400 Subject: [PATCH] feat(langgraph): move stream_v2/astream_v2 onto Pregel, drop StreamingHandler Compile-time `transformers=` on `StateGraph.compile` now registers transformer factories directly on the compiled graph. `stream_v2` and `astream_v2` live on Pregel and read the stashed list, so callers no longer need a separate wrapper to drive the transformer pipeline. --- libs/langgraph/langgraph/graph/state.py | 7 + libs/langgraph/langgraph/pregel/main.py | 132 +++++++++++++ libs/langgraph/langgraph/stream/__init__.py | 7 +- .../langgraph/stream/streaming_handler.py | 149 -------------- ...ng_handler.py => test_pregel_stream_v2.py} | 185 +++++++++--------- 5 files changed, 234 insertions(+), 246 deletions(-) delete mode 100644 libs/langgraph/langgraph/stream/streaming_handler.py rename libs/langgraph/tests/{test_streaming_handler.py => test_pregel_stream_v2.py} (92%) diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index b1c24de2b..4e3bde4c7 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -1045,6 +1045,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): interrupt_after: All | list[str] | None = None, debug: bool = False, name: str | None = None, + transformers: Sequence[Callable[[], Any]] | None = None, ) -> CompiledStateGraph[StateT, ContextT, InputT, OutputT]: """Compiles the `StateGraph` into a `CompiledStateGraph` object. @@ -1077,6 +1078,11 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): interrupt_after: An optional list of node names to interrupt after. debug: A flag indicating whether to enable debug mode. name: The name to use for the compiled graph. + transformers: Optional sequence of zero-arg factories returning + `StreamTransformer` instances. Registered on the compiled + graph and instantiated per-run whenever `stream_v2` / + `astream_v2` is called. Appended after the built-in + `ValuesTransformer` and `MessagesTransformer`. Returns: CompiledStateGraph: The compiled `StateGraph`. @@ -1159,6 +1165,7 @@ class StateGraph(Generic[StateT, ContextT, InputT, OutputT]): store=store, cache=cache, name=name or "LangGraph", + stream_transformers=transformers, ) compiled._serde_allowlist = serde_allowlist diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index c440e77f9..ce38ae1bf 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -671,6 +671,7 @@ class Pregel( config: RunnableConfig | None = None, trigger_to_nodes: Mapping[str, Sequence[str]] | None = None, name: str = "LangGraph", + stream_transformers: Sequence[Callable[[], Any]] | None = None, **deprecated_kwargs: Unpack[DeprecatedKwargs], ) -> None: if ( @@ -717,6 +718,9 @@ class Pregel( self.config = config self.trigger_to_nodes = trigger_to_nodes or {} self.name = name + self._stream_transformers: tuple[Callable[[], Any], ...] = tuple( + stream_transformers or () + ) self._serde_allowlist: set[tuple[str, ...]] | None = None if auto_validate: self.validate() @@ -3237,6 +3241,134 @@ class Pregel( await asyncio.shield(run_manager.on_chain_error(e)) raise + def stream_v2( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + transformers: Sequence[Any] | None = None, + ) -> Any: + """Start a sync v2 streaming run driven by transformer projections. + + Builds a `StreamMux` from the built-in `ValuesTransformer` / + `MessagesTransformer`, this graph's compile-time + `stream_transformers`, and any additional `transformers=` + supplied at the call site. Returns a `GraphRunStream` that the + caller drives by iterating any projection — no background + thread. + + Args: + input: Graph input. + config: Optional runnable config forwarded to the graph. + interrupt_before: Nodes to interrupt before, if any. + interrupt_after: Nodes to interrupt after, if any. + transformers: Extra transformer instances appended after + compile-time `stream_transformers`. + + Returns: + A `GraphRunStream` the caller iterates to drive the run. + """ + from langgraph.stream._mux import StreamMux + from langgraph.stream.run_stream import GraphRunStream + from langgraph.stream.transformers import ( + MessagesTransformer, + ValuesTransformer, + ) + + values_t = ValuesTransformer() + compiled_instances = [f() for f in self._stream_transformers] + mux = StreamMux( + [ + values_t, + MessagesTransformer(), + *compiled_instances, + *(transformers or ()), + ], + is_async=False, + ) + graph_iter = iter( + self.stream( + input, + config, + stream_mode=[ + "values", + "updates", + "messages", + "custom", + "checkpoints", + "tasks", + "debug", + ], + subgraphs=True, + version="v2", + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + ) + ) + return GraphRunStream(graph_iter, mux, values_t) + + async def astream_v2( + self, + input: InputT | Command | None, + config: RunnableConfig | None = None, + *, + interrupt_before: All | Sequence[str] | None = None, + interrupt_after: All | Sequence[str] | None = None, + transformers: Sequence[Any] | None = None, + ) -> Any: + """Async counterpart to `stream_v2`. + + Returns an `AsyncGraphRunStream` whose projections can be awaited + concurrently; each subscribed cursor drives the pump when its + buffer is empty. + + Args: + input: Graph input. + config: Optional runnable config forwarded to the graph. + interrupt_before: Nodes to interrupt before, if any. + interrupt_after: Nodes to interrupt after, if any. + transformers: Extra transformer instances appended after + compile-time `stream_transformers`. + """ + from langgraph.stream._mux import StreamMux + from langgraph.stream.run_stream import AsyncGraphRunStream + from langgraph.stream.transformers import ( + MessagesTransformer, + ValuesTransformer, + ) + + values_t = ValuesTransformer() + compiled_instances = [f() for f in self._stream_transformers] + mux = StreamMux( + [ + values_t, + MessagesTransformer(), + *compiled_instances, + *(transformers or ()), + ], + is_async=True, + ) + graph_aiter = self.astream( + input, + config, + stream_mode=[ + "values", + "updates", + "messages", + "custom", + "checkpoints", + "tasks", + "debug", + ], + subgraphs=True, + version="v2", + interrupt_before=interrupt_before, + interrupt_after=interrupt_after, + ).__aiter__() + return AsyncGraphRunStream(graph_aiter, mux, values_t) + @overload def invoke( self, diff --git a/libs/langgraph/langgraph/stream/__init__.py b/libs/langgraph/langgraph/stream/__init__.py index 5a4633c71..ca5a26010 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 -ergonomic streaming projections through a transformer pipeline. +Compile a graph with `transformers=[...]` and call `graph.stream_v2()` / +`graph.astream_v2()` to drive a transformer pipeline that projects the +graph's raw events into ergonomic per-channel streams. """ from langgraph.stream._event_log import EventLog from langgraph.stream._types import ProtocolEvent, StreamTransformer from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream from langgraph.stream.stream_channel import StreamChannel -from langgraph.stream.streaming_handler import StreamingHandler __all__ = [ "AsyncGraphRunStream", @@ -17,5 +17,4 @@ __all__ = [ "ProtocolEvent", "StreamChannel", "StreamTransformer", - "StreamingHandler", ] diff --git a/libs/langgraph/langgraph/stream/streaming_handler.py b/libs/langgraph/langgraph/stream/streaming_handler.py deleted file mode 100644 index ff031016e..000000000 --- a/libs/langgraph/langgraph/stream/streaming_handler.py +++ /dev/null @@ -1,149 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence -from typing import Any - -from langchain_core.runnables import RunnableConfig - -from langgraph.pregel import Pregel -from langgraph.stream._mux import StreamMux -from langgraph.stream._types import StreamTransformer -from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream -from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer -from langgraph.types import All, StreamMode - -# All stream modes to request from the graph. -STREAM_V2_MODES: list[StreamMode] = [ - "values", - "updates", - "messages", - "custom", - "checkpoints", - "tasks", - "debug", -] - - -class StreamingHandler: - """Wrap a compiled graph with ergonomic streaming projections. - - Example: - ```python - handler = StreamingHandler(graph) - - # Sync - run = handler.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) - async for state in run.values: - print(state) - output = await run.output() - ``` - """ - - def __init__(self, graph: Pregel) -> None: - """Initialize the handler. - - Args: - graph: A compiled LangGraph graph to stream from. - """ - self._graph = graph - - def stream( - self, - input: Any, - config: RunnableConfig | None = None, - *, - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - transformers: list[StreamTransformer] | None = None, - ) -> GraphRunStream: - """Start a sync streaming run. - - Returns a GraphRunStream immediately. The caller's iteration on - any projection drives the graph forward — no background thread - is used. This matches v1's model where the caller's `for` loop - is the pump. - - Args: - input: Graph input. - config: Optional runnable config forwarded to the graph. - interrupt_before: Nodes to interrupt before, if any. - interrupt_after: Nodes to interrupt after, if any. - transformers: User transformers appended after the built-in - `ValuesTransformer` and `MessagesTransformer`. - - Returns: - A GraphRunStream the caller can iterate to drive the run. - """ - values_t = ValuesTransformer() - mux = StreamMux( - [values_t, MessagesTransformer(), *(transformers or ())], - is_async=False, - ) - - graph_iter = iter( - self._graph.stream( - input, - config, - stream_mode=STREAM_V2_MODES, - subgraphs=True, - version="v2", - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - ) - ) - - return GraphRunStream(graph_iter, mux, values_t) - - async def astream( - self, - input: Any, - config: RunnableConfig | None = None, - *, - interrupt_before: All | Sequence[str] | None = None, - interrupt_after: All | Sequence[str] | None = None, - transformers: list[StreamTransformer] | None = None, - ) -> AsyncGraphRunStream: - """Start an async streaming run. - - Returns an AsyncGraphRunStream immediately. The caller's - iteration on any projection drives the graph forward — there - is no background task. Concurrent consumers share a - single-flight pump via an internal `asyncio.Lock`. - - Args: - input: Graph input. - config: Optional runnable config forwarded to the graph. - interrupt_before: Nodes to interrupt before, if any. - interrupt_after: Nodes to interrupt after, if any. - transformers: User transformers appended after the built-in - `ValuesTransformer` and `MessagesTransformer`. - - Returns: - An AsyncGraphRunStream whose projections can be awaited - concurrently; each subscribed cursor drives the pump when - its buffer is empty. - """ - values_t = ValuesTransformer() - mux = StreamMux( - [values_t, MessagesTransformer(), *(transformers or ())], - is_async=True, - ) - - graph_aiter = self._graph.astream( - input, - config, - stream_mode=STREAM_V2_MODES, - subgraphs=True, - version="v2", - interrupt_before=interrupt_before, - interrupt_after=interrupt_after, - ).__aiter__() - - return AsyncGraphRunStream(graph_aiter, mux, values_t) diff --git a/libs/langgraph/tests/test_streaming_handler.py b/libs/langgraph/tests/test_pregel_stream_v2.py similarity index 92% rename from libs/langgraph/tests/test_streaming_handler.py rename to libs/langgraph/tests/test_pregel_stream_v2.py index 835158ed5..e0d47126a 100644 --- a/libs/langgraph/tests/test_streaming_handler.py +++ b/libs/langgraph/tests/test_pregel_stream_v2.py @@ -1,4 +1,4 @@ -"""Tests for the StreamingHandler and its supporting infrastructure.""" +"""Tests for `Pregel.stream_v2` / `astream_v2` and the transformer pipeline.""" from __future__ import annotations @@ -17,7 +17,6 @@ from langgraph.graph import StateGraph from langgraph.stream import ( EventLog, StreamChannel, - StreamingHandler, StreamTransformer, ) from langgraph.stream._convert import convert_to_protocol_event @@ -404,15 +403,15 @@ class TestStreamChannel: # --------------------------------------------------------------------------- -# StreamingHandler sync tests +# stream_v2 sync tests # --------------------------------------------------------------------------- -class TestStreamingHandlerSync: +class TestStreamV2Sync: def test_values_projection(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) snapshots = list(run.values) # Should have at least the initial + per-node snapshots. assert len(snapshots) >= 1 @@ -423,8 +422,8 @@ class TestStreamingHandlerSync: def test_output(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) output = run.output assert output is not None assert output["value"] == "xAB" @@ -432,8 +431,8 @@ class TestStreamingHandlerSync: def test_raw_event_iteration(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) events = list(run) assert len(events) > 0 for event in events: @@ -444,8 +443,8 @@ class TestStreamingHandlerSync: def test_extensions_has_native_keys(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) # Drain events so the run completes. _ = run.output assert "values" in run.extensions @@ -457,8 +456,8 @@ 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) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) with pytest.raises(TypeError): run.extensions["new_key"] = object() # type: ignore[index] with pytest.raises(TypeError): @@ -466,8 +465,8 @@ class TestStreamingHandlerSync: def test_custom_stream_events(self) -> None: graph = _build_custom_stream_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) custom_events = [e for e in run if e["method"] == "custom"] assert len(custom_events) == 2 assert custom_events[0]["params"]["data"] == {"step": "start"} @@ -475,8 +474,8 @@ class TestStreamingHandlerSync: def test_interleave_values_and_messages(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) tagged = list(run.interleave("values", "messages")) names = [name for name, _ in tagged] @@ -489,8 +488,8 @@ class TestStreamingHandlerSync: def test_abort_marks_exhausted_and_closes_mux(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) values_iter = iter(run.values) # Consume one item so the pump advances. _ = next(values_iter) @@ -503,64 +502,64 @@ class TestStreamingHandlerSync: def test_context_manager_calls_abort_on_exit(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - with handler.stream({"value": "x", "items": []}) as run: + handler = graph + with handler.stream_v2({"value": "x", "items": []}) as run: values_iter = iter(run.values) _ = next(values_iter) assert run._exhausted is True def test_interleave_unknown_projection(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) with pytest.raises(KeyError): list(run.interleave("values", "does_not_exist")) -class TestStreamingHandlerSyncErrors: +class TestStreamV2SyncErrors: def test_error_propagation_output(self) -> None: graph = _build_error_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"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) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"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) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): list(run) 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) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): _ = run.interrupted def test_error_propagation_interrupts(self) -> None: """`run.interrupts` should raise on a failed run.""" graph = _build_error_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): _ = run.interrupts -class TestStreamingHandlerSyncInterrupt: +class TestStreamV2SyncInterrupt: def test_interrupted(self) -> None: graph = _build_interrupt_graph() - handler = StreamingHandler(graph) - run = handler.stream( + handler = graph + run = handler.stream_v2( {"value": "x", "items": []}, {"configurable": {"thread_id": "t1"}}, ) @@ -570,17 +569,17 @@ class TestStreamingHandlerSyncInterrupt: # --------------------------------------------------------------------------- -# StreamingHandler async tests +# astream_v2 async tests # --------------------------------------------------------------------------- -class TestStreamingHandlerAsync: +class TestStreamV2Async: @pytest.mark.anyio @NEEDS_CONTEXTVARS async def test_values_projection(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) snapshots = [s async for s in run.values] assert len(snapshots) >= 1 last = snapshots[-1] @@ -591,8 +590,8 @@ class TestStreamingHandlerAsync: @NEEDS_CONTEXTVARS async def test_output(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) output = await run.output() assert output is not None assert output["value"] == "xAB" @@ -602,8 +601,8 @@ class TestStreamingHandlerAsync: @NEEDS_CONTEXTVARS async def test_raw_event_iteration(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) events = [e async for e in run] assert len(events) > 0 for event in events: @@ -613,8 +612,8 @@ class TestStreamingHandlerAsync: @NEEDS_CONTEXTVARS async def test_abort_marks_exhausted_and_closes_mux(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) values_iter = aiter(run.values) _ = await anext(values_iter) await run.abort() @@ -628,8 +627,8 @@ class TestStreamingHandlerAsync: @NEEDS_CONTEXTVARS async def test_async_context_manager_calls_abort_on_exit(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) async with run: values_iter = aiter(run.values) _ = await anext(values_iter) @@ -639,8 +638,8 @@ class TestStreamingHandlerAsync: @NEEDS_CONTEXTVARS async def test_extensions_has_native_keys(self) -> None: graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) _ = await run.output() assert "values" in run.extensions assert "messages" in run.extensions @@ -648,13 +647,13 @@ class TestStreamingHandlerAsync: assert run.messages is run.extensions["messages"] -class TestStreamingHandlerAsyncErrors: +class TestStreamV2AsyncErrors: @pytest.mark.anyio @NEEDS_CONTEXTVARS async def test_error_propagation_output(self) -> None: graph = _build_error_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): await run.output() @@ -662,8 +661,8 @@ class TestStreamingHandlerAsyncErrors: @NEEDS_CONTEXTVARS async def test_error_propagation_values(self) -> None: graph = _build_error_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): async for _ in run.values: pass @@ -672,8 +671,8 @@ class TestStreamingHandlerAsyncErrors: @NEEDS_CONTEXTVARS async def test_error_propagation_raw_events(self) -> None: graph = _build_error_graph() - handler = StreamingHandler(graph) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): async for _ in run: pass @@ -683,8 +682,8 @@ 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) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): await run.interrupted() @@ -693,19 +692,19 @@ 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) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) with pytest.raises(ValueError, match="boom"): await run.interrupts() -class TestStreamingHandlerAsyncInterrupt: +class TestStreamV2AsyncInterrupt: @pytest.mark.anyio @NEEDS_CONTEXTVARS async def test_interrupted(self) -> None: graph = _build_interrupt_graph() - handler = StreamingHandler(graph) - run = await handler.astream( + handler = graph + run = await handler.astream_v2( {"value": "x", "items": []}, {"configurable": {"thread_id": "t2"}}, ) @@ -714,13 +713,13 @@ class TestStreamingHandlerAsyncInterrupt: assert len(await run.interrupts()) > 0 -class TestStreamingHandlerAsyncCustom: +class TestStreamV2AsyncCustom: @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 = graph + run = await handler.astream_v2({"value": "x", "items": []}) events = [e async for e in run] custom_events = [e for e in events if e["method"] == "custom"] assert len(custom_events) == 2 @@ -1082,9 +1081,9 @@ class TestCustomTransformer: return True graph = _build_simple_graph() - handler = StreamingHandler(graph) + handler = graph counter_t = CounterTransformer() - run = handler.stream({"value": "x", "items": []}, transformers=[counter_t]) + run = handler.stream_v2({"value": "x", "items": []}, transformers=[counter_t]) assert "counter" in run.extensions # Subscribe before driving the run so channel pushes are retained. counter_iter = iter(run.extensions["counter"]) @@ -1113,9 +1112,9 @@ class TestCustomTransformer: return True graph = _build_simple_graph() - handler = StreamingHandler(graph) + handler = graph foo_t = FooTransformer() - run = handler.stream({"value": "x", "items": []}, transformers=[foo_t]) + run = handler.stream_v2({"value": "x", "items": []}, transformers=[foo_t]) # Subscribe before driving the run. foo_iter = iter(run.foo) _ = run.output @@ -1143,8 +1142,8 @@ class TestCustomTransformer: return True graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream( + handler = graph + run = handler.stream_v2( {"value": "x", "items": []}, transformers=[EmitterTransformer()] ) events = list(run) @@ -1202,12 +1201,12 @@ class TestCustomTransformer: return True graph = _build_simple_graph() - handler = StreamingHandler(graph) + handler = graph with pytest.raises( ValueError, match=r"conflict.*'values'.*ValuesTransformer", ): - handler.stream( + handler.stream_v2( {"value": "x", "items": []}, transformers=[ConflictTransformer()], ) @@ -1304,9 +1303,9 @@ class TestEventLogAutoLifecycle: return True graph = _build_simple_graph() - handler = StreamingHandler(graph) + handler = graph t = MinimalTransformer() - run = handler.stream({"value": "x", "items": []}, transformers=[t]) + run = handler.stream_v2({"value": "x", "items": []}, transformers=[t]) minimal_iter = iter(run.extensions["minimal"]) _ = run.output items = list(minimal_iter) @@ -1605,8 +1604,8 @@ class TestAsyncTransformerLane: self._log.close() graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = await handler.astream( + handler = graph + run = await handler.astream_v2( {"value": "x", "items": []}, transformers=[Scorer()], ) @@ -1631,8 +1630,8 @@ 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) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) events_iter = iter(run) max_buffered = 0 count = 0 @@ -1649,8 +1648,8 @@ class TestMemoryBounds: """Projections without a subscriber drop pushes silently — their buffers stay empty regardless of run length.""" graph = _build_simple_graph() - handler = StreamingHandler(graph) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) # Subscribe to main events only; leave values and messages unsubscribed. list(run) values_log = run.extensions["values"] @@ -1665,8 +1664,8 @@ 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) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) _ = run.output values_log = run.extensions["values"] assert len(values_log._items) == 0 @@ -1675,8 +1674,8 @@ 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) - run = handler.stream({"value": "x", "items": []}) + handler = graph + run = handler.stream_v2({"value": "x", "items": []}) values_log = run.extensions["values"] list(run.values) assert len(values_log._items) == 0 @@ -1686,8 +1685,8 @@ 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) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) max_buffered = 0 count = 0 async for _ in run: @@ -1701,8 +1700,8 @@ 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) - run = await handler.astream({"value": "x", "items": []}) + handler = graph + run = await handler.astream_v2({"value": "x", "items": []}) _ = await run.output() values_log = run.extensions["values"] messages_log = run.extensions["messages"]