From 5e45d85a3abe1f74bf243f9667fe8a63747c52f2 Mon Sep 17 00:00:00 2001 From: Nick Hollon Date: Mon, 20 Apr 2026 16:39:29 -0400 Subject: [PATCH] feat(langchain): register ToolCallTransformer on compile, drop AgentStreamer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `create_agent` now passes `transformers=[ToolCallTransformer]` straight through to `graph.compile`, so callers drive the transformer pipeline with `agent.stream_v2()` / `agent.astream_v2()` on the compiled graph. The `AgentStreamer` / `AgentRunStream` / `AsyncAgentRunStream` wrapper classes are gone — projection attributes still bind on `GraphRunStream` by key when the matching transformer is registered. --- .../langchain_v1/langchain/agents/__init__.py | 8 -- .../langchain/agents/_streaming.py | 79 ------------------- libs/langchain_v1/langchain/agents/factory.py | 2 + .../unit_tests/agents/test_agent_streaming.py | 49 ++++-------- 4 files changed, 17 insertions(+), 121 deletions(-) delete mode 100644 libs/langchain_v1/langchain/agents/_streaming.py diff --git a/libs/langchain_v1/langchain/agents/__init__.py b/libs/langchain_v1/langchain/agents/__init__.py index fbce646b131..5e5c545fe21 100644 --- a/libs/langchain_v1/langchain/agents/__init__.py +++ b/libs/langchain_v1/langchain/agents/__init__.py @@ -1,17 +1,9 @@ """Entrypoint to building [Agents](https://docs.langchain.com/oss/python/langchain/agents) with LangChain.""" # noqa: E501 -from langchain.agents._streaming import ( - AgentRunStream, - AgentStreamer, - AsyncAgentRunStream, -) from langchain.agents.factory import create_agent from langchain.agents.middleware.types import AgentState __all__ = [ - "AgentRunStream", "AgentState", - "AgentStreamer", - "AsyncAgentRunStream", "create_agent", ] diff --git a/libs/langchain_v1/langchain/agents/_streaming.py b/libs/langchain_v1/langchain/agents/_streaming.py deleted file mode 100644 index 247c5100ab8..00000000000 --- a/libs/langchain_v1/langchain/agents/_streaming.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Streaming entry point for `create_agent` graphs. - -`AgentStreamer` pre-registers `ToolCallTransformer` so every agent run -exposes `run.tool_calls` without the caller opting in. - -Example: - ```python - from langchain.agents import AgentStreamer, create_agent - - agent = create_agent(model, tools) - - run = AgentStreamer(agent).stream({"messages": [...]}) - for tc in run.tool_calls: - for delta in tc.output_deltas: - print(delta, end="") - print(run.output) - ``` -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, ClassVar - -from langgraph.prebuilt import ToolCallTransformer -from langgraph.stream import GraphStreamer -from langgraph.stream.run_stream import AsyncGraphRunStream, GraphRunStream - -if TYPE_CHECKING: - from collections.abc import AsyncIterator, Iterator - - from langgraph.stream._mux import StreamMux, TransformerFactory - - -class AgentRunStream(GraphRunStream): - """Sync run stream for a `create_agent` graph. - - Native projections (`tool_calls`, `messages`, `values`) are bound as - instance attributes by `BaseRunStream.__init__` whenever the - matching transformer is registered — this subclass exists for - `isinstance` checks and as an extension point for downstream - streamers (e.g. a deepagents-layer `DeepAgentRunStream`). - """ - - -class AsyncAgentRunStream(AsyncGraphRunStream): - """Async counterpart to `AgentRunStream`.""" - - -class AgentStreamer(GraphStreamer): - """`GraphStreamer` pre-configured for `create_agent` graphs. - - Extends `GraphStreamer.builtin_factories` with `ToolCallTransformer` - so `run.tool_calls` is populated on every run without the caller - passing `transformers=[ToolCallTransformer]`. Returns - `AgentRunStream` / `AsyncAgentRunStream` for `isinstance` checks. - - Caller-supplied `transformers=[...]` on `stream()` / `astream()` - are appended after the built-ins, matching `GraphStreamer`'s - behavior — they add to, rather than replace, the agent defaults. - """ - - builtin_factories: ClassVar[tuple[TransformerFactory, ...]] = ( - *GraphStreamer.builtin_factories, - ToolCallTransformer, - ) - - def _make_run_stream( - self, - graph_iter: Iterator[Any], - mux: StreamMux, - ) -> AgentRunStream: - return AgentRunStream(graph_iter, mux) - - def _make_async_run_stream( - self, - graph_aiter: AsyncIterator[Any], - mux: StreamMux, - ) -> AsyncAgentRunStream: - return AsyncAgentRunStream(graph_aiter, mux) diff --git a/libs/langchain_v1/langchain/agents/factory.py b/libs/langchain_v1/langchain/agents/factory.py index a72ea8d5492..7a7a6b0a3c8 100644 --- a/libs/langchain_v1/langchain/agents/factory.py +++ b/libs/langchain_v1/langchain/agents/factory.py @@ -22,6 +22,7 @@ from langchain_core.tools import BaseTool from langgraph._internal._runnable import RunnableCallable from langgraph.constants import END, START from langgraph.graph.state import StateGraph +from langgraph.prebuilt import ToolCallTransformer from langgraph.prebuilt.tool_node import ToolCallWithContext, ToolNode from langgraph.types import Command, Send from langsmith import traceable @@ -1660,6 +1661,7 @@ def create_agent( debug=debug, name=name, cache=cache, + transformers=[ToolCallTransformer], ).with_config(config) diff --git a/libs/langchain_v1/tests/unit_tests/agents/test_agent_streaming.py b/libs/langchain_v1/tests/unit_tests/agents/test_agent_streaming.py index 45cd3b22099..bebe0e35758 100644 --- a/libs/langchain_v1/tests/unit_tests/agents/test_agent_streaming.py +++ b/libs/langchain_v1/tests/unit_tests/agents/test_agent_streaming.py @@ -1,4 +1,4 @@ -"""Unit tests for `AgentStreamer` and `AgentRunStream`.""" +"""Unit tests for create_agent graphs streaming via stream_v2.""" from __future__ import annotations @@ -8,14 +8,9 @@ import pytest from langchain_core.messages import HumanMessage from langchain_core.tools import tool from langgraph.config import emit_tool_output_delta -from langgraph.stream import EventLog, GraphStreamer, StreamTransformer +from langgraph.stream import EventLog, StreamTransformer -from langchain.agents import ( - AgentRunStream, - AgentStreamer, - AsyncAgentRunStream, - create_agent, -) +from langchain.agents import create_agent from tests.unit_tests.agents.model import FakeToolCallingModel if TYPE_CHECKING: @@ -60,13 +55,12 @@ def _single_tool_call_script(name: str, **args: Any) -> list[list[dict[str, Any] ] -class TestAgentStreamerSync: +class TestAgentStreamV2Sync: def test_stream_returns_agent_run_stream(self) -> None: model = FakeToolCallingModel(tool_calls=_single_tool_call_script("echo", text="x")) agent = create_agent(model, [echo]) - run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]}) - assert isinstance(run, AgentRunStream) + run = agent.stream_v2({"messages": [HumanMessage("hi")]}) # Drain so the run closes cleanly. list(run.tool_calls) @@ -76,7 +70,7 @@ class TestAgentStreamerSync: model = FakeToolCallingModel(tool_calls=_single_tool_call_script("echo", text="x")) agent = create_agent(model, [echo]) - run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]}) + run = agent.stream_v2({"messages": [HumanMessage("hi")]}) collected: list[ToolCallStream] = list(run.tool_calls) assert len(collected) == 1 @@ -90,7 +84,7 @@ class TestAgentStreamerSync: model = FakeToolCallingModel(tool_calls=_single_tool_call_script("streamer", text="x")) agent = create_agent(model, [streamer]) - run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]}) + run = agent.stream_v2({"messages": [HumanMessage("hi")]}) tool_calls: list[ToolCallStream] = [] for tc in run.tool_calls: @@ -103,7 +97,7 @@ class TestAgentStreamerSync: model = FakeToolCallingModel() # no tool calls scripted agent = create_agent(model, []) - run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]}) + run = agent.stream_v2({"messages": [HumanMessage("hi")]}) assert list(run.tool_calls) == [] assert run.output is not None @@ -112,7 +106,7 @@ class TestAgentStreamerSync: model = FakeToolCallingModel(tool_calls=_single_tool_call_script("echo", text="x")) agent = create_agent(model, [echo]) - run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]}) + run = agent.stream_v2({"messages": [HumanMessage("hi")]}) # The native `messages` projection is bound as an instance attribute # by `BaseRunStream.__init__` whenever `MessagesTransformer` is # registered. Content population is covered by langgraph tests — @@ -143,7 +137,7 @@ class TestAgentStreamerSync: model = FakeToolCallingModel(tool_calls=_single_tool_call_script("echo", text="x")) agent = create_agent(model, [echo]) - run = AgentStreamer(agent).stream( + run = agent.stream_v2( {"messages": [HumanMessage("hi")]}, transformers=[_Marker], ) @@ -162,7 +156,7 @@ class TestAgentStreamerSync: model = FakeToolCallingModel(tool_calls=_single_tool_call_script("boom")) agent = create_agent(model, [boom]) - run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]}) + run = agent.stream_v2({"messages": [HumanMessage("hi")]}) collected: list[ToolCallStream] = [] @@ -178,37 +172,24 @@ class TestAgentStreamerSync: assert "nope" in collected[0].error assert collected[0].completed is True - def test_plain_graph_streamer_still_works(self) -> None: - """Base `GraphStreamer` on an agent graph works; just no `tool_calls`.""" - model = FakeToolCallingModel(tool_calls=_single_tool_call_script("echo", text="x")) - agent = create_agent(model, [echo]) - run = GraphStreamer(agent).stream({"messages": [HumanMessage("hi")]}) - # Without `ToolCallTransformer` the projection is absent. - assert "tool_calls" not in run._mux.extensions # type: ignore[attr-defined] - assert run.output is not None - - -class TestAgentStreamerAsync: +class TestAgentStreamV2Async: @pytest.mark.anyio async def test_astream_returns_async_agent_run_stream(self) -> None: model = FakeToolCallingModel(tool_calls=_single_tool_call_script("echo", text="x")) agent = create_agent(model, [echo]) - run = await AgentStreamer(agent).astream({"messages": [HumanMessage("hi")]}) - assert isinstance(run, AsyncAgentRunStream) + run = await agent.astream_v2({"messages": [HumanMessage("hi")]}) async for tc in run.tool_calls: async for _ in tc.output_deltas: pass @pytest.mark.anyio async def test_async_tool_deltas_flow(self) -> None: - model = FakeToolCallingModel( - tool_calls=_single_tool_call_script("astreamer", text="hi") - ) + model = FakeToolCallingModel(tool_calls=_single_tool_call_script("astreamer", text="hi")) agent = create_agent(model, [astreamer]) - run = await AgentStreamer(agent).astream({"messages": [HumanMessage("hi")]}) + run = await agent.astream_v2({"messages": [HumanMessage("hi")]}) collected: list[ToolCallStream] = [] async for tc in run.tool_calls: