From a3feaef2eb0a25ccdc0abbeedd8107c2b1b8e3b8 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Dec 2024 16:45:07 -0800 Subject: [PATCH 1/3] lib: Handle Command in RemoteGraph --- libs/langgraph/langgraph/pregel/remote.py | 15 ++++++++++++++- libs/sdk-py/langgraph_sdk/client.py | 15 +++++++++++++++ 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index abe27eb28..8ee7ec884 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -1,3 +1,4 @@ +from dataclasses import asdict from typing import ( Any, AsyncIterator, @@ -41,7 +42,7 @@ from langgraph.constants import ( 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, StreamProtocol +from langgraph.types import Command, Interrupt, StreamProtocol from langgraph.utils.config import merge_configs @@ -597,11 +598,17 @@ class RemoteGraph(PregelProtocol): stream_modes, requested, req_single, stream = self._get_stream_modes( stream_mode, config ) + if isinstance(input, Command): + command: dict[str, Any] = asdict(input) + input = None + else: + command = None for chunk in sync_client.runs.stream( thread_id=sanitized_config["configurable"].get("thread_id"), assistant_id=self.name, input=input, + command=command, config=sanitized_config, stream_mode=stream_modes, interrupt_before=interrupt_before, @@ -680,11 +687,17 @@ class RemoteGraph(PregelProtocol): stream_modes, requested, req_single, stream = self._get_stream_modes( stream_mode, config ) + if isinstance(input, Command): + command: dict[str, Any] = asdict(input) + input = None + else: + command = None async for chunk in client.runs.stream( thread_id=sanitized_config["configurable"].get("thread_id"), assistant_id=self.name, input=input, + command=command, config=sanitized_config, stream_mode=stream_modes, interrupt_before=interrupt_before, diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index 81e8a8506..be909d336 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -3303,6 +3303,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3326,6 +3327,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3346,6 +3348,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3370,6 +3373,7 @@ class SyncRunsClient: assistant_id: The assistant ID or graph name to stream from. If using graph name, will default to first assistant created from that graph. input: The input to the graph. + command: The command to execute. stream_mode: The stream mode(s) to use. stream_subgraphs: Whether to stream output from subgraphs. metadata: Metadata to assign to the run. @@ -3420,6 +3424,7 @@ class SyncRunsClient: """ # noqa: E501 payload = { "input": input, + "command": command, "config": config, "metadata": metadata, "stream_mode": stream_mode, @@ -3453,6 +3458,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3472,6 +3478,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3492,6 +3499,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values", stream_subgraphs: bool = False, metadata: Optional[dict] = None, @@ -3514,6 +3522,7 @@ class SyncRunsClient: assistant_id: The assistant ID or graph name to stream from. If using graph name, will default to first assistant created from that graph. input: The input to the graph. + command: The command to execute. stream_mode: The stream mode(s) to use. stream_subgraphs: Whether to stream output from subgraphs. metadata: Metadata to assign to the run. @@ -3600,6 +3609,7 @@ class SyncRunsClient: """ # noqa: E501 payload = { "input": input, + "command": command, "stream_mode": stream_mode, "stream_subgraphs": stream_subgraphs, "config": config, @@ -3637,6 +3647,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, @@ -3657,6 +3668,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, metadata: Optional[dict] = None, config: Optional[Config] = None, interrupt_before: Optional[Union[All, Sequence[str]]] = None, @@ -3674,6 +3686,7 @@ class SyncRunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, @@ -3695,6 +3708,7 @@ class SyncRunsClient: assistant_id: The assistant ID or graph name to run. If using graph name, will default to first assistant created from that graph. input: The input to the graph. + command: The command to execute. metadata: Metadata to assign to the run. config: The configuration for the assistant. checkpoint: The checkpoint to resume from. @@ -3761,6 +3775,7 @@ class SyncRunsClient: """ # noqa: E501 payload = { "input": input, + "command": command, "config": config, "metadata": metadata, "assistant_id": assistant_id, From d36e6ceaaf3c9c9d57ae751bf1ab9e965fec370e Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Dec 2024 16:59:38 -0800 Subject: [PATCH 2/3] Fix --- libs/langgraph/tests/test_pregel_async.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 6c85e57ed..378ca1ac1 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -12862,6 +12862,10 @@ async def test_parent_command(checkpointer_name: str) -> None: ) +@pytest.mark.skipif( + sys.version_info < (3, 11), + reason="Python 3.11+ is required for async contextvars support", +) @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC) async def test_interrupt_subgraph(checkpointer_name: str): class State(TypedDict): From 0071bd1e1cf91f84a9c364ee3eae8eacd9ad1001 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Mon, 2 Dec 2024 17:01:12 -0800 Subject: [PATCH 3/3] Lint --- libs/langgraph/langgraph/pregel/remote.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index 8ee7ec884..d45cdb310 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/libs/langgraph/langgraph/pregel/remote.py @@ -28,6 +28,7 @@ from langgraph_sdk.client import ( get_sync_client, ) from langgraph_sdk.schema import Checkpoint, ThreadState +from langgraph_sdk.schema import Command as CommandSDK from langgraph_sdk.schema import StreamMode as StreamModeSDK from typing_extensions import Self @@ -599,7 +600,7 @@ class RemoteGraph(PregelProtocol): stream_mode, config ) if isinstance(input, Command): - command: dict[str, Any] = asdict(input) + command: Optional[CommandSDK] = cast(CommandSDK, asdict(input)) input = None else: command = None @@ -688,7 +689,7 @@ class RemoteGraph(PregelProtocol): stream_mode, config ) if isinstance(input, Command): - command: dict[str, Any] = asdict(input) + command: Optional[CommandSDK] = cast(CommandSDK, asdict(input)) input = None else: command = None