From bc404d14f6841c1b130613db959f36269568103c Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 27 May 2025 15:35:32 -0700 Subject: [PATCH] Format --- libs/langgraph/langgraph/pregel/__init__.py | 38 +++++- libs/langgraph/langgraph/pregel/messages.py | 36 +++-- libs/langgraph/tests/test_large_cases.py | 122 ++++++++++++++++- .../langgraph/tests/test_large_cases_async.py | 128 +++++++++++++++++- libs/langgraph/tests/test_pregel_async.py | 1 + 5 files changed, 298 insertions(+), 27 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index dd60f0e76..73ada5aff 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -34,6 +34,7 @@ from langgraph.checkpoint.base import ( CheckpointTuple, copy_checkpoint, ) +from langgraph.config import get_config from langgraph.constants import ( CACHE_NS_WRITES, CONF, @@ -2432,13 +2433,27 @@ class Pregel(PregelProtocol): # set up messages stream mode if "messages" in stream_modes: run_manager.inheritable_handlers.append( - StreamMessagesHandler(stream.put) + StreamMessagesHandler(stream.put, subgraphs) ) # set up custom stream mode if "custom" in stream_modes: config[CONF][CONFIG_KEY_STREAM_WRITER] = lambda c: stream.put( - ((), "custom", c) + ( + tuple( + get_config()[CONF][CONFIG_KEY_CHECKPOINT_NS].split(NS_SEP)[ + :-1 + ] + ), + "custom", + c, + ) ) + elif ( + CONFIG_KEY_STREAM not in config[CONF] + and CONFIG_KEY_STREAM_WRITER in config[CONF] + ): + # remove parent graph stream writer if subgraph streaming not requested + del config[CONF][CONFIG_KEY_STREAM_WRITER] # set checkpointing mode for subgraphs if checkpoint_during is not None: config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during @@ -2658,15 +2673,30 @@ class Pregel(PregelProtocol): # set up messages stream mode if "messages" in stream_modes: run_manager.inheritable_handlers.append( - StreamMessagesHandler(stream_put) + StreamMessagesHandler(stream_put, subgraphs) ) # set up custom stream mode if "custom" in stream_modes: config[CONF][CONFIG_KEY_STREAM_WRITER] = ( lambda c: aioloop.call_soon_threadsafe( - stream.put_nowait, ((), "custom", c) + stream.put_nowait, + ( + tuple( + get_config()[CONF][CONFIG_KEY_CHECKPOINT_NS].split( + NS_SEP + )[:-1] + ), + "custom", + c, + ), ) ) + elif ( + CONFIG_KEY_STREAM not in config[CONF] + and CONFIG_KEY_STREAM_WRITER in config[CONF] + ): + # remove parent graph stream writer if subgraph streaming not requested + del config[CONF][CONFIG_KEY_STREAM_WRITER] # set checkpointing mode for subgraphs if checkpoint_during is not None: config[CONF][CONFIG_KEY_CHECKPOINT_DURING] = checkpoint_during diff --git a/libs/langgraph/langgraph/pregel/messages.py b/libs/langgraph/langgraph/pregel/messages.py index ef82bd483..5014d1cbb 100644 --- a/libs/langgraph/langgraph/pregel/messages.py +++ b/libs/langgraph/langgraph/pregel/messages.py @@ -11,7 +11,7 @@ from uuid import UUID, uuid4 from langchain_core.callbacks import BaseCallbackHandler from langchain_core.messages import BaseMessage -from langchain_core.outputs import ChatGenerationChunk, LLMResult +from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, LLMResult from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM, TAG_NOSTREAM_ALT from langgraph.types import Command, StreamChunk @@ -32,8 +32,9 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler): run_inline = True """We want this callback to run in the main thread, to avoid order/locking issues.""" - def __init__(self, stream: Callable[[StreamChunk], None]): + def __init__(self, stream: Callable[[StreamChunk], None], subgraphs: bool): self.stream = stream + self.subgraphs = subgraphs self.metadata: dict[UUID, Meta] = {} self.seen: set[Union[int, str]] = set() @@ -96,10 +97,15 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler): if metadata and ( not tags or (TAG_NOSTREAM not in tags and TAG_NOSTREAM_ALT not in tags) ): - self.metadata[run_id] = ( - tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)), - metadata, - ) + ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[ + :-1 + ] + if not self.subgraphs and len(ns) > 0: + return + if tags: + if filtered_tags := [t for t in tags if not t.startswith("seq:step")]: + metadata["tags"] = filtered_tags + self.metadata[run_id] = (ns, metadata) def on_llm_new_token( self, @@ -114,9 +120,6 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler): if not isinstance(chunk, ChatGenerationChunk): return if meta := self.metadata.get(run_id): - filtered_tags = [t for t in (tags or []) if not t.startswith("seq:step")] - if filtered_tags: - meta[1]["tags"] = filtered_tags self._emit(meta, chunk.message) def on_llm_end( @@ -127,6 +130,11 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler): parent_run_id: Optional[UUID] = None, **kwargs: Any, ) -> Any: + if meta := self.metadata.get(run_id): + if response.generations and response.generations[0]: + gen = response.generations[0][0] + if isinstance(gen, ChatGeneration): + self._emit(meta, gen.message, dedupe=True) self.metadata.pop(run_id, None) def on_llm_error( @@ -155,10 +163,12 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler): 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, - ) + ns = tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP))[ + :-1 + ] + if not self.subgraphs and len(ns) > 0: + return + self.metadata[run_id] = (ns, metadata) if isinstance(inputs, dict): for key, value in inputs.items(): if isinstance(value, BaseMessage): diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 2a58fccbc..6cc21979e 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -9279,12 +9279,16 @@ def test_weather_subgraph( inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} # run with custom output - assert [c for c in graph.stream(inputs, thread2, stream_mode="custom")] == [ - "I'm", - " very", + assert [ + c for c in graph.stream(inputs, thread2, stream_mode="custom", subgraphs=True) + ] == [ + ((), "I'm"), + ((AnyStr("weather_graph:"),), " very"), ] - assert [c for c in graph.stream(None, thread2, stream_mode="custom")] == [ - " good", + assert [ + c for c in graph.stream(None, thread2, stream_mode="custom", subgraphs=True) + ] == [ + ((AnyStr("weather_graph:"),), " good"), ] # run until interrupt @@ -9602,3 +9606,111 @@ def test_weather_subgraph( }, ), ] + + # run with custom output, without subgraph streaming, should omit subgraph chunks + assert [ + c + for c in graph.stream( + inputs, {"configurable": {"thread_id": "3"}}, stream_mode="custom" + ) + ] == [ + "I'm", + ] + + # run with messages output, with subgraph streaming, should inc subgraph messages + assert [ + c + for c in graph.stream( + inputs, + {"configurable": {"thread_id": "4"}}, + stream_mode="messages", + subgraphs=True, + ) + ] == [ + ( + (), + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ), + { + "thread_id": "4", + "langgraph_step": 1, + "langgraph_node": "router_node", + "langgraph_triggers": ("branch:to:router_node",), + "langgraph_path": ("__pregel_pull", "router_node"), + "langgraph_checkpoint_ns": AnyStr("router_node:"), + "checkpoint_ns": AnyStr("router_node:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ), + ( + (AnyStr("weather_graph:"),), + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="get_weather", + args={"city": "San Francisco"}, + ) + ], + ), + { + "thread_id": "4", + "langgraph_step": 1, + "langgraph_node": "model_node", + "langgraph_triggers": ("branch:to:model_node",), + "langgraph_path": ("__pregel_pull", "model_node"), + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_ns": AnyStr("weather_graph:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ), + ] + + # run with messages output, without subgraph streaming, should exc subgraph messages + assert [ + c + for c in graph.stream( + inputs, + {"configurable": {"thread_id": "5"}}, + stream_mode="messages", + ) + ] == [ + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ), + { + "thread_id": "5", + "langgraph_step": 1, + "langgraph_node": "router_node", + "langgraph_triggers": ("branch:to:router_node",), + "langgraph_path": ("__pregel_pull", "router_node"), + "langgraph_checkpoint_ns": AnyStr("router_node:"), + "checkpoint_ns": AnyStr("router_node:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ] diff --git a/libs/langgraph/tests/test_large_cases_async.py b/libs/langgraph/tests/test_large_cases_async.py index b73483f0a..d4daa7367 100644 --- a/libs/langgraph/tests/test_large_cases_async.py +++ b/libs/langgraph/tests/test_large_cases_async.py @@ -6452,12 +6452,22 @@ async def test_weather_subgraph( inputs = {"messages": [{"role": "user", "content": "what's the weather in sf"}]} # run with custom output - assert [c async for c in graph.astream(inputs, thread2, stream_mode="custom")] == [ - "I'm", - " very", + assert [ + c + async for c in graph.astream( + inputs, thread2, stream_mode="custom", subgraphs=True + ) + ] == [ + ((), "I'm"), + ((AnyStr("weather_graph:"),), " very"), ] - assert [c async for c in graph.astream(None, thread2, stream_mode="custom")] == [ - " good", + assert [ + c + async for c in graph.astream( + None, thread2, stream_mode="custom", subgraphs=True + ) + ] == [ + ((AnyStr("weather_graph:"),), " good"), ] # run until interrupt @@ -6777,3 +6787,111 @@ async def test_weather_subgraph( }, ), ] + + # run with custom output, without subgraph streaming, should omit subgraph chunks + assert [ + c + async for c in graph.astream( + inputs, {"configurable": {"thread_id": "3"}}, stream_mode="custom" + ) + ] == [ + "I'm", + ] + + # run with messages output, with subgraph streaming, should inc subgraph messages + assert [ + c + async for c in graph.astream( + inputs, + {"configurable": {"thread_id": "4"}}, + stream_mode="messages", + subgraphs=True, + ) + ] == [ + ( + (), + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ), + { + "thread_id": "4", + "langgraph_step": 1, + "langgraph_node": "router_node", + "langgraph_triggers": ("branch:to:router_node",), + "langgraph_path": ("__pregel_pull", "router_node"), + "langgraph_checkpoint_ns": AnyStr("router_node:"), + "checkpoint_ns": AnyStr("router_node:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ), + ( + (AnyStr("weather_graph:"),), + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="get_weather", + args={"city": "San Francisco"}, + ) + ], + ), + { + "thread_id": "4", + "langgraph_step": 1, + "langgraph_node": "model_node", + "langgraph_triggers": ("branch:to:model_node",), + "langgraph_path": ("__pregel_pull", "model_node"), + "langgraph_checkpoint_ns": AnyStr("weather_graph:"), + "checkpoint_ns": AnyStr("weather_graph:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ), + ] + + # run with messages output, without subgraph streaming, should exc subgraph messages + assert [ + c + async for c in graph.astream( + inputs, + {"configurable": {"thread_id": "5"}}, + stream_mode="messages", + ) + ] == [ + ( + _AnyIdAIMessage( + content="", + tool_calls=[ + ToolCall( + id="tool_call123", + name="router", + args={"dest": "weather"}, + ) + ], + ), + { + "thread_id": "5", + "langgraph_step": 1, + "langgraph_node": "router_node", + "langgraph_triggers": ("branch:to:router_node",), + "langgraph_path": ("__pregel_pull", "router_node"), + "langgraph_checkpoint_ns": AnyStr("router_node:"), + "checkpoint_ns": AnyStr("router_node:"), + "ls_provider": "fakemessageslistchatmodel", + "ls_model_type": "chat", + }, + ), + ] diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 5577709e6..b9712f4ab 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -5372,6 +5372,7 @@ async def test_stream_subgraphs_during_execution( ] +@NEEDS_CONTEXTVARS async def test_stream_buffering_single_node( async_checkpointer: BaseCheckpointSaver, ) -> None: