mirror of
https://github.com/hwchase17/langchain.git
synced 2026-07-12 19:31:24 +00:00
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:
@@ -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",
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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 == []
|
||||
|
||||
Reference in New Issue
Block a user