wehub-resource-sync
a7d6d88f6f
CI / changes (push) Has been cancelled
CI / cd libs/checkpoint (push) Has been cancelled
CI / cd libs/checkpoint-conformance (push) Has been cancelled
CI / cd libs/checkpoint-postgres (push) Has been cancelled
CI / cd libs/checkpoint-sqlite (push) Has been cancelled
CI / cd libs/cli (push) Has been cancelled
CI / cd libs/prebuilt (push) Has been cancelled
CI / cd libs/sdk-py (push) Has been cancelled
CI / cd libs/langgraph (push) Has been cancelled
CI / Check SDK methods matching (push) Has been cancelled
CI / Check CLI schema hasn't changed #3.13 (push) Has been cancelled
CI / CLI integration test (push) Has been cancelled
CI / sdk-py integration test (push) Has been cancelled
CI / CI Success (push) Has been cancelled
baseline / benchmark (push) Has been cancelled
Deploy Redirects to GitHub Pages / deploy (push) Has been cancelled
941 lines
33 KiB
Python
941 lines
33 KiB
Python
"""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
|