Merge pull request #2597 from langchain-ai/nc/2dec/remote-command

lib: Handle Command in RemoteGraph

libs/langgraph/langgraph/pregel/remote.py

@@ -1,3 +1,4 @@
+from dataclasses import asdict
 from typing import (
     Any,
     AsyncIterator,
@@ -27,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
 
@@ -41,7 +43,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 +599,17 @@ class RemoteGraph(PregelProtocol):
         stream_modes, requested, req_single, stream = self._get_stream_modes(
             stream_mode, config
         )
+        if isinstance(input, Command):
+            command: Optional[CommandSDK] = cast(CommandSDK, 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 +688,17 @@ class RemoteGraph(PregelProtocol):
         stream_modes, requested, req_single, stream = self._get_stream_modes(
             stream_mode, config
         )
+        if isinstance(input, Command):
+            command: Optional[CommandSDK] = cast(CommandSDK, 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,

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):

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,