From 03785c7d83f764bde65d067c26c5b6dcd4ef50b7 Mon Sep 17 00:00:00 2001
From: Nuno Campos <nuno@boringbits.io>
Date: Sat, 24 Aug 2024 21:34:35 -0700
Subject: [PATCH] sdk: Add on_completion param

---
 libs/sdk-py/langgraph_sdk/client.py | 10 ++++++++++
 libs/sdk-py/langgraph_sdk/schema.py |  2 ++
 2 files changed, 12 insertions(+)

diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py
index 24b2ef3b0..8578680d7 100644
--- a/libs/sdk-py/langgraph_sdk/client.py
+++ b/libs/sdk-py/langgraph_sdk/client.py
@@ -29,6 +29,7 @@ from langgraph_sdk.schema import (
     GraphSchema,
     Metadata,
     MultitaskStrategy,
+    OnCompletionBehavior,
     OnConflictBehavior,
     Run,
     RunCreate,
@@ -985,6 +986,7 @@ class RunsClient:
         feedback_keys: Optional[list[str]] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         webhook: Optional[str] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> AsyncIterator[StreamPart]:
         ...
 
@@ -1004,6 +1006,7 @@ class RunsClient:
         on_disconnect: Optional[DisconnectMode] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> AsyncIterator[StreamPart]:
         """Create a run and stream the results.
 
@@ -1070,6 +1073,7 @@ class RunsClient:
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
             "on_disconnect": on_disconnect,
+            "on_completion": on_completion,
         }
         endpoint = (
             f"/threads/{thread_id}/runs/stream"
@@ -1092,6 +1096,7 @@ class RunsClient:
         interrupt_before: Optional[list[str]] = None,
         interrupt_after: Optional[list[str]] = None,
         webhook: Optional[str] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Run:
         ...
 
@@ -1125,6 +1130,7 @@ class RunsClient:
         interrupt_after: Optional[list[str]] = None,
         webhook: Optional[str] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Run:
         """Create a background run.
 
@@ -1221,6 +1227,7 @@ class RunsClient:
             "webhook": webhook,
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
+            "on_completion": on_completion,
         }
         payload = {k: v for k, v in payload.items() if v is not None}
         if thread_id:
@@ -1268,6 +1275,7 @@ class RunsClient:
         interrupt_after: Optional[list[str]] = None,
         webhook: Optional[str] = None,
         on_disconnect: Optional[DisconnectMode] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Union[list[dict], dict[str, Any]]:
         ...
 
@@ -1285,6 +1293,7 @@ class RunsClient:
         webhook: Optional[str] = None,
         on_disconnect: Optional[DisconnectMode] = None,
         multitask_strategy: Optional[MultitaskStrategy] = None,
+        on_completion: Optional[OnCompletionBehavior] = None,
     ) -> Union[list[dict], dict[str, Any]]:
         """Create a run, wait until it finishes and return the final state.
 
@@ -1364,6 +1373,7 @@ class RunsClient:
             "checkpoint_id": checkpoint_id,
             "multitask_strategy": multitask_strategy,
             "on_disconnect": on_disconnect,
+            "on_completion": on_completion,
         }
         endpoint = (
             f"/threads/{thread_id}/runs/wait" if thread_id is not None else "/runs/wait"
diff --git a/libs/sdk-py/langgraph_sdk/schema.py b/libs/sdk-py/langgraph_sdk/schema.py
index c3232c88e..29fa903a9 100644
--- a/libs/sdk-py/langgraph_sdk/schema.py
+++ b/libs/sdk-py/langgraph_sdk/schema.py
@@ -15,6 +15,8 @@ MultitaskStrategy = Literal["reject", "interrupt", "rollback", "enqueue"]
 
 OnConflictBehavior = Literal["raise", "do_nothing"]
 
+OnCompletionBehavior = Literal["delete", "keep"]
+
 All = Literal["*"]