From 229a9e19a89aa7ee1c0ca946ded1ff296b334ee9 Mon Sep 17 00:00:00 2001 From: Nuno Campos Date: Wed, 13 Nov 2024 17:06:56 -0800 Subject: [PATCH] sdk-py: Add command arg for creating runs --- libs/langgraph/langgraph/pregel/io.py | 10 +--------- libs/sdk-py/langgraph_sdk/client.py | 16 ++++++++++++++++ libs/sdk-py/langgraph_sdk/schema.py | 11 +++++++++++ 3 files changed, 28 insertions(+), 9 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/io.py b/libs/langgraph/langgraph/pregel/io.py index 7c26731f1..6695e1ce0 100644 --- a/libs/langgraph/langgraph/pregel/io.py +++ b/libs/langgraph/langgraph/pregel/io.py @@ -69,19 +69,11 @@ def map_command( ) -> Iterator[tuple[str, str, Any]]: """Map input chunk to a sequence of pending writes in the form (channel, value).""" if cmd.send: - if isinstance(cmd.send, (tuple, list)) and all( - isinstance(x, Send) - or isinstance(x, (list, tuple)) - and len(x) == 2 - and isinstance(x[0], str) - for x in cmd.send - ): + if isinstance(cmd.send, (tuple, list)): sends = cmd.send else: sends = [cmd.send] for send in sends: - if isinstance(send, tuple) and len(send) == 2 and isinstance(send[0], str): - send = Send(*send) if not isinstance(send, Send): raise TypeError( f"In Command.send, expected Send, got {type(send).__name__}" diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index a057d452f..6a3bb6c9c 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -36,6 +36,7 @@ from langgraph_sdk.schema import ( AssistantVersion, CancelAction, Checkpoint, + Command, Config, Cron, DisconnectMode, @@ -1174,6 +1175,7 @@ class RunsClient: 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, @@ -1197,6 +1199,7 @@ class RunsClient: 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, @@ -1217,6 +1220,7 @@ class RunsClient: 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, @@ -1241,6 +1245,7 @@ class RunsClient: 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: A command to execute. Cannot be combined with input. stream_mode: The stream mode(s) to use. stream_subgraphs: Whether to stream output from subgraphs. metadata: Metadata to assign to the run. @@ -1291,6 +1296,7 @@ class RunsClient: """ # noqa: E501 payload = { "input": input, + "command": command, "config": config, "metadata": metadata, "stream_mode": stream_mode, @@ -1324,6 +1330,7 @@ class RunsClient: 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, @@ -1343,6 +1350,7 @@ class RunsClient: 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, @@ -1363,6 +1371,7 @@ class RunsClient: 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, @@ -1385,6 +1394,7 @@ class RunsClient: 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: A command to execute. Cannot be combined with input. stream_mode: The stream mode(s) to use. stream_subgraphs: Whether to stream output from subgraphs. metadata: Metadata to assign to the run. @@ -1471,6 +1481,7 @@ class RunsClient: """ # noqa: E501 payload = { "input": input, + "command": command, "stream_mode": stream_mode, "stream_subgraphs": stream_subgraphs, "config": config, @@ -1508,6 +1519,7 @@ class RunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, @@ -1529,6 +1541,7 @@ class RunsClient: 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, @@ -1547,6 +1560,7 @@ class RunsClient: assistant_id: str, *, input: Optional[dict] = None, + command: Optional[Command] = None, metadata: Optional[dict] = None, config: Optional[Config] = None, checkpoint: Optional[Checkpoint] = None, @@ -1569,6 +1583,7 @@ class RunsClient: 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: A command to execute. Cannot be combined with input. metadata: Metadata to assign to the run. config: The configuration for the assistant. checkpoint: The checkpoint to resume from. @@ -1635,6 +1650,7 @@ class RunsClient: """ # noqa: E501 payload = { "input": input, + "command": command, "config": config, "metadata": metadata, "assistant_id": assistant_id, diff --git a/libs/sdk-py/langgraph_sdk/schema.py b/libs/sdk-py/langgraph_sdk/schema.py index 9583a1c1b..5264ce709 100644 --- a/libs/sdk-py/langgraph_sdk/schema.py +++ b/libs/sdk-py/langgraph_sdk/schema.py @@ -339,3 +339,14 @@ class StreamPart(NamedTuple): """The type of event for this stream part.""" data: dict """The data payload associated with the event.""" + + +class Send(TypedDict): + node: str + input: Optional[dict[str, Any]] + + +class Command(TypedDict, total=False): + send: Union[Send, Sequence[Send]] + update: dict[str, Any] + resume: Any