sdk: Add action query param to cancel Run methods (#2284)

This commit is contained in:
Andrew Nguonly
2024-11-01 13:32:49 -07:00
committed by GitHub
parent 509291fc1a
commit ecb584cbe0
4 changed files with 40 additions and 6 deletions
+4
View File
@@ -1,6 +1,7 @@
import {
Assistant,
AssistantGraph,
CancelAction,
Config,
DefaultValues,
GraphSchema,
@@ -928,17 +929,20 @@ export class RunsClient extends BaseClient {
* @param threadId The ID of the thread.
* @param runId The ID of the run.
* @param wait Whether to block when canceling
* @param action Action to take when cancelling the run. Possible values are `interrupt` or `rollback`. Default is `interrupt`.
* @returns
*/
async cancel(
threadId: string,
runId: string,
wait: boolean = false,
action: CancelAction = "interrupt",
): Promise<void> {
return this.fetch<void>(`/threads/${threadId}/runs/${runId}/cancel`, {
method: "POST",
params: {
wait: wait ? "1" : "0",
action: action,
},
});
}
+2
View File
@@ -14,6 +14,8 @@ type ThreadStatus = "idle" | "busy" | "interrupted" | "error";
type MultitaskStrategy = "reject" | "interrupt" | "rollback" | "enqueue";
export type CancelAction = "interrupt" | "rollback";
export interface Config {
/**
* Tags for this call and any sub-calls (eg. a Chain calling an LLM).
+27 -6
View File
@@ -34,6 +34,7 @@ from langgraph_sdk.schema import (
All,
Assistant,
AssistantVersion,
CancelAction,
Checkpoint,
Config,
Cron,
@@ -1712,13 +1713,22 @@ class RunsClient:
return await self.http.get(f"/threads/{thread_id}/runs/{run_id}")
async def cancel(self, thread_id: str, run_id: str, *, wait: bool = False) -> None:
async def cancel(
self,
thread_id: str,
run_id: str,
*,
wait: bool = False,
action: CancelAction = "interrupt",
) -> None:
"""Get a run.
Args:
thread_id: The thread ID to cancel.
run_id: The run ID to cancek.
wait: Whether to wait until run has completed.
action: Action to take when cancelling the run. Possible values
are `interrupt` or `rollback`. Default is `interrupt`.
Returns:
None
@@ -1728,12 +1738,13 @@ class RunsClient:
await client.runs.cancel(
thread_id="thread_id_to_cancel",
run_id="run_id_to_cancel",
wait=True
wait=True,
action="interrupt"
)
""" # noqa: E501
return await self.http.post(
f"/threads/{thread_id}/runs/{run_id}/cancel?wait={1 if wait else 0}",
f"/threads/{thread_id}/runs/{run_id}/cancel?wait={1 if wait else 0}&action={action}",
json=None,
)
@@ -3792,13 +3803,22 @@ class SyncRunsClient:
return self.http.get(f"/threads/{thread_id}/runs/{run_id}")
def cancel(self, thread_id: str, run_id: str, *, wait: bool = False) -> None:
def cancel(
self,
thread_id: str,
run_id: str,
*,
wait: bool = False,
action: CancelAction = "interrupt",
) -> None:
"""Get a run.
Args:
thread_id: The thread ID to cancel.
run_id: The run ID to cancek.
wait: Whether to wait until run has completed.
action: Action to take when cancelling the run. Possible values
are `interrupt` or `rollback`. Default is `interrupt`.
Returns:
None
@@ -3808,12 +3828,13 @@ class SyncRunsClient:
client.runs.cancel(
thread_id="thread_id_to_cancel",
run_id="run_id_to_cancel",
wait=True
wait=True,
action="interrupt"
)
""" # noqa: E501
return self.http.post(
f"/threads/{thread_id}/runs/{run_id}/cancel?wait={1 if wait else 0}",
f"/threads/{thread_id}/runs/{run_id}/cancel?wait={1 if wait else 0}&action={action}",
json=None,
)
+7
View File
@@ -78,6 +78,13 @@ Specifies behavior if the thread doesn't exist:
- "reject": Reject the operation if the thread doesn't exist.
"""
CancelAction = Literal["interrupt", "rollback"]
"""
Action to take when cancelling the run.
- "interrupt": Simply cancel the run.
- "rollback": Cancel the run. Then delete the run and associated checkpoints.
"""
class Config(TypedDict, total=False):
"""Configuration options for a call."""