From e2a36982500b07ced72b00d8e43bd1b26d86e723 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 17 Oct 2024 15:28:14 -0700
Subject: [PATCH 1/2] For stream_mode=messages skip any nodes/llms with tag
 nostream

---
 libs/langgraph/langgraph/constants.py       |  2 ++
 libs/langgraph/langgraph/pregel/messages.py | 10 +++++++---
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py
index 564d2cac7..52eb72c38 100644
--- a/libs/langgraph/langgraph/constants.py
+++ b/libs/langgraph/langgraph/constants.py
@@ -12,6 +12,8 @@ EMPTY_MAP: Mapping[str, Any] = MappingProxyType({})
 EMPTY_SEQ: tuple[str, ...] = tuple()
 
 # --- Public constants ---
+TAG_NOSTREAM = sys.intern("langsmith:nostream")
+"""Tag to disable streaming for a chat model."""
 TAG_HIDDEN = sys.intern("langsmith:hidden")
 """Tag to hide a node/edge from certain tracing/streaming environments."""
 START = sys.intern("__start__")
diff --git a/libs/langgraph/langgraph/pregel/messages.py b/libs/langgraph/langgraph/pregel/messages.py
index 08327805f..2c31de279 100644
--- a/libs/langgraph/langgraph/pregel/messages.py
+++ b/libs/langgraph/langgraph/pregel/messages.py
@@ -17,7 +17,7 @@ from langchain_core.messages import BaseMessage
 from langchain_core.outputs import ChatGenerationChunk, LLMResult
 from langchain_core.tracers._streaming import T, _StreamingCallbackHandler
 
-from langgraph.constants import NS_SEP
+from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM
 from langgraph.pregel.loop import StreamChunk
 
 Meta = tuple[tuple[str, ...], dict[str, Any]]
@@ -63,7 +63,7 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
         metadata: Optional[dict[str, Any]] = None,
         **kwargs: Any,
     ) -> Any:
-        if metadata:
+        if metadata and (not tags or TAG_NOSTREAM not in tags):
             self.metadata[run_id] = (
                 tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),
                 metadata,
@@ -114,7 +114,11 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
         metadata: Optional[Dict[str, Any]] = None,
         **kwargs: Any,
     ) -> Any:
-        if metadata and kwargs.get("name") == metadata.get("langgraph_node"):
+        if (
+            metadata
+            and kwargs.get("name") == metadata.get("langgraph_node")
+            and (not tags or TAG_HIDDEN not in tags)
+        ):
             self.metadata[run_id] = (
                 tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),
                 metadata,

From 74a17a6d4ce4dbcf09bfa8b577ccfec4b7de3db5 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@langchain.dev>
Date: Thu, 17 Oct 2024 15:52:17 -0700
Subject: [PATCH 2/2] Update tests

---
 libs/langgraph/tests/test_pregel.py       | 12 ------------
 libs/langgraph/tests/test_pregel_async.py | 12 ------------
 2 files changed, 24 deletions(-)

diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py
index a73083290..ded80e20d 100644
--- a/libs/langgraph/tests/test_pregel.py
+++ b/libs/langgraph/tests/test_pregel.py
@@ -4078,18 +4078,6 @@ def test_prebuilt_tool_chat(snapshot: SnapshotAssertion) -> None:
             stream_mode="messages",
         )
     ] == [
-        (
-            _AnyIdHumanMessage(
-                content="what is weather in sf",
-            ),
-            {
-                "langgraph_step": 0,
-                "langgraph_node": "__start__",
-                "langgraph_triggers": ["__start__"],
-                "langgraph_path": ("__pregel_pull", "__start__"),
-                "langgraph_checkpoint_ns": AnyStr("__start__:"),
-            },
-        ),
         (
             _AnyIdAIMessageChunk(
                 content="",
diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py
index 1c45d414f..169311035 100644
--- a/libs/langgraph/tests/test_pregel_async.py
+++ b/libs/langgraph/tests/test_pregel_async.py
@@ -3999,18 +3999,6 @@ async def test_prebuilt_tool_chat() -> None:
             stream_mode="messages",
         )
     ] == [
-        (
-            _AnyIdHumanMessage(
-                content="what is weather in sf",
-            ),
-            {
-                "langgraph_step": 0,
-                "langgraph_node": "__start__",
-                "langgraph_triggers": ["__start__"],
-                "langgraph_path": ("__pregel_pull", "__start__"),
-                "langgraph_checkpoint_ns": AnyStr("__start__:"),
-            },
-        ),
         (
             _AnyIdAIMessageChunk(
                 content="",