From 43d9567c5d19b9dd70a3aab0800442035cda5b39 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 30 Oct 2024 13:47:27 -0700 Subject: [PATCH 1/5] fix: RemoteGraph should propagate subgraphs streaming --- libs/langgraph/langgraph/pregel/remote.py | 41 +++++++++++++++++------ 1 file changed, 31 insertions(+), 10 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index e47f5f50b..36467e8b0 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -31,11 +31,11 @@ from langgraph_sdk.schema import StreamMode as StreamModeSDK from typing_extensions import Self from langgraph.checkpoint.base import CheckpointMetadata -from langgraph.constants import INTERRUPT +from langgraph.constants import CONF, CONFIG_KEY_STREAM, INTERRUPT from langgraph.errors import GraphInterrupt from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode -from langgraph.types import Interrupt +from langgraph.types import Interrupt, StreamProtocol from langgraph.utils.config import merge_configs @@ -557,18 +557,31 @@ class RemoteGraph(PregelProtocol): merged_config = merge_configs(self.config, config) sanitized_config = self._sanitize_config(merged_config) stream_modes, req_updates, req_single = self._get_stream_modes(stream_mode) + stream: Optional[StreamProtocol] = config.get(CONF, {}).get(CONFIG_KEY_STREAM) + stream_modes_ext: list[StreamMode] = ( + [*stream_modes, *stream.modes] if stream else stream_modes + ) for chunk in sync_client.runs.stream( thread_id=sanitized_config["configurable"].get("thread_id"), assistant_id=self.name, input=input, config=sanitized_config, - stream_mode=stream_modes, + stream_mode=stream_modes_ext, interrupt_before=interrupt_before, interrupt_after=interrupt_after, - stream_subgraphs=subgraphs, + stream_subgraphs=subgraphs or stream is not None, if_not_exists="create", ): + if "|" in chunk.event: + mode, ns_ = chunk.event.split("|", 1) + ns = tuple(ns_.split("|")) + else: + mode, ns = chunk.event, () + if stream is not None and chunk.event in stream.modes: + stream((ns, mode, chunk.data)) + if chunk.event not in stream_modes: + continue if chunk.event.startswith("updates"): if isinstance(chunk.data, dict) and INTERRUPT in chunk.data: raise GraphInterrupt(chunk.data[INTERRUPT]) @@ -622,18 +635,31 @@ class RemoteGraph(PregelProtocol): merged_config = merge_configs(self.config, config) sanitized_config = self._sanitize_config(merged_config) stream_modes, req_updates, req_single = self._get_stream_modes(stream_mode) + stream: Optional[StreamProtocol] = config.get(CONF, {}).get(CONFIG_KEY_STREAM) + stream_modes_ext: list[StreamMode] = ( + [*stream_modes, *stream.modes] if stream else stream_modes + ) async for chunk in client.runs.stream( thread_id=sanitized_config["configurable"].get("thread_id"), assistant_id=self.name, input=input, config=sanitized_config, - stream_mode=stream_modes, + stream_mode=stream_modes_ext, interrupt_before=interrupt_before, interrupt_after=interrupt_after, stream_subgraphs=subgraphs, if_not_exists="create", ): + if "|" in chunk.event: + mode, ns_ = chunk.event.split("|", 1) + ns = tuple(ns_.split("|")) + else: + mode, ns = chunk.event, () + if stream is not None and chunk.event in stream.modes: + stream((ns, mode, chunk.data)) + if chunk.event not in stream_modes: + continue if chunk.event.startswith("updates"): if isinstance(chunk.data, dict) and INTERRUPT in chunk.data: raise GraphInterrupt(chunk.data[INTERRUPT]) @@ -642,11 +668,6 @@ class RemoteGraph(PregelProtocol): elif chunk.event.startswith("error"): raise RemoteException(chunk.data) if subgraphs: - if "|" in chunk.event: - mode, ns_ = chunk.event.split("|", 1) - ns = tuple(ns_.split("|")) - else: - mode, ns = chunk.event, () if req_single: yield ns, chunk.data else: From 61abee4bd46d8417e785b7e535f45966a78e62a9 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 30 Oct 2024 13:59:31 -0700 Subject: [PATCH 2/5] Lint --- libs/langgraph/langgraph/pregel/loop.py | 4 ++-- libs/langgraph/langgraph/pregel/remote.py | 18 +++++++++++++----- libs/langgraph/langgraph/types.py | 6 ++++-- libs/sdk-py/langgraph_sdk/schema.py | 4 +++- 4 files changed, 22 insertions(+), 10 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/loop.py b/libs/langgraph/langgraph/pregel/loop.py index 29a2d3ccb..2ba8af3a6 100644 --- a/libs/langgraph/langgraph/pregel/loop.py +++ b/libs/langgraph/langgraph/pregel/loop.py @@ -116,7 +116,7 @@ def DuplexStream(*streams: StreamProtocol) -> StreamProtocol: def __call__(value: StreamChunk) -> None: for stream in streams: if value[1] in stream.modes: - stream(value) # type: ignore + stream(value) return StreamProtocol(__call__, {mode for s in streams for mode in s.modes}) @@ -587,7 +587,7 @@ class PregelLoop(LoopProtocol): if mode not in self.stream.modes: return for v in values(*args, **kwargs): - self.stream((self.checkpoint_ns, mode, v)) # type: ignore + self.stream((self.checkpoint_ns, mode, v)) def _output_writes( self, task_id: str, writes: Sequence[tuple[str, Any]], *, cached: bool = False diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 36467e8b0..d5e9416c1 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -507,7 +507,7 @@ class RemoteGraph(PregelProtocol): 'updates' mode is added to the list of stream modes so that interrupts can be detected in the remote graph. """ - updated_stream_modes: list[StreamMode] = [] + updated_stream_modes: list[StreamModeSDK] = [] req_updates = False req_single = True # coerce to list, or add default stream mode @@ -519,6 +519,10 @@ class RemoteGraph(PregelProtocol): updated_stream_modes.extend(stream_mode) else: updated_stream_modes.append(default) + # map "messages" to "messages-tuple" + if "messages" in updated_stream_modes: + updated_stream_modes.remove("messages") + updated_stream_modes.append("messages-tuple") # add 'updates' mode if not present if "updates" in updated_stream_modes: req_updates = True @@ -557,8 +561,10 @@ class RemoteGraph(PregelProtocol): merged_config = merge_configs(self.config, config) sanitized_config = self._sanitize_config(merged_config) stream_modes, req_updates, req_single = self._get_stream_modes(stream_mode) - stream: Optional[StreamProtocol] = config.get(CONF, {}).get(CONFIG_KEY_STREAM) - stream_modes_ext: list[StreamMode] = ( + stream: Optional[StreamProtocol] = ( + (config or {}).get(CONF, {}).get(CONFIG_KEY_STREAM) + ) + stream_modes_ext: list[StreamModeSDK] = ( [*stream_modes, *stream.modes] if stream else stream_modes ) @@ -635,8 +641,10 @@ class RemoteGraph(PregelProtocol): merged_config = merge_configs(self.config, config) sanitized_config = self._sanitize_config(merged_config) stream_modes, req_updates, req_single = self._get_stream_modes(stream_mode) - stream: Optional[StreamProtocol] = config.get(CONF, {}).get(CONFIG_KEY_STREAM) - stream_modes_ext: list[StreamMode] = ( + stream: Optional[StreamProtocol] = ( + (config or {}).get(CONF, {}).get(CONFIG_KEY_STREAM) + ) + stream_modes_ext: list[StreamModeSDK] = ( [*stream_modes, *stream.modes] if stream else stream_modes ) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index c29cc1480..a4831c96b 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -7,9 +7,11 @@ from typing import ( Literal, NamedTuple, Optional, + Self, Sequence, Type, Union, + cast, ) from langchain_core.runnables import Runnable, RunnableConfig @@ -227,14 +229,14 @@ class StreamProtocol: modes: set[StreamMode] - __call__: Callable[[StreamChunk], None] + __call__: Callable[[Self, StreamChunk], None] def __init__( self, __call__: Callable[[StreamChunk], None], modes: set[StreamMode], ) -> None: - self.__call__ = __call__ + self.__call__ = cast(Callable[[Self, StreamChunk], None], __call__) self.modes = modes diff --git a/libs/sdk-py/langgraph_sdk/schema.py b/libs/sdk-py/langgraph_sdk/schema.py index 3bc43f035..ac34adfe2 100644 --- a/libs/sdk-py/langgraph_sdk/schema.py +++ b/libs/sdk-py/langgraph_sdk/schema.py @@ -26,7 +26,9 @@ Represents the status of a thread: - "error": An exception occurred during task processing. """ -StreamMode = Literal["values", "messages", "updates", "events", "debug", "custom"] +StreamMode = Literal[ + "values", "messages", "updates", "events", "debug", "custom", "messages-tuple" +] """ Defines the mode of streaming: - "values": Stream only the values. From 9c41ce0b1c51a654dd89437f5978d891d0342740 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 30 Oct 2024 14:07:27 -0700 Subject: [PATCH 3/5] Lint --- libs/langgraph/langgraph/types.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/types.py b/libs/langgraph/langgraph/types.py index a4831c96b..279594dec 100644 --- a/libs/langgraph/langgraph/types.py +++ b/libs/langgraph/langgraph/types.py @@ -7,7 +7,6 @@ from typing import ( Literal, NamedTuple, Optional, - Self, Sequence, Type, Union, @@ -15,6 +14,7 @@ from typing import ( ) from langchain_core.runnables import Runnable, RunnableConfig +from typing_extensions import Self from langgraph.checkpoint.base import BaseCheckpointSaver, CheckpointMetadata From 5cb6ff79ca788c3e1378acebc48495f63f79bbd0 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 30 Oct 2024 14:09:43 -0700 Subject: [PATCH 4/5] Prepend ns --- libs/langgraph/langgraph/pregel/remote.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index d5e9416c1..9b36c1f7f 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -31,7 +31,12 @@ from langgraph_sdk.schema import StreamMode as StreamModeSDK from typing_extensions import Self from langgraph.checkpoint.base import CheckpointMetadata -from langgraph.constants import CONF, CONFIG_KEY_STREAM, INTERRUPT +from langgraph.constants import ( + CONF, + CONFIG_KEY_CHECKPOINT_NS, + CONFIG_KEY_STREAM, + INTERRUPT, +) from langgraph.errors import GraphInterrupt from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode @@ -584,6 +589,8 @@ class RemoteGraph(PregelProtocol): ns = tuple(ns_.split("|")) else: mode, ns = chunk.event, () + if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS): + ns = caller_ns + ns if stream is not None and chunk.event in stream.modes: stream((ns, mode, chunk.data)) if chunk.event not in stream_modes: @@ -664,6 +671,8 @@ class RemoteGraph(PregelProtocol): ns = tuple(ns_.split("|")) else: mode, ns = chunk.event, () + if caller_ns := (config or {}).get(CONF, {}).get(CONFIG_KEY_CHECKPOINT_NS): + ns = caller_ns + ns if stream is not None and chunk.event in stream.modes: stream((ns, mode, chunk.data)) if chunk.event not in stream_modes: From 59e8d59b7c6598cd5930d5d6a68316791144140e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 30 Oct 2024 14:17:50 -0700 Subject: [PATCH 5/5] Fix --- libs/langgraph/langgraph/pregel/remote.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 9b36c1f7f..ccf4306f2 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -593,8 +593,6 @@ class RemoteGraph(PregelProtocol): ns = caller_ns + ns if stream is not None and chunk.event in stream.modes: stream((ns, mode, chunk.data)) - if chunk.event not in stream_modes: - continue if chunk.event.startswith("updates"): if isinstance(chunk.data, dict) and INTERRUPT in chunk.data: raise GraphInterrupt(chunk.data[INTERRUPT]) @@ -602,6 +600,8 @@ class RemoteGraph(PregelProtocol): continue elif chunk.event.startswith("error"): raise RemoteException(chunk.data) + if chunk.event.split("|", 1)[0] not in stream_modes: + continue if subgraphs: if "|" in chunk.event: mode, ns_ = chunk.event.split("|", 1) @@ -675,8 +675,6 @@ class RemoteGraph(PregelProtocol): ns = caller_ns + ns if stream is not None and chunk.event in stream.modes: stream((ns, mode, chunk.data)) - if chunk.event not in stream_modes: - continue if chunk.event.startswith("updates"): if isinstance(chunk.data, dict) and INTERRUPT in chunk.data: raise GraphInterrupt(chunk.data[INTERRUPT]) @@ -684,6 +682,8 @@ class RemoteGraph(PregelProtocol): continue elif chunk.event.startswith("error"): raise RemoteException(chunk.data) + if chunk.event.split("|", 1)[0] not in stream_modes: + continue if subgraphs: if req_single: yield ns, chunk.data