langgraph: use tuples for streamed message events in RemoteGraph (#4589)

Co-authored-by: William FH <13333726+hinthornw@users.noreply.github.com>
This commit is contained in:
Vadym Barda
2025-05-08 03:20:23 +00:00
committed by GitHub
co-authored by William FH
parent e3146d8050
commit b84ae660b8
2 changed files with 86 additions and 0 deletions
@@ -695,6 +695,10 @@ class RemoteGraph(PregelProtocol):
# filter for what was actually requested
if mode not in requested:
continue
if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data)) # type: ignore
# emit chunk
if subgraphs:
if NS_SEP in chunk.event:
@@ -790,6 +794,10 @@ class RemoteGraph(PregelProtocol):
# filter for what was actually requested
if mode not in requested:
continue
if chunk.event.startswith("messages"):
chunk = chunk._replace(data=tuple(chunk.data)) # type: ignore
# emit chunk
if subgraphs:
if NS_SEP in chunk.event:
+78
View File
@@ -432,6 +432,21 @@ def test_stream():
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(
event="messages",
data=[
{
"content": [{"text": "Hello", "type": "text", "index": 0}],
"type": "AIMessageChunk",
},
{
"langgraph_step": 1,
"langgraph_node": "call_llm",
"langgraph_triggers": ["branch:to:call_llm"],
"langgraph_path": ["__pregel_pull", "call_llm"],
},
],
),
StreamPart(
event="updates",
data={
@@ -489,6 +504,30 @@ def test_stream():
{"chunk": "data3"},
]
# stream_mode messages
stream_parts = []
for stream_part in remote_pregel.stream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode="messages",
):
stream_parts.append(stream_part)
assert stream_parts == [
(
{
"content": [{"text": "Hello", "type": "text", "index": 0}],
"type": "AIMessageChunk",
},
{
"langgraph_step": 1,
"langgraph_node": "call_llm",
"langgraph_triggers": ["branch:to:call_llm"],
"langgraph_path": ["__pregel_pull", "call_llm"],
},
),
]
mock_sync_client.runs.stream.return_value = [
StreamPart(event="updates", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
@@ -566,6 +605,21 @@ async def test_astream():
StreamPart(event="values", data={"chunk": "data2"}),
StreamPart(event="values", data={"chunk": "data3"}),
StreamPart(event="updates", data={"chunk": "data4"}),
StreamPart(
event="messages",
data=[
{
"content": [{"text": "Hello", "type": "text", "index": 0}],
"type": "AIMessageChunk",
},
{
"langgraph_step": 1,
"langgraph_node": "call_llm",
"langgraph_triggers": ["branch:to:call_llm"],
"langgraph_path": ["__pregel_pull", "call_llm"],
},
],
),
StreamPart(
event="updates",
data={
@@ -624,6 +678,30 @@ async def test_astream():
{"chunk": "data3"},
]
# stream_mode messages
stream_parts = []
async for stream_part in remote_pregel.astream(
{"input": "data"},
config={"configurable": {"thread_id": "thread_1"}},
stream_mode="messages",
):
stream_parts.append(stream_part)
assert stream_parts == [
(
{
"content": [{"text": "Hello", "type": "text", "index": 0}],
"type": "AIMessageChunk",
},
{
"langgraph_step": 1,
"langgraph_node": "call_llm",
"langgraph_triggers": ["branch:to:call_llm"],
"langgraph_path": ["__pregel_pull", "call_llm"],
},
),
]
async_iter = MagicMock()
async_iter.__aiter__.return_value = [
StreamPart(event="updates", data={"chunk": "data3"}),