mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-12 12:47:53 +02:00
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:
co-authored by
William FH
parent
e3146d8050
commit
b84ae660b8
@@ -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:
|
||||
|
||||
@@ -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"}),
|
||||
|
||||
Reference in New Issue
Block a user