From 58cf0c6a6e1420efb7d4c360f1f6b7631e9bf617 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Tue, 22 Oct 2024 10:46:37 -0700 Subject: [PATCH] chore: Switch s3 client utils from httpx client to curl client --- libs/langgraph/langgraph/pregel/remote.py | 64 +++++++------ libs/sdk-py/langgraph_sdk/client.py | 108 +++++++++++----------- 2 files changed, 90 insertions(+), 82 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 28d9abd83..509f048d3 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -351,7 +351,8 @@ class RemoteGraph(PregelProtocol, Runnable): def _get_stream_modes( self, stream_mode: Optional[Union[StreamMode, list[StreamMode]]], - ) -> tuple[list[StreamModeSDK], bool]: + default: StreamMode = "updates", + ) -> tuple[list[StreamModeSDK], bool, bool]: """Return a tuple of the final list of stream modes sent to the remote graph and a boolean flag indicating if stream mode 'updates' was present in the original list of stream modes. @@ -359,23 +360,24 @@ class RemoteGraph(PregelProtocol, Runnable): 'updates' mode is added to the list of stream modes so that interrupts can be detected in the remote graph. """ - updated_stream_modes = [] - updates_mode = False - + updated_stream_modes: list[StreamMode] = [] + req_updates = False + req_single = True + # coerce to list, or add default stream mode if stream_mode: if isinstance(stream_mode, str): updated_stream_modes.append(stream_mode) else: + req_single = False updated_stream_modes.extend(stream_mode) - - if "updates" in updated_stream_modes: - updates_mode = True - else: - updated_stream_modes.append("updates") else: - updated_stream_modes.extend(["values", "updates"]) - - return (updated_stream_modes, updates_mode) + updated_stream_modes.append(default) + # add 'updates' mode if not present + if "updates" in updated_stream_modes: + req_updates = True + else: + updated_stream_modes.append("updates") + return (updated_stream_modes, req_updates, req_single) def stream( self, @@ -389,25 +391,28 @@ class RemoteGraph(PregelProtocol, Runnable): ) -> Iterator[Union[dict[str, Any], Any]]: merged_config = merge_configs(self.config, config) sanitized_config = self._sanitize_config(merged_config) - updated_stream_modes, include_updates = self._get_stream_modes(stream_mode) + stream_modes, req_updates, req_single = self._get_stream_modes(stream_mode) + # TODO if req_subgraphs transform chunk to match Pregel for chunk in self.sync_client.runs.stream( thread_id=sanitized_config["configurable"]["thread_id"], assistant_id=self.graph_id, input=input, config=sanitized_config, - stream_mode=updated_stream_modes, + stream_mode=stream_modes, interrupt_before=interrupt_before, # type: ignore interrupt_after=interrupt_after, # type: ignore stream_subgraphs=subgraphs, ): if chunk.event == "updates": - if INTERRUPT in chunk.data: - raise GraphInterrupt() - if not include_updates: + if isinstance(chunk.data, dict) and INTERRUPT in chunk.data: + raise GraphInterrupt(chunk.data[INTERRUPT]) + if not req_updates: continue - - yield chunk + if req_single: + yield chunk.data + else: + yield chunk async def astream( self, @@ -421,25 +426,27 @@ class RemoteGraph(PregelProtocol, Runnable): ) -> AsyncIterator[Union[dict[str, Any], Any]]: merged_config = merge_configs(self.config, config) sanitized_config = self._sanitize_config(merged_config) - updated_stream_modes, include_updates = self._get_stream_modes(stream_mode) + stream_modes, req_updates, req_single = self._get_stream_modes(stream_mode) async for chunk in self.client.runs.stream( thread_id=sanitized_config["configurable"]["thread_id"], assistant_id=self.graph_id, input=input, config=sanitized_config, - stream_mode=updated_stream_modes, + stream_mode=stream_modes, interrupt_before=interrupt_before, # type: ignore interrupt_after=interrupt_after, # type: ignore stream_subgraphs=subgraphs, ): if chunk.event == "updates": - if INTERRUPT in chunk.data: - raise GraphInterrupt() - if not include_updates: + if isinstance(chunk.data, dict) and INTERRUPT in chunk.data: + raise GraphInterrupt(chunk.data[INTERRUPT]) + if not req_updates: continue - - yield chunk + if req_single: + yield chunk.data + else: + yield chunk async def astream_events( self, @@ -451,10 +458,11 @@ class RemoteGraph(PregelProtocol, Runnable): sanitized_config = self._sanitize_config(merged_config) # manually add 'events' to stream modes list - stream_mode: list[StreamMode] = kwargs.get("stream_mode", []) - updated_stream_modes, include_updates = self._get_stream_modes(stream_mode) + stream_mode: Union[StreamMode, list[StreamMode]] = kwargs.get("stream_mode", []) + updated_stream_modes, include_updates, _ = self._get_stream_modes(stream_mode) if "events" not in updated_stream_modes: updated_stream_modes.append("events") + # TODO bundle main stream events back into StreamEvent async for chunk in self.client.runs.stream( thread_id=sanitized_config["configurable"]["thread_id"], diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index f89a0f6a5..8800d77cf 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -1169,15 +1169,15 @@ class RunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, - feedback_keys: Optional[list[str]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + feedback_keys: Optional[Sequence[str]] = None, on_disconnect: Optional[DisconnectMode] = None, webhook: Optional[str] = None, multitask_strategy: Optional[MultitaskStrategy] = None, @@ -1191,13 +1191,13 @@ class RunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, - feedback_keys: Optional[list[str]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + feedback_keys: Optional[Sequence[str]] = None, on_disconnect: Optional[DisconnectMode] = None, on_completion: Optional[OnCompletionBehavior] = None, webhook: Optional[str] = None, @@ -1210,15 +1210,15 @@ class RunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, - feedback_keys: Optional[list[str]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + feedback_keys: Optional[Sequence[str]] = None, on_disconnect: Optional[DisconnectMode] = None, on_completion: Optional[OnCompletionBehavior] = None, webhook: Optional[str] = None, @@ -1313,12 +1313,12 @@ class RunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, on_completion: Optional[OnCompletionBehavior] = None, after_seconds: Optional[int] = None, @@ -1331,14 +1331,14 @@ class RunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, multitask_strategy: Optional[MultitaskStrategy] = None, after_seconds: Optional[int] = None, @@ -1350,14 +1350,14 @@ class RunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, multitask_strategy: Optional[MultitaskStrategy] = None, on_completion: Optional[OnCompletionBehavior] = None, @@ -1495,8 +1495,8 @@ class RunsClient: config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, on_disconnect: Optional[DisconnectMode] = None, multitask_strategy: Optional[MultitaskStrategy] = None, @@ -1512,8 +1512,8 @@ class RunsClient: input: Optional[dict] = None, metadata: Optional[dict] = None, config: Optional[Config] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, on_disconnect: Optional[DisconnectMode] = None, on_completion: Optional[OnCompletionBehavior] = None, @@ -1530,8 +1530,8 @@ class RunsClient: config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, on_disconnect: Optional[DisconnectMode] = None, on_completion: Optional[OnCompletionBehavior] = None, @@ -3232,15 +3232,15 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, - feedback_keys: Optional[list[str]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + feedback_keys: Optional[Sequence[str]] = None, on_disconnect: Optional[DisconnectMode] = None, webhook: Optional[str] = None, multitask_strategy: Optional[MultitaskStrategy] = None, @@ -3254,13 +3254,13 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, - feedback_keys: Optional[list[str]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + feedback_keys: Optional[Sequence[str]] = None, on_disconnect: Optional[DisconnectMode] = None, on_completion: Optional[OnCompletionBehavior] = None, webhook: Optional[str] = None, @@ -3273,15 +3273,15 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, - feedback_keys: Optional[list[str]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, + feedback_keys: Optional[Sequence[str]] = None, on_disconnect: Optional[DisconnectMode] = None, on_completion: Optional[OnCompletionBehavior] = None, webhook: Optional[str] = None, @@ -3376,12 +3376,12 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, on_completion: Optional[OnCompletionBehavior] = None, after_seconds: Optional[int] = None, @@ -3394,14 +3394,14 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, multitask_strategy: Optional[MultitaskStrategy] = None, after_seconds: Optional[int] = None, @@ -3413,14 +3413,14 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, - stream_mode: Union[StreamMode, list[StreamMode]] = "values", + stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, multitask_strategy: Optional[MultitaskStrategy] = None, on_completion: Optional[OnCompletionBehavior] = None, @@ -3558,8 +3558,8 @@ class SyncRunsClient: config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, on_disconnect: Optional[DisconnectMode] = None, multitask_strategy: Optional[MultitaskStrategy] = None, @@ -3575,8 +3575,8 @@ class SyncRunsClient: input: Optional[dict] = None, metadata: Optional[dict] = None, config: Optional[Config] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, on_disconnect: Optional[DisconnectMode] = None, on_completion: Optional[OnCompletionBehavior] = None, @@ -3593,8 +3593,8 @@ class SyncRunsClient: config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, checkpoint_id: Optional[str] = None, - interrupt_before: Optional[Union[All, list[str]]] = None, - interrupt_after: Optional[Union[All, list[str]]] = None, + interrupt_before: Optional[Union[All, Sequence[str]]] = None, + interrupt_after: Optional[Union[All, Sequence[str]]] = None, webhook: Optional[str] = None, on_disconnect: Optional[DisconnectMode] = None, on_completion: Optional[OnCompletionBehavior] = None,