diff --git a/libs/langgraph/langgraph/stream/transformers.py b/libs/langgraph/langgraph/stream/transformers.py index 29e247b5e..615b4ebc3 100644 --- a/libs/langgraph/langgraph/stream/transformers.py +++ b/libs/langgraph/langgraph/stream/transformers.py @@ -8,7 +8,7 @@ from langchain_core.language_models.chat_model_stream import ( AsyncChatModelStream, ChatModelStream, ) -from langchain_core.messages import AIMessageChunk, BaseMessage, ToolMessage +from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage from langchain_protocol.protocol import LifecycleCause, MessagesData from typing_extensions import NotRequired, TypedDict @@ -277,11 +277,7 @@ class MessagesTransformer(StreamTransformer): self._route_protocol_event( cast("MessagesData", payload), run_id=run_id, node=node ) - elif ( - isinstance(payload, BaseMessage) - and not isinstance(payload, AIMessageChunk) - and not isinstance(payload, ToolMessage) - ): + elif isinstance(payload, AIMessage) and not isinstance(payload, AIMessageChunk): self._route_whole_message(payload, node=node) # Legacy AIMessageChunk tuples (from on_llm_new_token) are ignored; # v1 streaming callers must switch to stream_events(version="v3") to populate this diff --git a/libs/langgraph/tests/test_stream_messages_transformer.py b/libs/langgraph/tests/test_stream_messages_transformer.py index 30443b773..db2d39334 100644 --- a/libs/langgraph/tests/test_stream_messages_transformer.py +++ b/libs/langgraph/tests/test_stream_messages_transformer.py @@ -14,7 +14,13 @@ from langchain_core.language_models.chat_model_stream import ( AsyncChatModelStream, ChatModelStream, ) -from langchain_core.messages import AIMessage, AIMessageChunk, ToolMessage +from langchain_core.messages import ( + AIMessage, + AIMessageChunk, + HumanMessage, + RemoveMessage, + ToolMessage, +) from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult from langchain_core.runnables import RunnableConfig from typing_extensions import TypedDict @@ -317,6 +323,29 @@ class TestWholeMessageFallback: log.close() assert _unstamped(log._items) == [] + @pytest.mark.parametrize( + "message", + [ + HumanMessage(content="summary", id="human-msg-1"), + RemoveMessage(id="removed-msg-1"), + ], + ) + def test_whole_non_ai_message_is_ignored(self, message: Any) -> None: + t, log = _make_sync_transformer() + t.process( + { + "type": "event", + "method": "messages", + "params": { + "namespace": [], + "timestamp": TS, + "data": (message, {"langgraph_node": "summarize"}), + }, + } + ) + 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")) @@ -586,6 +615,28 @@ class TestEndToEnd: (stream,) = list(run.messages) assert stream.output.text == "hardcoded" + def test_human_message_returned_from_node_is_not_projected(self) -> None: + def summarize(state: MessagesState) -> dict[str, Any]: + return {"messages": HumanMessage(content="summary", id="summary-1")} + + def answer(state: MessagesState) -> dict[str, Any]: + return {"messages": AIMessage(content="answer", id="answer-1")} + + graph = ( + StateGraph(MessagesState) + .add_node("summarize", summarize) + .add_node("answer", answer) + .add_edge(START, "summarize") + .add_edge("summarize", "answer") + .add_edge("answer", END) + .compile() + ) + + run = graph.stream_events({"messages": "hi"}, version="v3") + (stream,) = list(run.messages) + assert stream.node == "answer" + assert stream.output.text == "answer" + @pytest.mark.anyio async def test_async_node_calling_astream_v2(self) -> None: model = GenericFakeChatModel(messages=iter(["async answer"]))