mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-27 12:04:58 +02:00
fix(langgraph): support emitting messages from subgraphs when messages mode explicitly requested (#5836)
Reproduces: https://github.com/langchain-ai/langgraph/issues/5249#issuecomment-3156519635 Caused after this change: https://github.com/langchain-ai/langgraph/pull/4843 Fix to allow emitting messages from subgraphs if the subgraphs explicitly used a stream mode "messages". ```python def node_in_parent(...): # subgraph was called as a function. # messages are explicitly requested. for event in subgraph.stream(..., stream_mode="messages"): # something is done with `event` return ... # subgraphs = False! parent_graph.invoke(..., subgraphs=False) ``` The code above should continue to work correctly regardless of the value of subgraphs as streaming messages was requested explicitly in the parent node! --------- Co-authored-by: Sydney Runkle <54324534+sydney-runkle@users.noreply.github.com>
This commit is contained in:
co-authored by
Sydney Runkle
parent
e365b2b8bd
commit
4571b708d9
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user