feat: ✨ support deleting assistant (#258)

* feat: ✨ support deleting assistant * fix: ♻️ assistant route description should be optional --------- Co-authored-by: benshuk <[email protected]>
AI21Labs · Jan 14, 2025 · 0195f62 · 0195f62
1 parent 45d874c
commit 0195f62
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 5 deletions.
diff --git a/ai21/clients/common/beta/assistant/assistants.py b/ai21/clients/common/beta/assistant/assistants.py
@@ -6,7 +6,7 @@
 from ai21.clients.common.beta.assistant.plans import BasePlans
 from ai21.clients.common.beta.assistant.routes import BaseRoutes
 from ai21.models.assistant.assistant import Optimization, Tool, ToolResources
-from ai21.models.responses.assistant_response import AssistantResponse, ListAssistant
+from ai21.models.responses.assistant_response import AssistantResponse, ListAssistant, DeletedAssistantResponse
 from ai21.types import NotGiven, NOT_GIVEN
 from ai21.utils.typing import remove_not_given
 
@@ -75,3 +75,7 @@ def modify(
         tool_resources: ToolResources | NotGiven = NOT_GIVEN,
     ) -> AssistantResponse:
         pass
+
+    @abstractmethod
+    def delete(self, assistant_id: str) -> DeletedAssistantResponse:
+        pass
diff --git a/ai21/clients/common/beta/assistant/routes.py b/ai21/clients/common/beta/assistant/routes.py
@@ -18,8 +18,8 @@ def create(
         assistant_id: str,
         plan_id: str,
         name: str,
-        description: str,
         examples: List[str],
+        description: str | NotGiven = NOT_GIVEN,
         **kwargs,
     ) -> RouteResponse:
         pass

diff --git a/ai21/clients/studio/resources/beta/assistant/assistant.py b/ai21/clients/studio/resources/beta/assistant/assistant.py
@@ -12,7 +12,7 @@
 from ai21.http_client.async_http_client import AsyncAI21HTTPClient
 from ai21.http_client.http_client import AI21HTTPClient
 from ai21.models.assistant.assistant import Tool, ToolResources
-from ai21.models.responses.assistant_response import AssistantResponse, ListAssistant
+from ai21.models.responses.assistant_response import AssistantResponse, ListAssistant, DeletedAssistantResponse
 from ai21.types import NotGiven, NOT_GIVEN
 
 
@@ -76,6 +76,9 @@ def modify(
 
         return self._patch(path=f"/{self._module_name}/{assistant_id}", body=body, response_cls=AssistantResponse)
 
+    def delete(self, assistant_id: str) -> DeletedAssistantResponse:
+        return self._delete(path=f"/{self._module_name}/{assistant_id}", response_cls=DeletedAssistantResponse)
+
 
 class AsyncAssistants(AsyncStudioResource, BaseAssistants):
     def __init__(self, client: AsyncAI21HTTPClient):
@@ -136,3 +139,6 @@ async def modify(
         )
 
         return await self._patch(path=f"/{self._module_name}/{assistant_id}", body=body, response_cls=AssistantResponse)
+
+    async def delete(self, assistant_id: str) -> DeletedAssistantResponse:
+        return await self._delete(path=f"/{self._module_name}/{assistant_id}", response_cls=DeletedAssistantResponse)
diff --git a/ai21/clients/studio/resources/beta/assistant/assistant_routes.py b/ai21/clients/studio/resources/beta/assistant/assistant_routes.py
@@ -18,7 +18,7 @@ def create(
         assistant_id: str,
         plan_id: str,
         name: str,
-        description: str,
+        description: str | NotGiven = NOT_GIVEN,
         examples: List[str],
         **kwargs,
     ) -> RouteResponse:
@@ -85,7 +85,7 @@ async def create(
         assistant_id: str,
         plan_id: str,
         name: str,
-        description: str,
+        description: str | NotGiven = NOT_GIVEN,
         examples: List[str],
         **kwargs,
     ) -> RouteResponse:

diff --git a/ai21/models/responses/assistant_response.py b/ai21/models/responses/assistant_response.py
@@ -24,3 +24,9 @@ class AssistantResponse(AI21BaseModel):
 
 class ListAssistant(AI21BaseModel):
     results: List[AssistantResponse]
+
+
+class DeletedAssistantResponse(AI21BaseModel):
+    object: Literal["assistant"] = "assistant"
+    deleted: bool = True
+    id: str