diff --git a/libs/langgraph/langgraph/pregel/_messages.py b/libs/langgraph/langgraph/pregel/_messages.py index 550ea789c..3ed677150 100644 --- a/libs/langgraph/langgraph/pregel/_messages.py +++ b/libs/langgraph/langgraph/pregel/_messages.py @@ -29,16 +29,51 @@ Meta = tuple[tuple[str, ...], dict[str, Any]] class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler): """A callback handler that implements stream_mode=messages. - Collects messages from (1) chat model stream events and (2) node outputs.""" + + Collects messages from: + (1) chat model stream events; and + (2) node outputs. + """ run_inline = True - """We want this callback to run in the main thread, to avoid order/locking issues.""" + """We want this callback to run in the main thread to avoid order/locking issues.""" - def __init__(self, stream: Callable[[StreamChunk], None], subgraphs: bool): + def __init__( + self, + stream: Callable[[StreamChunk], None], + subgraphs: bool, + *, + parent_ns: tuple[str, ...] | None = None, + ) -> None: + """Configure the handler to stream messages from LLMs and nodes. + + Args: + stream: A callable that takes a StreamChunk and emits it. + subgraphs: Whether to emit messages from subgraphs. + parent_ns: The namespace where the handler was created. + We keep track of this namespace to allow calls to subgraphs that + were explicitly requested as a stream with `messages` mode + configured. + + Example: + parent_ns is used to handle scenarios where the subgraph is explicitly + streamed with `stream_mode="messages"`. + + ```python + def parent_graph_node(): + # This node is in the parent graph. + async for event in some_subgraph(..., stream_mode="messages"): + do something with event # <-- these events will be emitted + return ... + + parent_graph.invoke(subgraphs=False) + ``` + """ self.stream = stream self.subgraphs = subgraphs self.metadata: dict[UUID, Meta] = {} self.seen: set[int | str] = set() + self.parent_ns = parent_ns def _emit(self, meta: Meta, message: BaseMessage, *, dedupe: bool = False) -> None: if dedupe and message.id in self.seen: @@ -100,7 +135,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler): ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[ :-1 ] - if not self.subgraphs and len(ns) > 0: + if not self.subgraphs and len(ns) > 0 and ns != self.parent_ns: return if tags: if filtered_tags := [t for t in tags if not t.startswith("seq:step")]: diff --git a/libs/langgraph/langgraph/pregel/main.py b/libs/langgraph/langgraph/pregel/main.py index 37bd069c2..e8205aee6 100644 --- a/libs/langgraph/langgraph/pregel/main.py +++ b/libs/langgraph/langgraph/pregel/main.py @@ -11,7 +11,7 @@ from collections.abc import AsyncIterator, Iterator, Mapping, Sequence from dataclasses import is_dataclass from functools import partial from inspect import isclass -from typing import Any, Callable, Generic, Union, cast, get_type_hints +from typing import Any, Callable, Generic, Optional, Union, cast, get_type_hints from uuid import UUID, uuid5 from langchain_core.globals import get_debug @@ -2534,8 +2534,13 @@ class Pregel( config[CONF][CONFIG_KEY_CHECKPOINT_NS] = recast_checkpoint_ns(ns) # set up messages stream mode if "messages" in stream_modes: + ns_ = cast(Optional[str], config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)) run_manager.inheritable_handlers.append( - StreamMessagesHandler(stream.put, subgraphs) + StreamMessagesHandler( + stream.put, + subgraphs, + parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, + ) ) # set up custom stream mode @@ -2814,8 +2819,14 @@ class Pregel( config[CONF][CONFIG_KEY_CHECKPOINT_NS] = recast_checkpoint_ns(ns) # set up messages stream mode if "messages" in stream_modes: + # namespace can be None in a root level graph? + ns_ = cast(Optional[str], config[CONF].get(CONFIG_KEY_CHECKPOINT_NS)) run_manager.inheritable_handlers.append( - StreamMessagesHandler(stream_put, subgraphs) + StreamMessagesHandler( + stream_put, + subgraphs, + parent_ns=tuple(ns_.split(NS_SEP)) if ns_ else None, + ) ) # set up custom stream mode diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index a1ec594fd..71283b673 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -27,7 +27,7 @@ from langsmith import traceable from pydantic import BaseModel, ConfigDict, Field, ValidationError from pytest_mock import MockerFixture from syrupy import SnapshotAssertion -from typing_extensions import TypedDict +from typing_extensions import NotRequired, TypedDict from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL from langgraph.cache.base import BaseCache @@ -8266,3 +8266,53 @@ def test_fork_and_update_task_results(sync_checkpointer: BaseCheckpointSaver) -> ], ], ] + + +def test_subgraph_streaming_sync() -> None: + """Test subgraph streaming when used as a node in sync version""" + + # Create a fake chat model that returns a simple response + model = GenericFakeChatModel(messages=iter(["The weather is sunny today."])) + + # Create a subgraph that uses the fake chat model + def call_model_node(state: MessagesState, config: RunnableConfig) -> MessagesState: + """Node that calls the model with the last message.""" + messages = state["messages"] + last_message = messages[-1].content if messages else "" + response = model.invoke([("user", last_message)], config) + return {"messages": [response]} + + # Build the subgraph + subgraph = StateGraph(MessagesState) + subgraph.add_node("call_model", call_model_node) + subgraph.add_edge(START, "call_model") + compiled_subgraph = subgraph.compile() + + class SomeCustomState(TypedDict): + last_chunk: NotRequired[str] + num_chunks: NotRequired[int] + + # Will invoke a subgraph as a function + def parent_node(state: SomeCustomState, config: RunnableConfig) -> dict: + """Node that runs the subgraph.""" + msgs = {"messages": [("user", "What is the weather in Tokyo?")]} + events = [] + for event in compiled_subgraph.stream(msgs, config, stream_mode="messages"): + events.append(event) + ai_msg_chunks = [ai_msg_chunk for ai_msg_chunk, _ in events] + return { + "last_chunk": ai_msg_chunks[-1], + "num_chunks": len(ai_msg_chunks), + } + + # Build the main workflow + workflow = StateGraph(SomeCustomState) + workflow.add_node("subgraph", parent_node) + workflow.add_edge(START, "subgraph") + compiled_workflow = workflow.compile() + + # Test the basic functionality + result = compiled_workflow.invoke({}) + + assert result["last_chunk"].content == "today." + assert result["num_chunks"] == 9 diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index c834ebbdd..6d42a0244 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -26,7 +26,7 @@ from langchain_core.utils.aiter import aclosing from pydantic import BaseModel, ConfigDict, Field, ValidationError from pytest_mock import MockerFixture from syrupy import SnapshotAssertion -from typing_extensions import TypedDict +from typing_extensions import NotRequired, TypedDict from langgraph._internal._constants import CONFIG_KEY_NODE_FINISHED, ERROR, PULL from langgraph.cache.base import BaseCache @@ -9053,3 +9053,57 @@ async def test_fork_and_update_task_results( ], ], ] + + +async def test_subgraph_streaming_async() -> None: + """Test subgraph streaming when used as a node in async version""" + + # Create a fake chat model that returns a simple response + model = GenericFakeChatModel(messages=iter(["The weather is sunny today."])) + + # Create a subgraph that uses the fake chat model + async def call_model_node( + state: MessagesState, config: RunnableConfig + ) -> MessagesState: + """Node that calls the model with the last message.""" + messages = state["messages"] + last_message = messages[-1].content if messages else "" + response = await model.ainvoke([("user", last_message)], config) + return {"messages": [response]} + + # Build the subgraph + subgraph = StateGraph(MessagesState) + subgraph.add_node("call_model", call_model_node) + subgraph.add_edge(START, "call_model") + compiled_subgraph = subgraph.compile() + + class SomeCustomState(TypedDict): + last_chunk: NotRequired[str] + num_chunks: NotRequired[int] + + # Will invoke a subgraph as a function + async def parent_node(state: SomeCustomState, config: RunnableConfig) -> dict: + """Node that runs the subgraph.""" + msgs = {"messages": [("user", "What is the weather in Tokyo?")]} + events = [] + async for event in compiled_subgraph.astream( + msgs, config, stream_mode="messages" + ): + events.append(event) + ai_msg_chunks = [ai_msg_chunk for ai_msg_chunk, _ in events] + return { + "last_chunk": ai_msg_chunks[-1], + "num_chunks": len(ai_msg_chunks), + } + + # Build the main workflow + workflow = StateGraph(SomeCustomState) + workflow.add_node("subgraph", parent_node) + workflow.add_edge(START, "subgraph") + compiled_workflow = workflow.compile() + + # Test the basic functionality + result = await compiled_workflow.ainvoke({}) + + assert result["last_chunk"].content == "today." + assert result["num_chunks"] == 9 diff --git a/libs/prebuilt/tests/test_react_agent.py b/libs/prebuilt/tests/test_react_agent.py index f70e1a88b..cff4c13c3 100644 --- a/libs/prebuilt/tests/test_react_agent.py +++ b/libs/prebuilt/tests/test_react_agent.py @@ -24,7 +24,7 @@ from langchain_core.messages import ( ToolCall, ToolMessage, ) -from langchain_core.runnables import RunnableLambda +from langchain_core.runnables import RunnableConfig, RunnableLambda from langchain_core.tools import InjectedToolCallId, ToolException from langchain_core.tools import tool as dec_tool from pydantic import BaseModel, Field @@ -1236,6 +1236,190 @@ def test_tool_node_stream_writer() -> None: ] +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +def test_react_agent_subgraph_streaming_sync(version: Literal["v1", "v2"]) -> None: + """Test React agent streaming when used as a subgraph node sync version""" + + @dec_tool + def get_weather(city: str) -> str: + """Get the weather of a city.""" + return f"The weather of {city} is sunny." + + # Create a React agent + model = FakeToolCallingModel( + tool_calls=[ + [{"args": {"city": "Tokyo"}, "id": "1", "name": "get_weather"}], + [], + ] + ) + + agent = create_react_agent( + model, + tools=[get_weather], + prompt="You are a helpful travel assistant.", + version=version, + ) + + # Create a subgraph that uses the React agent as a node + def react_agent_node(state: MessagesState, config: RunnableConfig) -> MessagesState: + """Node that runs the React agent and collects streaming output.""" + collected_content = "" + + # Stream the agent output and collect content + for msg_chunk, msg_metadata in agent.stream( + {"messages": [("user", state["messages"][-1].content)]}, + config, + stream_mode="messages", + ): + if hasattr(msg_chunk, "content") and msg_chunk.content: + collected_content += msg_chunk.content + + return {"messages": [("assistant", collected_content)]} + + # Create the main workflow with the React agent as a subgraph node + workflow = StateGraph(MessagesState) + workflow.add_node("react_agent", react_agent_node) + workflow.add_edge(START, "react_agent") + workflow.add_edge("react_agent", "__end__") + compiled_workflow = workflow.compile() + + # Test the streaming functionality + result = compiled_workflow.invoke( + {"messages": [("user", "What is the weather in Tokyo?")]} + ) + + # Verify the result contains expected structure + assert len(result["messages"]) == 2 + assert result["messages"][0].content == "What is the weather in Tokyo?" + assert "assistant" in str(result["messages"][1]) + + # Test streaming with subgraphs = True + result = compiled_workflow.invoke( + {"messages": [("user", "What is the weather in Tokyo?")]}, + subgraphs=True, + ) + assert len(result["messages"]) == 2 + + events = [] + for event in compiled_workflow.stream( + {"messages": [("user", "What is the weather in Tokyo?")]}, + stream_mode="messages", + subgraphs=False, + ): + events.append(event) + + assert len(events) == 0 + + events = [] + for event in compiled_workflow.stream( + {"messages": [("user", "What is the weather in Tokyo?")]}, + stream_mode="messages", + subgraphs=True, + ): + events.append(event) + + assert len(events) == 3 + namespace, (msg, metadata) = events[0] + # FakeToolCallingModel returns a single AIMessage with tool calls + # The content of the AIMessage reflects the input message + assert msg.content.startswith("You are a helpful travel assistant") + namespace, (msg, metadata) = events[1] # ToolMessage + assert msg.content.startswith("The weather of Tokyo is sunny.") + + +@pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) +async def test_react_agent_subgraph_streaming(version: Literal["v1", "v2"]) -> None: + """Test React agent streaming when used as a subgraph node.""" + + @dec_tool + def get_weather(city: str) -> str: + """Get the weather of a city.""" + return f"The weather of {city} is sunny." + + # Create a React agent + model = FakeToolCallingModel( + tool_calls=[ + [{"args": {"city": "Tokyo"}, "id": "1", "name": "get_weather"}], + [], + ] + ) + + agent = create_react_agent( + model, + tools=[get_weather], + prompt="You are a helpful travel assistant.", + version=version, + ) + + # Create a subgraph that uses the React agent as a node + async def react_agent_node( + state: MessagesState, config: RunnableConfig + ) -> MessagesState: + """Node that runs the React agent and collects streaming output.""" + collected_content = "" + + # Stream the agent output and collect content + async for msg_chunk, msg_metadata in agent.astream( + {"messages": [("user", state["messages"][-1].content)]}, + config, + stream_mode="messages", + ): + if hasattr(msg_chunk, "content") and msg_chunk.content: + collected_content += msg_chunk.content + + return {"messages": [("assistant", collected_content)]} + + # Create the main workflow with the React agent as a subgraph node + workflow = StateGraph(MessagesState) + workflow.add_node("react_agent", react_agent_node) + workflow.add_edge(START, "react_agent") + workflow.add_edge("react_agent", "__end__") + compiled_workflow = workflow.compile() + + # Test the streaming functionality + result = await compiled_workflow.ainvoke( + {"messages": [("user", "What is the weather in Tokyo?")]} + ) + + # Verify the result contains expected structure + assert len(result["messages"]) == 2 + assert result["messages"][0].content == "What is the weather in Tokyo?" + assert "assistant" in str(result["messages"][1]) + + # Test streaming with subgraphs = True + result = await compiled_workflow.ainvoke( + {"messages": [("user", "What is the weather in Tokyo?")]}, + subgraphs=True, + ) + assert len(result["messages"]) == 2 + + events = [] + async for event in compiled_workflow.astream( + {"messages": [("user", "What is the weather in Tokyo?")]}, + stream_mode="messages", + subgraphs=False, + ): + events.append(event) + + assert len(events) == 0 + + events = [] + async for event in compiled_workflow.astream( + {"messages": [("user", "What is the weather in Tokyo?")]}, + stream_mode="messages", + subgraphs=True, + ): + events.append(event) + + assert len(events) == 3 + namespace, (msg, metadata) = events[0] + # FakeToolCallingModel returns a single AIMessage with tool calls + # The content of the AIMessage reflects the input message + assert msg.content.startswith("You are a helpful travel assistant") + namespace, (msg, metadata) = events[1] # ToolMessage + assert msg.content.startswith("The weather of Tokyo is sunny.") + + @pytest.mark.parametrize("version", REACT_TOOL_CALL_VERSIONS) def test_tool_node_node_interrupt( sync_checkpointer: BaseCheckpointSaver, version: str