diff --git a/libs/langgraph/langgraph/pregel/messages.py b/libs/langgraph/langgraph/pregel/messages.py index 5766c2b63..16d0904db 100644 --- a/libs/langgraph/langgraph/pregel/messages.py +++ b/libs/langgraph/langgraph/pregel/messages.py @@ -18,7 +18,7 @@ from langchain_core.messages import BaseMessage from langchain_core.outputs import ChatGenerationChunk, LLMResult from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM -from langgraph.types import StreamChunk +from langgraph.types import Command, StreamChunk try: from langchain_core.tracers._streaming import _StreamingCallbackHandler @@ -153,6 +153,17 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler): **kwargs: Any, ) -> Any: if meta := self.metadata.pop(run_id, None): + if isinstance(response, Command): + response = response.update + + if isinstance(response, Sequence) and any( + isinstance(value, Command) for value in response + ): + response = [ + value.update if isinstance(value, Command) else value + for value in response + ] + if isinstance(response, BaseMessage): self._emit(meta, response, dedupe=True) elif isinstance(response, Sequence): diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 82ba6ff96..cf1c70dab 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -7168,6 +7168,52 @@ def test_tags_stream_mode_messages() -> None: ] +def test_stream_mode_messages_command() -> None: + from langchain_core.messages import HumanMessage + + def my_node(state): + return {"messages": HumanMessage(content="foo")} + + def my_other_node(state): + return Command(update={"messages": HumanMessage(content="bar")}) + + graph = ( + StateGraph(MessagesState) + .add_sequence([my_node, my_other_node]) + .add_edge(START, "my_node") + .compile() + ) + assert list( + graph.stream( + { + "messages": [], + }, + stream_mode="messages", + ) + ) == [ + ( + _AnyIdHumanMessage(content="foo"), + { + "langgraph_step": 1, + "langgraph_node": "my_node", + "langgraph_triggers": ("branch:to:my_node",), + "langgraph_path": ("__pregel_pull", "my_node"), + "langgraph_checkpoint_ns": AnyStr("my_node:"), + }, + ), + ( + _AnyIdHumanMessage(content="bar"), + { + "langgraph_step": 2, + "langgraph_node": "my_other_node", + "langgraph_triggers": ("branch:to:my_other_node",), + "langgraph_path": ("__pregel_pull", "my_other_node"), + "langgraph_checkpoint_ns": AnyStr("my_other_node:"), + }, + ), + ] + + def test_node_destinations() -> None: class State(TypedDict): foo: Annotated[str, operator.add] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 0e74f0b74..987732b90 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -7894,6 +7894,53 @@ async def test_tags_stream_mode_messages() -> None: ] +async def test_stream_mode_messages_command() -> None: + from langchain_core.messages import HumanMessage + + async def my_node(state): + return {"messages": HumanMessage(content="foo")} + + async def my_other_node(state): + return Command(update={"messages": HumanMessage(content="bar")}) + + graph = ( + StateGraph(MessagesState) + .add_sequence([my_node, my_other_node]) + .add_edge(START, "my_node") + .compile() + ) + assert [ + c + async for c in graph.astream( + { + "messages": [], + }, + stream_mode="messages", + ) + ] == [ + ( + _AnyIdHumanMessage(content="foo"), + { + "langgraph_step": 1, + "langgraph_node": "my_node", + "langgraph_triggers": ("branch:to:my_node",), + "langgraph_path": ("__pregel_pull", "my_node"), + "langgraph_checkpoint_ns": AnyStr("my_node:"), + }, + ), + ( + _AnyIdHumanMessage(content="bar"), + { + "langgraph_step": 2, + "langgraph_node": "my_other_node", + "langgraph_triggers": ("branch:to:my_other_node",), + "langgraph_path": ("__pregel_pull", "my_other_node"), + "langgraph_checkpoint_ns": AnyStr("my_other_node:"), + }, + ), + ] + + async def test_stream_messages_dedupe_inputs() -> None: from langchain_core.messages import AIMessage