Merge branch 'nh/agent-streamer' into nh/middleware-transformer

# Conflicts:
#	libs/langchain_v1/langchain/agents/__init__.py
#	libs/langchain_v1/langchain/agents/_streaming.py
This commit is contained in:
Nick Hollon
2026-04-20 16:41:45 -04:00
6 changed files with 28 additions and 145 deletions

View File

@@ -5,19 +5,11 @@ from langchain.agents._middleware_transformer import (
MiddlewarePhase,
MiddlewareTransformer,
)
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",
"MiddlewareEvent",
"MiddlewarePhase",
"MiddlewareTransformer",

View File

@@ -30,15 +30,11 @@ if TYPE_CHECKING:
from langgraph.stream._types import ProtocolEvent
MiddlewarePhase = Literal[
"before_agent", "before_model", "after_model", "after_agent"
]
MiddlewarePhase = Literal["before_agent", "before_model", "after_model", "after_agent"]
"""Lifecycle phase a middleware node represents."""
_PHASES: frozenset[str] = frozenset(
("before_agent", "before_model", "after_model", "after_agent")
)
_PHASES: frozenset[str] = frozenset(("before_agent", "before_model", "after_model", "after_agent"))
@dataclass(frozen=True)

View File

@@ -1,83 +0,0 @@
"""Streaming entry point for `create_agent` graphs.
`AgentStreamer` pre-registers `ToolCallTransformer` and
`MiddlewareTransformer` so every agent run exposes `run.tool_calls`
and `run.middleware` 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
from langchain.agents._middleware_transformer import MiddlewareTransformer
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`, `middleware`, `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`
and `MiddlewareTransformer` so `run.tool_calls` and `run.middleware`
are populated on every run without the caller opting in. 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,
MiddlewareTransformer,
)
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)

View File

@@ -22,11 +22,13 @@ 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
from typing_extensions import NotRequired, Required, TypedDict
from langchain.agents._middleware_transformer import MiddlewareTransformer
from langchain.agents.middleware.types import (
AgentMiddleware,
AgentState,
@@ -1660,6 +1662,7 @@ def create_agent(
debug=debug,
name=name,
cache=cache,
transformers=[ToolCallTransformer, MiddlewareTransformer],
).with_config(config)

View File

@@ -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:

View File

@@ -7,7 +7,6 @@ from typing import TYPE_CHECKING, Any
from langchain_core.messages import AIMessage, HumanMessage
from langchain.agents import (
AgentStreamer,
MiddlewareEvent,
MiddlewareTransformer,
create_agent,
@@ -56,9 +55,7 @@ class TestMiddlewareTransformerProcess:
def test_emits_event_for_middleware_node(self) -> None:
tr = _fresh_transformer()
kept = tr.process(
_updates_event({"MyMW.before_model": {"messages": ["hi"]}})
)
kept = tr.process(_updates_event({"MyMW.before_model": {"messages": ["hi"]}}))
assert kept is True # event passes through
events = list(tr._log._items)
assert len(events) == 1
@@ -144,7 +141,7 @@ class TestMiddlewareTransformerProcess:
# ---------------------------------------------------------------------------
# Integration: AgentStreamer auto-registers it; run.middleware works
# Integration: create_agent auto-registers it; run.middleware works
# ---------------------------------------------------------------------------
@@ -153,17 +150,15 @@ class _MarkerMiddleware(AgentMiddleware):
name = "MarkerMW"
def before_model(
self, state: AgentState, runtime: Any
) -> dict[str, Any] | None:
def before_model(self, state: AgentState, runtime: Any) -> dict[str, Any] | None:
return {"messages": []}
class TestAgentStreamerRegistersMiddlewareTransformer:
class TestCreateAgentRegistersMiddlewareTransformer:
def test_middleware_projection_present(self) -> None:
model = FakeToolCallingModel(tool_calls=[[], []])
agent = create_agent(model, [], middleware=[_MarkerMiddleware()])
run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]})
run = agent.stream_v2({"messages": [HumanMessage("hi")]})
assert "middleware" in run.extensions
assert hasattr(run, "middleware")
# Drive to completion.
@@ -176,11 +171,10 @@ class TestAgentStreamerRegistersMiddlewareTransformer:
responses=[AIMessage(content="done", id="a1", tool_calls=[])],
)
agent = create_agent(model, [], middleware=[_MarkerMiddleware()])
run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]})
run = agent.stream_v2({"messages": [HumanMessage("hi")]})
events: list[MiddlewareEvent] = list(run.middleware)
assert any(
ev.phase == "before_model" and ev.middleware_name == "MarkerMW"
for ev in events
ev.phase == "before_model" and ev.middleware_name == "MarkerMW" for ev in events
), f"expected a MarkerMW.before_model event; got {events}"
def test_no_events_when_no_middleware(self) -> None:
@@ -189,6 +183,6 @@ class TestAgentStreamerRegistersMiddlewareTransformer:
responses=[AIMessage(content="done", id="a1", tool_calls=[])],
)
agent = create_agent(model, [])
run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]})
run = agent.stream_v2({"messages": [HumanMessage("hi")]})
events = list(run.middleware)
assert events == []