diff --git a/libs/langchain_v1/langchain/agents/__init__.py b/libs/langchain_v1/langchain/agents/__init__.py index 5e5c545fe21..fbce646b131 100644 --- a/libs/langchain_v1/langchain/agents/__init__.py +++ b/libs/langchain_v1/langchain/agents/__init__.py @@ -1,9 +1,17 @@ """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 new file mode 100644 index 00000000000..247c5100ab8 --- /dev/null +++ b/libs/langchain_v1/langchain/agents/_streaming.py @@ -0,0 +1,79 @@ +"""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/tests/unit_tests/agents/test_agent_streaming.py b/libs/langchain_v1/tests/unit_tests/agents/test_agent_streaming.py new file mode 100644 index 00000000000..45cd3b22099 --- /dev/null +++ b/libs/langchain_v1/tests/unit_tests/agents/test_agent_streaming.py @@ -0,0 +1,219 @@ +"""Unit tests for `AgentStreamer` and `AgentRunStream`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +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 langchain.agents import ( + AgentRunStream, + AgentStreamer, + AsyncAgentRunStream, + create_agent, +) +from tests.unit_tests.agents.model import FakeToolCallingModel + +if TYPE_CHECKING: + from langgraph.prebuilt import ToolCallStream + from langgraph.stream._types import ProtocolEvent + + +@tool +def echo(text: str) -> str: + """Return the input unchanged.""" + return text + + +@tool +def streamer(text: str) -> str: + """Stream two chunks, then return the full text.""" + for chunk in ("one", "two"): + emit_tool_output_delta(chunk) + return text + + +@tool +async def astreamer(text: str) -> str: + """Async: stream two chunks, then return the full text.""" + emit_tool_output_delta(text) + emit_tool_output_delta(text + "!") + return text + + +@tool +def boom() -> str: + """Raise unconditionally.""" + msg = "nope" + raise ValueError(msg) + + +def _single_tool_call_script(name: str, **args: Any) -> list[list[dict[str, Any]]]: + """Script: one tool call on turn 0, finish on turn 1.""" + return [ + [{"name": name, "args": args, "id": "tc1"}], + [], + ] + + +class TestAgentStreamerSync: + 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) + + # Drain so the run closes cleanly. + list(run.tool_calls) + + def test_tool_calls_populated_without_opt_in(self) -> None: + """`ToolCallTransformer` is registered by default on the agent streamer.""" + model = FakeToolCallingModel(tool_calls=_single_tool_call_script("echo", text="x")) + agent = create_agent(model, [echo]) + + run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]}) + + collected: list[ToolCallStream] = list(run.tool_calls) + assert len(collected) == 1 + tc = collected[0] + assert tc.tool_name == "echo" + assert tc.tool_call_id == "tc1" + assert tc.completed is True + assert tc.error is None + + def test_tool_output_deltas_flow_through(self) -> None: + model = FakeToolCallingModel(tool_calls=_single_tool_call_script("streamer", text="x")) + agent = create_agent(model, [streamer]) + + run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]}) + + tool_calls: list[ToolCallStream] = [] + for tc in run.tool_calls: + tool_calls.append(tc) + assert list(tc.output_deltas) == ["one", "two"] + assert len(tool_calls) == 1 + + def test_no_tools_run_is_still_usable(self) -> None: + """`.tool_calls` is empty when the model never calls a tool.""" + model = FakeToolCallingModel() # no tool calls scripted + agent = create_agent(model, []) + + run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]}) + assert list(run.tool_calls) == [] + assert run.output is not None + + def test_messages_projection_present(self) -> None: + """`MessagesTransformer` is inherited from `GraphStreamer.builtin_factories`.""" + model = FakeToolCallingModel(tool_calls=_single_tool_call_script("echo", text="x")) + agent = create_agent(model, [echo]) + + run = AgentStreamer(agent).stream({"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 — + # here we only assert the agent streamer inherits the built-in. + assert "messages" in run._mux.extensions # type: ignore[attr-defined] + assert hasattr(run, "messages") + # Drain so the run closes cleanly. + for tc in run.tool_calls: + list(tc.output_deltas) + + def test_caller_transformers_appended_not_replaced(self) -> None: + """User-supplied transformers add to, rather than replace, the agent defaults.""" + + class _Marker(StreamTransformer): + required_stream_modes = () + + def __init__(self, scope: tuple[str, ...] = ()) -> None: + super().__init__(scope) + self._log: EventLog[int] = EventLog() + + def init(self) -> dict[str, Any]: + return {"marker": self._log} + + def process(self, event: ProtocolEvent) -> bool: + del event + return True + + model = FakeToolCallingModel(tool_calls=_single_tool_call_script("echo", text="x")) + agent = create_agent(model, [echo]) + + run = AgentStreamer(agent).stream( + {"messages": [HumanMessage("hi")]}, + transformers=[_Marker], + ) + # Both the agent default and the user transformer are registered. + assert "tool_calls" in run._mux.extensions # type: ignore[attr-defined] + assert "marker" in run._mux.extensions # type: ignore[attr-defined] + list(run.tool_calls) + + def test_tool_error_sets_error_field(self) -> None: + """Tool errors are surfaced on the `ToolCallStream.error` field. + + `create_agent`'s default tool-error handler re-raises, so the + overall run fails — the assertion here is that the error is + attached to the scoped `ToolCallStream` *before* the run raises. + """ + model = FakeToolCallingModel(tool_calls=_single_tool_call_script("boom")) + agent = create_agent(model, [boom]) + + run = AgentStreamer(agent).stream({"messages": [HumanMessage("hi")]}) + + collected: list[ToolCallStream] = [] + + def _drive() -> None: + for tc in run.tool_calls: + collected.append(tc) + list(tc.output_deltas) + + with pytest.raises(ValueError, match="nope"): + _drive() + assert len(collected) == 1 + assert collected[0].error is not None + 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: + @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) + 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") + ) + agent = create_agent(model, [astreamer]) + + run = await AgentStreamer(agent).astream({"messages": [HumanMessage("hi")]}) + + 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