WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-12 20:57:52 +02:00
Code Issues Packages Projects Releases Wiki Activity
Files
d707ab6650ae06cc8a434eb1ab9303668a14c19a
langgraph/libs/langgraph/tests/test_stream_messages_transformer.py
T
Christian BromannandGitHub ea44df3476 fix(langgraph): keep tool results out of v3 messages (#7838) 
## 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)
2026-05-19 11:36:16 -04:00

941 lines
33 KiB
Python
Raw Blame History

"""Tests for MessagesTransformer: protocol event routing, whole-message fallback,
legacy v1 chunk filtering, and end-to-end via stream_events(version="v3") / astream_events(version="v3")."""
from __future__ import annotations
import time
from typing import Any
import pytest
from langchain_core.language_models import GenericFakeChatModel
from langchain_core.language_models.chat_model_stream import (
AsyncChatModelStream,
ChatModelStream,
)
from langchain_core.messages import AIMessage, AIMessageChunk, ToolMessage
from langchain_core.runnables import RunnableConfig
from typing_extensions import TypedDict
from langgraph.constants import END, START
from langgraph.graph import MessagesState, StateGraph
from langgraph.stream._mux import StreamMux
from langgraph.stream.run_stream import GraphRunStream
from langgraph.stream.stream_channel import StreamChannel
from langgraph.stream.transformers import MessagesTransformer, ValuesTransformer
TS = int(time.time() * 1000)
def _unstamped(items):
"""Strip push stamps from a StreamChannel's internal buffer."""
return [item for _stamp, item in items]
# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _proto_event(
event: dict[str, Any],
*,
run_id: str = "run-1",
node: str = "llm",
) -> dict[str, Any]:
"""Build a messages ProtocolEvent carrying a protocol event dict (v2 path)."""
return {
"type": "event",
"method": "messages",
"params": {
"namespace": [],
"timestamp": TS,
"data": (event, {"langgraph_node": node, "run_id": run_id}),
},
}
def _v1_chunk(
text: str,
msg_id: str = "msg-1",
*,
finish: bool = False,
node: str = "llm",
) -> dict[str, Any]:
"""Build a messages ProtocolEvent carrying a v1 AIMessageChunk tuple."""
rm: dict[str, Any] = {"finish_reason": "stop"} if finish else {}
return {
"type": "event",
"method": "messages",
"params": {
"namespace": [],
"timestamp": TS,
"data": (
AIMessageChunk(content=text, id=msg_id, response_metadata=rm),
{"langgraph_node": node},
),
},
}
def _whole_msg(
text: str,
msg_id: str = "msg-10",
*,
node: str = "node",
) -> dict[str, Any]:
"""Build a messages ProtocolEvent carrying a completed AIMessage."""
return {
"type": "event",
"method": "messages",
"params": {
"namespace": [],
"timestamp": TS,
"data": (AIMessage(content=text, id=msg_id), {"langgraph_node": node}),
},
}
def _make_sync_transformer() -> tuple[
MessagesTransformer, StreamChannel[ChatModelStream]
]:
t = MessagesTransformer()
log: StreamChannel[ChatModelStream] = t.init()["messages"]
log._bind(is_async=False)
# Subscribe up front so pushes during process() are retained.
log._subscribed = True
t._bind_pump(lambda: False)
return t, log
def _make_async_transformer() -> tuple[
MessagesTransformer, StreamChannel[ChatModelStream]
]:
t = MessagesTransformer()
log: StreamChannel[ChatModelStream] = t.init()["messages"]
log._bind(is_async=True)
log._subscribed = True
return t, log
def _lifecycle(
*, text: str = "hello world", message_id: str = "run-1"
) -> list[dict[str, Any]]:
"""Produce a valid protocol event lifecycle: start, delta, finish."""
half = len(text) // 2
first, second = text[:half], text[half:]
return [
{"event": "message-start", "role": "ai", "message_id": message_id},
{
"event": "content-block-start",
"index": 0,
"content_block": {"type": "text", "text": ""},
},
{
"event": "content-block-delta",
"index": 0,
"content_block": {"type": "text", "text": first},
},
{
"event": "content-block-delta",
"index": 0,
"content_block": {"type": "text", "text": second},
},
{
"event": "content-block-finish",
"index": 0,
"content_block": {"type": "text", "text": text},
},
{"event": "message-finish", "reason": "stop"},
]
def _simple_graph():
def call_model(state: MessagesState) -> dict[str, Any]:
model = GenericFakeChatModel(messages=iter(["hello world"]))
stream = model.stream_events(state["messages"], version="v3")
return {"messages": stream.output}
return (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
# ---------------------------------------------------------------------------
# Protocol event routing
# ---------------------------------------------------------------------------
class TestProtocolEventRouting:
def test_message_start_creates_stream(self) -> None:
t, log = _make_sync_transformer()
t.process(
_proto_event(
{"event": "message-start", "role": "ai", "message_id": "run-1"},
run_id="run-1",
)
)
log.close()
(stream,) = _unstamped(log._items)
assert isinstance(stream, ChatModelStream)
assert stream.message_id == "run-1"
def test_full_lifecycle_yields_done_stream(self) -> None:
t, log = _make_sync_transformer()
for evt in _lifecycle(text="hello world"):
t.process(_proto_event(evt, run_id="run-1"))
log.close()
(stream,) = _unstamped(log._items)
assert stream.done
assert stream.output.text == "hello world"
def test_message_finish_cleans_up_routing(self) -> None:
t, log = _make_sync_transformer()
for evt in _lifecycle():
t.process(_proto_event(evt, run_id="run-1"))
assert t._by_run == {}
def test_events_without_prior_start_are_ignored(self) -> None:
t, log = _make_sync_transformer()
t.process(
_proto_event(
{
"event": "content-block-delta",
"index": 0,
"content_block": {"type": "text", "text": "orphan"},
},
run_id="unknown",
)
)
log.close()
assert _unstamped(log._items) == []
def test_tool_role_protocol_events_are_ignored(self) -> None:
t, log = _make_sync_transformer()
for evt in [
{"event": "message-start", "role": "tool", "message_id": "tool-msg-1"},
{
"event": "content-block-delta",
"index": 0,
"content_block": {"type": "text", "text": "[]"},
},
{"event": "message-finish", "reason": "stop"},
]:
t.process(_proto_event(evt, run_id="tool-run"))
log.close()
assert _unstamped(log._items) == []
assert t._ignored_runs == set()
def test_concurrent_streams_routed_by_run_id(self) -> None:
t, log = _make_sync_transformer()
life_a = _lifecycle(text="aaaa", message_id="run-a")
life_b = _lifecycle(text="bbbb", message_id="run-b")
for a, b in zip(life_a, life_b):
t.process(_proto_event(a, run_id="run-a"))
t.process(_proto_event(b, run_id="run-b"))
log.close()
streams = _unstamped(log._items)
assert len(streams) == 2
by_id = {s.message_id: s for s in streams}
assert by_id["run-a"].output.text == "aaaa"
assert by_id["run-b"].output.text == "bbbb"
def test_text_deltas_accumulated_on_stream(self) -> None:
t, log = _make_sync_transformer()
for evt in _lifecycle(text="abcdef"):
t.process(_proto_event(evt))
log.close()
(stream,) = _unstamped(log._items)
assert "".join(stream._text_proj._deltas) == "abcdef"
def test_stream_pushed_on_message_start_not_finish(self) -> None:
# Consumer can see the stream before message-finish arrives.
t, log = _make_sync_transformer()
t.process(
_proto_event(
{"event": "message-start", "role": "ai", "message_id": "run-1"},
run_id="run-1",
)
)
assert len(log._items) == 1
def test_node_metadata_set_on_stream(self) -> None:
t, log = _make_sync_transformer()
t.process(
_proto_event(
{"event": "message-start", "role": "ai", "message_id": "run-1"},
run_id="run-1",
node="my_llm",
)
)
(stream,) = _unstamped(log._items)
assert stream.node == "my_llm"
# ---------------------------------------------------------------------------
# Whole-message fallback
# ---------------------------------------------------------------------------
class TestWholeMessageFallback:
def test_whole_ai_message_produces_complete_stream(self) -> None:
t, log = _make_sync_transformer()
t.process(_whole_msg("the full answer"))
log.close()
(stream,) = _unstamped(log._items)
assert stream.done
assert stream.output.text == "the full answer"
def test_whole_tool_message_is_ignored(self) -> None:
t, log = _make_sync_transformer()
t.process(
{
"type": "event",
"method": "messages",
"params": {
"namespace": [],
"timestamp": TS,
"data": (
ToolMessage(
content="[]",
id="tool-msg-1",
tool_call_id="call_1",
),
{"langgraph_node": "tools"},
),
},
}
)
log.close()
assert _unstamped(log._items) == []
def test_whole_message_has_full_lifecycle(self) -> None:
t, log = _make_sync_transformer()
t.process(_whole_msg("full"))
log.close()
(stream,) = _unstamped(log._items)
assert [e["event"] for e in stream._events] == [
"message-start",
"content-block-start",
"content-block-delta",
"content-block-finish",
"message-finish",
]
# ---------------------------------------------------------------------------
# Filtering
# ---------------------------------------------------------------------------
class TestFiltering:
def test_non_messages_events_pass_through(self) -> None:
t, _ = _make_sync_transformer()
assert (
t.process(
{
"type": "event",
"method": "values",
"params": {"namespace": [], "timestamp": TS, "data": {"x": 1}},
}
)
is True
)
def test_subgraph_namespace_dropped(self) -> None:
t, log = _make_sync_transformer()
t.process(
{
"type": "event",
"method": "messages",
"params": {
"namespace": ["subgraph"],
"timestamp": TS,
"data": (
{"event": "message-start", "message_id": "run-x"},
{"run_id": "run-x"},
),
},
}
)
log.close()
assert _unstamped(log._items) == []
def test_legacy_v1_chunks_ignored(self) -> None:
# v1 AIMessageChunk tuples (from on_llm_new_token) are not streamed
# into this projection; callers must migrate to stream_events(version="v3").
t, log = _make_sync_transformer()
t.process(_v1_chunk("hello"))
t.process(_v1_chunk(" world", finish=True))
log.close()
assert _unstamped(log._items) == []
# ---------------------------------------------------------------------------
# Lifecycle: fail / finalize
# ---------------------------------------------------------------------------
class TestLifecycle:
def test_fail_propagates_to_open_streams(self) -> None:
t, log = _make_sync_transformer()
t.process(
_proto_event(
{"event": "message-start", "message_id": "run-1"}, run_id="run-1"
)
)
streams = _unstamped(log._items)
err = RuntimeError("graph died")
t.fail(err)
assert t._by_run == {}
assert streams[0]._error is err
def test_finalize_clears_routing_state(self) -> None:
t, _ = _make_sync_transformer()
t.process(
_proto_event(
{"event": "message-start", "message_id": "run-1"}, run_id="run-1"
)
)
assert "run-1" in t._by_run
t.finalize()
assert t._by_run == {}
# ---------------------------------------------------------------------------
# Async mode
# ---------------------------------------------------------------------------
class TestAsyncMode:
def test_async_mode_creates_async_stream(self) -> None:
t, log = _make_async_transformer()
for evt in _lifecycle(text="async stream"):
t.process(_proto_event(evt))
assert isinstance(_unstamped(log._items)[0], AsyncChatModelStream)
@pytest.mark.anyio
async def test_text_projection_yields_deltas(self) -> None:
t, log = _make_async_transformer()
for evt in _lifecycle(text="hello world"):
t.process(_proto_event(evt))
(stream,) = _unstamped(log._items)
assert isinstance(stream, AsyncChatModelStream)
assert "".join([d async for d in stream.text]) == "hello world"
@pytest.mark.anyio
async def test_output_awaitable(self) -> None:
t, log = _make_async_transformer()
for evt in _lifecycle(text="async"):
t.process(_proto_event(evt))
(stream,) = _unstamped(log._items)
assert (await stream.output).text == "async"
# ---------------------------------------------------------------------------
# GraphRunStream integration
# ---------------------------------------------------------------------------
class TestWireRequestMore:
def test_bind_pump_called_on_wire(self) -> None:
values_t = ValuesTransformer()
messages_t = MessagesTransformer()
mux = StreamMux([values_t, messages_t], is_async=False)
assert messages_t._pump_fn is None
run = GraphRunStream(iter([]), mux)
assert messages_t._pump_fn is not None
assert messages_t._pump_fn() is False
assert run._exhausted
def test_created_streams_have_request_more(self) -> None:
values_t = ValuesTransformer()
messages_t = MessagesTransformer()
mux = StreamMux([values_t, messages_t], is_async=False)
GraphRunStream(iter([]), mux)
log: StreamChannel[ChatModelStream] = mux.extensions["messages"]
log._subscribed = True
for evt in _lifecycle():
messages_t.process(_proto_event(evt))
(stream,) = _unstamped(log._items)
assert stream._request_more is messages_t._pump_fn
# ---------------------------------------------------------------------------
# End-to-end via StreamMux
# ---------------------------------------------------------------------------
class TestViaMux:
def _make_mux(
self,
) -> tuple[MessagesTransformer, StreamMux, StreamChannel[ChatModelStream]]:
t = MessagesTransformer()
v = ValuesTransformer()
mux = StreamMux([v, t], is_async=False)
t._bind_pump(lambda: False)
log: StreamChannel[ChatModelStream] = mux.extensions["messages"]
log._subscribed = True
return t, mux, log
def test_streaming_via_mux(self) -> None:
t, mux, log = self._make_mux()
for evt in _lifecycle(text="mux stream"):
mux.push(_proto_event(evt))
mux.close()
(stream,) = _unstamped(log._items)
assert stream.output.text == "mux stream"
def test_whole_message_via_mux(self) -> None:
t, mux, log = self._make_mux()
mux.push(_whole_msg("result"))
mux.close()
(stream,) = _unstamped(log._items)
assert stream.output.text == "result"
@pytest.mark.anyio
async def test_async_streaming_via_mux(self) -> None:
t = MessagesTransformer()
v = ValuesTransformer()
mux = StreamMux([v, t], is_async=True)
log: StreamChannel[ChatModelStream] = mux.extensions["messages"]
log._subscribed = True
for evt in _lifecycle(text="async mux"):
await mux.apush(_proto_event(evt))
(stream,) = _unstamped(log._items)
assert (await stream.output).text == "async mux"
await mux.aclose()
# ---------------------------------------------------------------------------
# End-to-end: graph → stream_events(version="v3") → run.messages (node calls stream_events)
# ---------------------------------------------------------------------------
class TestEndToEnd:
"""stream_events(version="v3") path: node calls model.stream_events() explicitly."""
def test_node_calling_stream_v2_populates_messages(self) -> None:
model = GenericFakeChatModel(messages=iter(["hello world"]))
def call_model(state: MessagesState) -> dict[str, Any]:
stream = model.stream_events(state["messages"], version="v3")
return {"messages": stream.output}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = graph.stream_events({"messages": "hi"}, version="v3")
(stream,) = list(run.messages)
assert isinstance(stream, ChatModelStream)
assert stream.output.text == "hello world"
def test_node_stream_v2_text_deltas_iterate(self) -> None:
"""Consumer can iterate `.text` on the streamed message in real time."""
model = GenericFakeChatModel(messages=iter(["streamed answer"]))
def call_model(state: MessagesState) -> dict[str, Any]:
stream = model.stream_events(state["messages"], version="v3")
return {"messages": stream.output}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = graph.stream_events({"messages": "go"}, version="v3")
(stream,) = list(run.messages)
assert "".join(stream.text) == "streamed answer"
def test_non_llm_message_returned_from_node(self) -> None:
"""Whole-message fallback: node returns a finalized AIMessage directly."""
def return_message(state: MessagesState) -> dict[str, Any]:
return {"messages": AIMessage(content="hardcoded", id="msg-abc")}
graph = (
StateGraph(MessagesState)
.add_node("return_message", return_message)
.add_edge(START, "return_message")
.add_edge("return_message", END)
.compile()
)
run = graph.stream_events({"messages": "hi"}, version="v3")
(stream,) = list(run.messages)
assert stream.output.text == "hardcoded"
@pytest.mark.anyio
async def test_async_node_calling_astream_v2(self) -> None:
model = GenericFakeChatModel(messages=iter(["async answer"]))
async def call_model(state: MessagesState) -> dict[str, Any]:
stream = await model.astream_events(state["messages"], version="v3")
return {"messages": await stream}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = await graph.astream_events({"messages": "hi"}, version="v3")
streams = [s async for s in run.messages]
assert len(streams) == 1
assert isinstance(streams[0], AsyncChatModelStream)
assert (await streams[0].output).text == "async answer"
@pytest.mark.anyio
async def test_nested_async_iteration_yields_text_deltas(self) -> None:
"""Inner stream.text drives the shared graph pump via the async pump binding."""
import asyncio
model = GenericFakeChatModel(messages=iter(["hello world"]))
async def call_model(state: MessagesState) -> dict[str, Any]:
stream = await model.astream_events(state["messages"], version="v3")
return {"messages": await stream}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = await graph.astream_events({"messages": "hi"}, version="v3")
async def consume() -> list[str]:
collected: list[str] = []
async for stream in run.messages:
async for delta in stream.text:
collected.append(delta)
return collected
assert "".join(await asyncio.wait_for(consume(), timeout=2.0)) == "hello world"
# ---------------------------------------------------------------------------
# End-to-end: graph → stream_events(version="v3") → run.messages (node calls invoke)
# ---------------------------------------------------------------------------
class TestEndToEndV2Invoke:
"""Auto-routing path: stream_events(version="v3") injects CONFIG_KEY_STREAM_MESSAGES_V2,
causing BaseChatModel to drive the v2 protocol event generator even for
model.invoke()."""
def _graph(self, model):
def call_model(state: MessagesState) -> dict[str, Any]:
return {"messages": model.invoke(state["messages"])}
return (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
def test_invoke_populates_messages(self) -> None:
run = self._graph(
GenericFakeChatModel(messages=iter(["hello world"]))
).stream_events({"messages": "hi"}, version="v3")
(stream,) = list(run.messages)
assert isinstance(stream, ChatModelStream)
assert stream.output.text == "hello world"
def test_invoke_emits_protocol_events(self) -> None:
"""Iterating the stream yields the full v2 lifecycle, not v1 chunks."""
run = self._graph(
GenericFakeChatModel(messages=iter(["streamed answer"]))
).stream_events({"messages": "go"}, version="v3")
(stream,) = list(run.messages)
events = list(stream)
event_types = [e.get("event") for e in events]
assert "message-start" in event_types
assert "content-block-start" in event_types
assert "content-block-delta" in event_types
assert "content-block-finish" in event_types
assert "message-finish" in event_types
# Sanity: every event is a dict carrying an "event" key — not an
# AIMessageChunk tuple from the v1 path.
for event in events:
assert isinstance(event, dict)
assert "event" in event
# Typed projection still assembles the final text.
assert stream.output.text == "streamed answer"
def test_invoke_text_deltas_iterate(self) -> None:
run = self._graph(
GenericFakeChatModel(messages=iter(["delta streaming works"]))
).stream_events({"messages": "hi"}, version="v3")
(stream,) = list(run.messages)
assert "".join(stream.text) == "delta streaming works"
def test_invoke_two_nodes_two_streams(self) -> None:
model_a = GenericFakeChatModel(messages=iter(["alpha"]))
model_b = GenericFakeChatModel(messages=iter(["beta"]))
def node_a(state: MessagesState) -> dict[str, Any]:
return {"messages": model_a.invoke(state["messages"])}
def node_b(state: MessagesState) -> dict[str, Any]:
return {"messages": model_b.invoke(state["messages"])}
graph = (
StateGraph(MessagesState)
.add_node("node_a", node_a)
.add_node("node_b", node_b)
.add_edge(START, "node_a")
.add_edge("node_a", "node_b")
.add_edge("node_b", END)
.compile()
)
streams = list(graph.stream_events({"messages": "hi"}, version="v3").messages)
assert len(streams) == 2
assert {s.output.text for s in streams} == {"alpha", "beta"}
def test_invoke_plus_constructed_message_two_streams(self) -> None:
"""Live-streamed node + constructed-message node → two ChatModelStreams."""
model = GenericFakeChatModel(messages=iter(["live stream"]))
def streaming_node(state: MessagesState) -> dict[str, Any]:
return {"messages": model.invoke(state["messages"])}
def constructed_node(state: MessagesState) -> dict[str, Any]:
return {"messages": [AIMessage(content="hardcoded", id="constructed-1")]}
graph = (
StateGraph(MessagesState)
.add_node("streaming_node", streaming_node)
.add_node("constructed_node", constructed_node)
.add_edge(START, "streaming_node")
.add_edge("streaming_node", "constructed_node")
.add_edge("constructed_node", END)
.compile()
)
run = graph.stream_events({"messages": "hi"}, version="v3")
streams = list(run.messages)
assert len(streams) == 2
assert streams[0].node == "streaming_node"
assert streams[0].output.text == "live stream"
assert streams[1].node == "constructed_node"
assert streams[1].output.text == "hardcoded"
assert streams[1].message_id == "constructed-1"
@pytest.mark.anyio
async def test_ainvoke_populates_messages(self) -> None:
model = GenericFakeChatModel(messages=iter(["async invoke"]))
async def call_model(state: MessagesState) -> dict[str, Any]:
return {"messages": await model.ainvoke(state["messages"])}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
run = await graph.astream_events({"messages": "hi"}, version="v3")
streams = [s async for s in run.messages]
assert len(streams) == 1
assert isinstance(streams[0], AsyncChatModelStream)
assert (await streams[0].output).text == "async invoke"
# ---------------------------------------------------------------------------
# Regression: direct stream_mode="messages" must stay v1
# ---------------------------------------------------------------------------
class TestDirectMessagesModeStaysV1:
def test_direct_graph_stream_messages_yields_ai_message_chunks(self) -> None:
"""graph.stream(stream_mode="messages") must not leak v2 event dicts —
the v2 flag is only injected by stream_events(version="v3") / astream_events(version="v3")."""
model = GenericFakeChatModel(messages=iter(["legacy path"]))
def call_model(state: MessagesState) -> dict[str, Any]:
return {"messages": model.invoke(state["messages"])}
graph = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
parts = list(graph.stream({"messages": "hi"}, stream_mode="messages"))
assert parts, "expected stream_mode='messages' to emit tuples"
for payload, _metadata in parts:
assert isinstance(payload, AIMessageChunk)
assert (
"".join(p[0].content for p in parts if isinstance(p[0].content, str))
== "legacy path"
)
def test_nested_graph_stream_messages_stays_v1_under_outer_stream_events_v3(
self,
) -> None:
"""An outer `stream_events(version="v3")` run must not flip an inner direct
`stream_mode="messages"` call onto the v2 event protocol."""
model = GenericFakeChatModel(messages=iter(["nested legacy path"]))
def call_model(state: MessagesState) -> dict[str, Any]:
return {"messages": model.invoke(state["messages"])}
inner = (
StateGraph(MessagesState)
.add_node("call_model", call_model)
.add_edge(START, "call_model")
.add_edge("call_model", END)
.compile()
)
class OuterState(TypedDict, total=False):
saw_only_chunks: bool
first_payload_type: str
text: str
def call_subgraph(state: OuterState, config: RunnableConfig) -> dict[str, Any]:
parts = list(
inner.stream(
{"messages": "hi"},
config,
stream_mode="messages",
)
)
assert parts
payloads = [payload for payload, _metadata in parts]
return {
"saw_only_chunks": all(
isinstance(payload, AIMessageChunk) for payload in payloads
),
"first_payload_type": type(payloads[0]).__name__,
"text": "".join(
payload.content
for payload in payloads
if isinstance(payload, AIMessageChunk)
and isinstance(payload.content, str)
),
}
outer = (
StateGraph(OuterState)
.add_node("call_subgraph", call_subgraph)
.add_edge(START, "call_subgraph")
.add_edge("call_subgraph", END)
.compile()
)
result = outer.stream_events({}, version="v3").output
assert result is not None
assert result["saw_only_chunks"] is True
assert result["first_payload_type"] == "AIMessageChunk"
assert result["text"] == "nested legacy path"
# ---------------------------------------------------------------------------
# StreamMessagesHandlerV2 unit
# ---------------------------------------------------------------------------
class TestStreamMessagesHandlerV2Unit:
def test_on_llm_new_token_is_noop(self) -> None:
"""v2 handler must not emit v1 chunks even when on_llm_new_token fires."""
from uuid import uuid4
from langchain_core.outputs import ChatGenerationChunk
from langgraph.pregel._messages import StreamMessagesHandlerV2
emitted: list[Any] = []
handler = StreamMessagesHandlerV2(emitted.append, subgraphs=False)
run_id = uuid4()
handler.metadata[run_id] = ((), {"langgraph_node": "x"})
handler.on_llm_new_token(
"hello",
chunk=ChatGenerationChunk(message=AIMessageChunk(content="hello")),
run_id=run_id,
)
assert emitted == []
def test_on_chain_end_does_not_emit_tool_messages(self) -> None:
from uuid import uuid4
from langgraph.pregel._messages import StreamMessagesHandlerV2
emitted: list[Any] = []
handler = StreamMessagesHandlerV2(emitted.append, subgraphs=False)
run_id = uuid4()
handler.metadata[run_id] = ((), {"langgraph_node": "tools"})
handler.on_chain_end(
{"messages": [ToolMessage(content="[]", tool_call_id="call_1")]},
run_id=run_id,
)
assert emitted == []
def test_on_llm_end_dedupes_when_final_message_id_differs(self) -> None:
"""A streamed v2 message should not be emitted again from the final
AIMessage fallback when its final id does not match `message-start`."""
from uuid import uuid4
from langchain_core.outputs import ChatGeneration, LLMResult
from langgraph.pregel._messages import StreamMessagesHandlerV2
emitted: list[Any] = []
handler = StreamMessagesHandlerV2(emitted.append, subgraphs=False)
run_id = uuid4()
handler.metadata[run_id] = ((), {"langgraph_node": "x"})
handler.on_stream_event(
{"event": "message-start", "message_id": "stream-msg-1"},
run_id=run_id,
)
handler.on_llm_end(
LLMResult(
generations=[
[
ChatGeneration(
message=AIMessage(content="hello", id="final-msg-1")
)
]
]
),
run_id=run_id,
)
assert len(emitted) == 1
Reference in New Issue View Git Blame Copy Permalink