mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-22 23:52:23 +02:00
Christian BromannandGitHub
ea44df3476
## Summary
- Filter `ToolMessage` from v3 `run.messages` streaming (handler +
`MessagesTransformer`) so tool results do not appear as chat text
deltas.
- Normalize `ToolCallStream.output` in `ToolCallTransformer` so live and
serialized `ToolMessage` payloads resolve to raw `content`.
- Add regression tests for message filtering and tool-output unwrapping.
<details>
<summary>Reproducible script</summary>
```python
"""Repro: v3 tool results must not leak through run.messages."""
from __future__ import annotations
import asyncio
from collections.abc import Callable, Sequence
from typing import Any
from uuid import uuid4
from deepagents import create_deep_agent
from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models import BaseChatModel, LanguageModelInput
from langchain_core.messages import AIMessage, HumanMessage, ToolCall, ToolMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from langchain_core.runnables import Runnable
from langchain_core.tools import BaseTool, tool
from pydantic import Field
TOOL_RESULT_SENTINEL = "[]"
class ScriptedChatModel(BaseChatModel):
responses: list[AIMessage] = Field(default_factory=list)
tools: Sequence[dict[str, Any] | type | Callable | BaseTool] = ()
_idx: int = 0
@property
def _llm_type(self) -> str:
return "scripted"
def _generate(
self,
messages: Sequence[Any],
stop: list[str] | None = None,
run_manager: CallbackManagerForLLMRun | None = None,
**kwargs: Any,
) -> ChatResult:
del messages, stop, run_manager, kwargs
idx = min(self._idx, len(self.responses) - 1)
self._idx += 1
return ChatResult(generations=[ChatGeneration(message=self.responses[idx])])
def bind_tools(
self,
tools: Sequence[dict[str, Any] | type | Callable | BaseTool],
*,
tool_choice: str | None = None,
**kwargs: Any,
) -> Runnable[LanguageModelInput, AIMessage]:
del tool_choice, kwargs
self.tools = tools
return self
@tool
def list_items() -> str:
"""List available items."""
return TOOL_RESULT_SENTINEL
def _tool_call_message() -> AIMessage:
return AIMessage(
content="",
tool_calls=[
ToolCall(id="call_list", name="list_items", args={}),
],
)
def _extract_text_delta(event: Any) -> str | None:
if isinstance(event, dict):
if event.get("event") != "content-block-delta":
return None
delta = event.get("delta")
if isinstance(delta, dict) and delta.get("type") == "text-delta":
text = delta.get("text")
return text if isinstance(text, str) else None
elif getattr(event, "event", None) == "content-block-delta":
delta = getattr(event, "delta", None)
if isinstance(delta, dict) and delta.get("type") == "text-delta":
text = delta.get("text")
return text if isinstance(text, str) else None
text = getattr(delta, "text", None)
return text if isinstance(text, str) else None
return None
async def _tool_output(tool_call: Any) -> Any:
output = getattr(tool_call, "output", None)
if callable(output):
return await output()
if hasattr(output, "__await__"):
return await output
return output
async def main() -> None:
model = ScriptedChatModel(
responses=[
_tool_call_message(),
AIMessage(content="No items found."),
]
)
agent = create_deep_agent(model=model, tools=[list_items])
run = await agent.astream_events(
{"messages": [HumanMessage(content="List items")]},
version="v3",
configurable={"thread_id": str(uuid4())},
recursion_limit=50,
)
async def collect_message_texts() -> list[str]:
texts: list[str] = []
async for message_stream in run.messages:
async for event in message_stream:
text = _extract_text_delta(event)
if text is not None:
texts.append(text)
return texts
async def collect_tool_outputs() -> list[Any]:
outputs: list[Any] = []
async for tool_call in run.tool_calls:
outputs.append(await _tool_output(tool_call))
return outputs
message_texts, tool_outputs, final_state = await asyncio.gather(
collect_message_texts(),
collect_tool_outputs(),
run.output(),
)
final_messages = final_state["messages"]
tool_message = next((m for m in final_messages if isinstance(m, ToolMessage)), None)
print("run.messages text deltas:", message_texts)
print("run.tool_calls outputs:", tool_outputs)
print("final state message roles:", [m.type for m in final_messages])
if TOOL_RESULT_SENTINEL in message_texts:
raise AssertionError("Tool result leaked through run.messages.")
if TOOL_RESULT_SENTINEL not in tool_outputs:
raise AssertionError("Tool output was not surfaced through run.tool_calls.")
if tool_message is None or tool_message.tool_call_id != "call_list":
raise AssertionError("Final state does not contain the expected ToolMessage.")
print("Reproduction passed: tool output stayed out of run.messages.")
if __name__ == "__main__":
asyncio.run(main())
```
</details>
<details>
<summary>Current behavior</summary>
```text
run.messages text deltas: ['[]', 'No items found.']
run.tool_calls outputs: [ToolMessage(content='[]', ...)]
final state message roles: ['human', 'ai', 'tool', 'ai']
AssertionError: Tool result leaked through run.messages.
```
</details>
<details>
<summary>Expected behavior</summary>
```text
run.messages text deltas: ['No items found.']
run.tool_calls outputs: ['[]']
final state message roles: ['human', 'ai', 'tool', 'ai']
Reproduction passed: tool output stayed out of run.messages.
```
</details>
## Test plan
- [ ] `uv run --project libs/langgraph pytest
libs/langgraph/tests/test_stream_messages_transformer.py`
- [ ] `uv run --project libs/prebuilt pytest
libs/prebuilt/tests/test_tool_call_transformer.py`
- [ ] `uv run --project libs/langgraph ruff check` (touched files)
- [ ] `uv run --project libs/prebuilt ruff check` (touched files)
Related:
[langchain-ai/langchainjs#10900](https://github.com/langchain-ai/langchainjs/pull/10900)
428 lines
14 KiB
Python
428 lines
14 KiB
Python
"""Tests for ToolCallTransformer and the ToolCallStream projection."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import time
|
|
from typing import Annotated, Any
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, ToolMessage
|
|
from langchain_core.tools import tool
|
|
from langgraph.constants import END, START
|
|
from langgraph.graph import StateGraph
|
|
from langgraph.graph.message import add_messages
|
|
from langgraph.stream._mux import StreamMux
|
|
from langgraph.stream._types import ProtocolEvent
|
|
from langgraph.stream.stream_channel import StreamChannel
|
|
from langgraph.stream.transformers import (
|
|
MessagesTransformer,
|
|
ValuesTransformer,
|
|
)
|
|
from typing_extensions import TypedDict
|
|
|
|
from langgraph.prebuilt import (
|
|
ToolCallTransformer,
|
|
ToolNode,
|
|
ToolRuntime,
|
|
)
|
|
from langgraph.prebuilt._tool_call_stream import ToolCallStream
|
|
|
|
TS = int(time.time() * 1000)
|
|
|
|
|
|
def _unstamped(items):
|
|
"""Strip push stamps from a StreamChannel's internal buffer."""
|
|
return [item for _stamp, item in items]
|
|
|
|
|
|
def _tool_event(
|
|
event: str,
|
|
tool_call_id: str,
|
|
*,
|
|
tool_name: str = "",
|
|
input: dict[str, Any] | None = None,
|
|
delta: Any = None,
|
|
output: Any = None,
|
|
message: str = "",
|
|
namespace: list[str] | None = None,
|
|
) -> ProtocolEvent:
|
|
data: dict[str, Any] = {"event": event, "tool_call_id": tool_call_id}
|
|
if event == "tool-started":
|
|
data["tool_name"] = tool_name
|
|
if input is not None:
|
|
data["input"] = input
|
|
elif event == "tool-output-delta":
|
|
data["delta"] = delta
|
|
elif event == "tool-finished":
|
|
data["output"] = output
|
|
elif event == "tool-error":
|
|
data["message"] = message
|
|
return {
|
|
"type": "event",
|
|
"method": "tools",
|
|
"params": {
|
|
"namespace": namespace or [],
|
|
"timestamp": TS,
|
|
"data": data,
|
|
},
|
|
}
|
|
|
|
|
|
def _subscribe(log: StreamChannel) -> None:
|
|
log._subscribed = True
|
|
|
|
|
|
def _mux() -> tuple[StreamMux, ToolCallTransformer]:
|
|
transformer = ToolCallTransformer()
|
|
mux = StreamMux(
|
|
[
|
|
ValuesTransformer(),
|
|
MessagesTransformer(),
|
|
transformer,
|
|
],
|
|
is_async=False,
|
|
)
|
|
_subscribe(transformer._log)
|
|
return mux, transformer
|
|
|
|
|
|
class TestToolCallTransformerUnit:
|
|
def test_required_stream_modes_declares_tools(self) -> None:
|
|
assert ToolCallTransformer.required_stream_modes == ("tools",)
|
|
|
|
def test_tool_started_yields_handle(self) -> None:
|
|
mux, transformer = _mux()
|
|
mux.push(
|
|
_tool_event(
|
|
"tool-started",
|
|
"tc1",
|
|
tool_name="echo",
|
|
input={"text": "hi"},
|
|
)
|
|
)
|
|
handles = _unstamped(transformer._log._items)
|
|
assert len(handles) == 1
|
|
h = handles[0]
|
|
assert isinstance(h, ToolCallStream)
|
|
assert h.tool_call_id == "tc1"
|
|
assert h.tool_name == "echo"
|
|
assert h.input == {"text": "hi"}
|
|
assert h.completed is False
|
|
|
|
def test_delta_accumulates_on_active_stream(self) -> None:
|
|
mux, transformer = _mux()
|
|
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
|
|
_subscribe(transformer._active["tc1"]._output_deltas)
|
|
mux.push(_tool_event("tool-output-delta", "tc1", delta="a"))
|
|
mux.push(_tool_event("tool-output-delta", "tc1", delta="b"))
|
|
stream = transformer._active["tc1"]
|
|
assert _unstamped(stream._output_deltas._items) == ["a", "b"]
|
|
|
|
def test_finish_closes_stream(self) -> None:
|
|
mux, transformer = _mux()
|
|
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
|
|
stream = transformer._active["tc1"]
|
|
mux.push(_tool_event("tool-finished", "tc1", output="done"))
|
|
assert stream.completed is True
|
|
assert stream.output == "done"
|
|
assert stream.error is None
|
|
assert "tc1" not in transformer._active
|
|
|
|
def test_finish_unwraps_tool_message_output(self) -> None:
|
|
mux, transformer = _mux()
|
|
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
|
|
stream = transformer._active["tc1"]
|
|
mux.push(
|
|
_tool_event(
|
|
"tool-finished",
|
|
"tc1",
|
|
output=ToolMessage(content="done", tool_call_id="tc1"),
|
|
)
|
|
)
|
|
assert stream.completed is True
|
|
assert stream.output == "done"
|
|
|
|
def test_finish_unwraps_serialized_tool_message_output(self) -> None:
|
|
mux, transformer = _mux()
|
|
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
|
|
stream = transformer._active["tc1"]
|
|
mux.push(
|
|
_tool_event(
|
|
"tool-finished",
|
|
"tc1",
|
|
output={
|
|
"lc": 1,
|
|
"type": "constructor",
|
|
"id": ["langchain_core", "messages", "ToolMessage"],
|
|
"kwargs": {
|
|
"content": "serialized done",
|
|
"tool_call_id": "tc1",
|
|
},
|
|
},
|
|
)
|
|
)
|
|
assert stream.completed is True
|
|
assert stream.output == "serialized done"
|
|
|
|
def test_error_closes_stream(self) -> None:
|
|
mux, transformer = _mux()
|
|
mux.push(_tool_event("tool-started", "tc1", tool_name="boom"))
|
|
stream = transformer._active["tc1"]
|
|
mux.push(_tool_event("tool-error", "tc1", message="nope"))
|
|
assert stream.completed is True
|
|
assert stream.output is None
|
|
assert stream.error == "nope"
|
|
assert "tc1" not in transformer._active
|
|
|
|
def test_concurrent_tool_calls_do_not_bleed(self) -> None:
|
|
mux, transformer = _mux()
|
|
mux.push(_tool_event("tool-started", "a", tool_name="t"))
|
|
mux.push(_tool_event("tool-started", "b", tool_name="t"))
|
|
for tc in ("a", "b"):
|
|
_subscribe(transformer._active[tc]._output_deltas)
|
|
mux.push(_tool_event("tool-output-delta", "a", delta="A1"))
|
|
mux.push(_tool_event("tool-output-delta", "b", delta="B1"))
|
|
mux.push(_tool_event("tool-output-delta", "a", delta="A2"))
|
|
assert _unstamped(transformer._active["a"]._output_deltas._items) == [
|
|
"A1",
|
|
"A2",
|
|
]
|
|
assert _unstamped(transformer._active["b"]._output_deltas._items) == ["B1"]
|
|
|
|
def test_tools_event_passes_through_main_log(self) -> None:
|
|
mux, transformer = _mux()
|
|
_subscribe(mux._events)
|
|
mux.push(_tool_event("tool-started", "tc1", tool_name="echo"))
|
|
kept = [e for e in _unstamped(mux._events._items) if e["method"] == "tools"]
|
|
assert len(kept) == 1
|
|
|
|
def test_out_of_scope_event_skipped(self) -> None:
|
|
"""Subgraph-scoped `tools` events must not project into a parent
|
|
transformer's `tool_calls` log.
|
|
|
|
The parent's main event log keeps the event (so wire consumers
|
|
still see it) but the parent's `ToolCallTransformer` only owns
|
|
the projection at its own scope. Per-scope `ToolCallTransformer`
|
|
instances on child mini-muxes are responsible for projecting
|
|
events at their own depth.
|
|
"""
|
|
# Root-scope transformer (`scope == ()`).
|
|
mux, transformer = _mux()
|
|
_subscribe(mux._events)
|
|
mux.push(
|
|
_tool_event(
|
|
"tool-started",
|
|
"tc1",
|
|
tool_name="inner_echo",
|
|
namespace=["child:abc"],
|
|
)
|
|
)
|
|
# No `ToolCallStream` was projected into the root's log.
|
|
assert _unstamped(transformer._log._items) == []
|
|
assert "tc1" not in transformer._active
|
|
# The event still passes through the main event log so consumers
|
|
# of the raw `tools` channel see it untouched.
|
|
kept = [e for e in _unstamped(mux._events._items) if e["method"] == "tools"]
|
|
assert len(kept) == 1
|
|
|
|
def test_in_scope_event_projected_when_scope_set(self) -> None:
|
|
"""A non-root transformer projects only events at its own scope."""
|
|
scope: tuple[str, ...] = ("child:abc",)
|
|
transformer = ToolCallTransformer(scope=scope)
|
|
mux = StreamMux(
|
|
[ValuesTransformer(), MessagesTransformer(), transformer],
|
|
scope=scope,
|
|
is_async=False,
|
|
)
|
|
_subscribe(transformer._log)
|
|
# Event at this scope: projected.
|
|
mux.push(
|
|
_tool_event(
|
|
"tool-started",
|
|
"tc1",
|
|
tool_name="echo",
|
|
namespace=list(scope),
|
|
)
|
|
)
|
|
assert len(_unstamped(transformer._log._items)) == 1
|
|
# Event at a deeper scope: ignored.
|
|
mux.push(
|
|
_tool_event(
|
|
"tool-started",
|
|
"tc2",
|
|
tool_name="grandchild",
|
|
namespace=[*scope, "grand:xyz"],
|
|
)
|
|
)
|
|
assert len(_unstamped(transformer._log._items)) == 1
|
|
# Event at root (above this scope): ignored.
|
|
mux.push(
|
|
_tool_event(
|
|
"tool-started",
|
|
"tc3",
|
|
tool_name="root_tool",
|
|
namespace=[],
|
|
)
|
|
)
|
|
assert len(_unstamped(transformer._log._items)) == 1
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# End-to-end tests with a real graph
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _State(TypedDict):
|
|
messages: Annotated[list, add_messages]
|
|
|
|
|
|
def _build_graph(caller, tools):
|
|
sg = StateGraph(_State)
|
|
sg.add_node("caller", caller)
|
|
sg.add_node("tools", ToolNode(tools))
|
|
sg.add_edge(START, "caller")
|
|
sg.add_edge("caller", "tools")
|
|
sg.add_edge("tools", END)
|
|
return sg.compile()
|
|
|
|
|
|
class TestToolCallTransformerEndToEnd:
|
|
def test_sync_streaming_tool_populates_tool_calls(self) -> None:
|
|
@tool
|
|
def streamer(text: str, runtime: ToolRuntime) -> str:
|
|
"""streams chunks."""
|
|
for chunk in ("one", "two"):
|
|
runtime.emit_output_delta(chunk)
|
|
return text
|
|
|
|
def caller(state: _State) -> dict:
|
|
return {
|
|
"messages": [
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"name": "streamer", "args": {"text": "x"}, "id": "tc1"}
|
|
],
|
|
)
|
|
]
|
|
}
|
|
|
|
graph = _build_graph(caller, [streamer])
|
|
run = graph.stream_events(
|
|
{"messages": []}, transformers=[ToolCallTransformer], version="v3"
|
|
)
|
|
|
|
tool_calls: list[ToolCallStream] = []
|
|
for tc in run.tool_calls:
|
|
tool_calls.append(tc)
|
|
deltas = list(tc.output_deltas)
|
|
assert deltas == ["one", "two"]
|
|
assert len(tool_calls) == 1
|
|
tc = tool_calls[0]
|
|
assert tc.tool_call_id == "tc1"
|
|
assert tc.tool_name == "streamer"
|
|
assert tc.completed is True
|
|
assert tc.error is None
|
|
|
|
def test_stream_modes_union_includes_tools(self) -> None:
|
|
@tool
|
|
def echo(text: str) -> str:
|
|
"""echo."""
|
|
return text
|
|
|
|
def caller(state: _State) -> dict:
|
|
return {
|
|
"messages": [
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"name": "echo", "args": {"text": "x"}, "id": "tc1"}
|
|
],
|
|
)
|
|
]
|
|
}
|
|
|
|
graph = _build_graph(caller, [echo])
|
|
# Without ToolCallTransformer, no tool_calls projection is
|
|
# exposed and no `tools` events flow through (required_stream_modes
|
|
# omits it).
|
|
run_no_tc = graph.stream_events({"messages": []}, version="v3")
|
|
assert "tool_calls" not in run_no_tc._mux.extensions # type: ignore[attr-defined]
|
|
|
|
# With ToolCallTransformer, the projection is present.
|
|
run = graph.stream_events(
|
|
{"messages": []}, transformers=[ToolCallTransformer], version="v3"
|
|
)
|
|
assert "tool_calls" in run._mux.extensions # type: ignore[attr-defined]
|
|
# Drain so the run closes cleanly.
|
|
list(run.tool_calls)
|
|
|
|
@pytest.mark.anyio
|
|
async def test_async_streaming_tool_populates_tool_calls(self) -> None:
|
|
@tool
|
|
async def astreamer(text: str, runtime: ToolRuntime) -> str:
|
|
"""async streams."""
|
|
runtime.emit_output_delta(text)
|
|
runtime.emit_output_delta(text + "!")
|
|
return text
|
|
|
|
async def caller(state: _State) -> dict:
|
|
return {
|
|
"messages": [
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"name": "astreamer", "args": {"text": "hi"}, "id": "tc1"}
|
|
],
|
|
)
|
|
]
|
|
}
|
|
|
|
graph = _build_graph(caller, [astreamer])
|
|
run = await graph.astream_events(
|
|
{"messages": []}, version="v3", transformers=[ToolCallTransformer]
|
|
)
|
|
|
|
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
|
|
assert collected[0].error is None
|
|
|
|
def test_tool_error_populates_error_field(self) -> None:
|
|
@tool
|
|
def boom() -> str:
|
|
"""raises."""
|
|
raise ValueError("nope")
|
|
|
|
def caller(state: _State) -> dict:
|
|
return {
|
|
"messages": [
|
|
AIMessage(
|
|
content="",
|
|
tool_calls=[{"name": "boom", "args": {}, "id": "tc1"}],
|
|
)
|
|
]
|
|
}
|
|
|
|
graph = _build_graph(caller, [boom])
|
|
run = graph.stream_events(
|
|
{"messages": []}, transformers=[ToolCallTransformer], version="v3"
|
|
)
|
|
|
|
collected: list[ToolCallStream] = []
|
|
with pytest.raises(ValueError, match="nope"):
|
|
for tc in run.tool_calls:
|
|
collected.append(tc)
|
|
# Drain deltas so the error field is populated before we
|
|
# inspect it below.
|
|
list(tc.output_deltas)
|
|
|
|
assert len(collected) == 1
|
|
assert collected[0].error == "nope"
|
|
assert collected[0].output is None
|
|
assert collected[0].completed is True
|