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, 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="",