From b84ae660b8dd1074c32caf81e3462daae5024b45 Mon Sep 17 00:00:00 2001
From: Vadym Barda <vadym@langchain.dev>
Date: Wed, 7 May 2025 23:20:23 -0400
Subject: [PATCH] langgraph: use tuples for streamed message events in
 RemoteGraph (#4589)

Co-authored-by: William FH <13333726+hinthornw@users.noreply.github.com>
---
 libs/langgraph/langgraph/pregel/remote.py |  8 +++
 libs/langgraph/tests/test_remote_graph.py | 78 +++++++++++++++++++++++
 2 files changed, 86 insertions(+)

diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py
index 1805ce58e..7380aa81a 100644
--- a/libs/langgraph/langgraph/pregel/remote.py
+++ b/libs/langgraph/langgraph/pregel/remote.py
@@ -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:
diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py
index c219086a8..247bee68b 100644
--- a/libs/langgraph/tests/test_remote_graph.py
+++ b/libs/langgraph/tests/test_remote_graph.py
@@ -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"}),