Merge pull request #231 from AI21Labs/assistant

feat: ✨ add support for Assistant resource

ai21/clients/common/assistant/assistant.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+
+from ai21.models.responses.assistant_response import (
+    AssistantResponse,
+    Optimization,
+    ToolResources,
+    Tool,
+    ListAssistantResponse,
+)
+from ai21.types import NotGiven, NOT_GIVEN
+from ai21.utils.typing import remove_not_given
+
+
+class Assistant(ABC):
+    _module_name = "assistants"
+
+    @abstractmethod
+    def create(
+        self,
+        name: str,
+        *,
+        description: str | NotGiven = NOT_GIVEN,
+        optimization: Optimization | NotGiven = NOT_GIVEN,
+        avatar: str | NotGiven = NOT_GIVEN,
+        models: List[str] | NotGiven = NOT_GIVEN,
+        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tool_resources: ToolResources | NotGiven = NOT_GIVEN,
+        **kwargs,
+    ) -> AssistantResponse:
+        pass
+
+    def _create_body(
+        self,
+        *,
+        name: str,
+        description: str | NotGiven,
+        optimization: str | NotGiven,
+        avatar: str | NotGiven,
+        models: List[str] | NotGiven,
+        tools: List[str] | NotGiven,
+        tool_resources: dict | NotGiven,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        return remove_not_given(
+            {
+                "name": name,
+                "description": description,
+                "optimization": optimization,
+                "avatar": avatar,
+                "models": models,
+                "tools": tools,
+                "tool_resources": tool_resources,
+                **kwargs,
+            }
+        )
+
+    @abstractmethod
+    def list(self) -> ListAssistantResponse:
+        pass
+
+    @abstractmethod
+    def get(self, assistant_id: str) -> AssistantResponse:
+        pass
+
+    @abstractmethod
+    def modify(
+        self,
+        assistant_id: str,
+        *,
+        name: str | NotGiven = NOT_GIVEN,
+        description: str | NotGiven = NOT_GIVEN,
+        optimization: Optimization | NotGiven = NOT_GIVEN,
+        avatar: str | NotGiven = NOT_GIVEN,
+        is_archived: bool | NotGiven = NOT_GIVEN,
+        models: List[str] | NotGiven = NOT_GIVEN,
+        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tool_resources: ToolResources | NotGiven = NOT_GIVEN,
+    ) -> AssistantResponse:
+        pass

ai21/clients/studio/resources/assistant/studio_assistant.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from typing import List
+
+from ai21.clients.common.assistant.assistant import Assistant
+from ai21.clients.studio.resources.studio_resource import (
+    AsyncStudioResource,
+    StudioResource,
+)
+from ai21.models.responses.assistant_response import (
+    AssistantResponse,
+    Tool,
+    ToolResources,
+    ListAssistantResponse,
+)
+from ai21.types import NotGiven, NOT_GIVEN
+
+
+class StudioAssistant(StudioResource, Assistant):
+    def create(
+        self,
+        name: str,
+        *,
+        description: str | NotGiven = NOT_GIVEN,
+        optimization: str | NotGiven = NOT_GIVEN,
+        avatar: str | NotGiven = NOT_GIVEN,
+        models: List[str] | NotGiven = NOT_GIVEN,
+        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tool_resources: ToolResources | NotGiven = NOT_GIVEN,
+        **kwargs,
+    ) -> AssistantResponse:
+        body = self._create_body(
+            name=name,
+            description=description,
+            optimization=optimization,
+            avatar=avatar,
+            models=models,
+            tools=tools,
+            tool_resources=tool_resources,
+            **kwargs,
+        )
+
+        return self._post(path=f"/{self._module_name}", body=body, response_cls=AssistantResponse)
+
+    def get(self, assistant_id: str) -> AssistantResponse:
+        return self._get(path=f"/{self._module_name}/{assistant_id}", response_cls=AssistantResponse)
+
+    def list(self) -> ListAssistantResponse:
+        return self._get(path=f"/{self._module_name}", response_cls=ListAssistantResponse)
+
+    def modify(
+        self,
+        assistant_id: str,
+        *,
+        name: str | NotGiven = NOT_GIVEN,
+        description: str | NotGiven = NOT_GIVEN,
+        optimization: str | NotGiven = NOT_GIVEN,
+        avatar: str | NotGiven = NOT_GIVEN,
+        is_archived: bool | NotGiven = NOT_GIVEN,
+        models: List[str] | NotGiven = NOT_GIVEN,
+        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tool_resources: ToolResources | NotGiven = NOT_GIVEN,
+    ) -> AssistantResponse:
+        body = self._create_body(
+            name=name,
+            description=description,
+            optimization=optimization,
+            avatar=avatar,
+            is_archived=is_archived,
+            models=models,
+            tools=tools,
+            tool_resources=tool_resources,
+        )
+
+        return self._patch(path=f"/{self._module_name}/{assistant_id}", body=body, response_cls=AssistantResponse)
+
+
+class AsyncStudioAssistant(AsyncStudioResource, Assistant):
+    async def create(
+        self,
+        name: str,
+        *,
+        description: str | NotGiven = NOT_GIVEN,
+        optimization: str | NotGiven = NOT_GIVEN,
+        avatar: str | NotGiven = NOT_GIVEN,
+        models: List[str] | NotGiven = NOT_GIVEN,
+        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tool_resources: ToolResources | NotGiven = NOT_GIVEN,
+        **kwargs,
+    ) -> AssistantResponse:
+        body = self._create_body(
+            name=name,
+            description=description,
+            optimization=optimization,
+            avatar=avatar,
+            models=models,
+            tools=tools,
+            tool_resources=tool_resources,
+            **kwargs,
+        )
+
+        return self._post(path=f"/{self._module_name}", body=body, response_cls=AssistantResponse)
+
+    async def get(self, assistant_id: str) -> AssistantResponse:
+        return await self._get(path=f"/{self._module_name}/{assistant_id}", response_cls=AssistantResponse)
+
+    async def list(self) -> ListAssistantResponse:
+        return await self._get(path=f"/{self._module_name}", response_cls=ListAssistantResponse)
+
+    async def modify(
+        self,
+        assistant_id: str,
+        *,
+        name: str | NotGiven = NOT_GIVEN,
+        description: str | NotGiven = NOT_GIVEN,
+        optimization: str | NotGiven = NOT_GIVEN,
+        avatar: str | NotGiven = NOT_GIVEN,
+        is_archived: bool | NotGiven = NOT_GIVEN,
+        models: List[str] | NotGiven = NOT_GIVEN,
+        tools: List[Tool] | NotGiven = NOT_GIVEN,
+        tool_resources: ToolResources | NotGiven = NOT_GIVEN,
+    ) -> AssistantResponse:
+        body = self._create_body(
+            name=name,
+            description=description,
+            optimization=optimization,
+            avatar=avatar,
+            is_archived=is_archived,
+            models=models,
+            tools=tools,
+            tool_resources=tool_resources,
+        )
+
+        return await self._patch(path=f"/{self._module_name}/{assistant_id}", body=body, response_cls=AssistantResponse)

ai21/clients/studio/resources/beta/async_beta.py

@@ -1,3 +1,4 @@
+from ai21.clients.studio.resources.assistant.studio_assistant import AsyncStudioAssistant
 from ai21.clients.studio.resources.studio_conversational_rag import AsyncStudioConversationalRag
 from ai21.clients.studio.resources.studio_resource import AsyncStudioResource
 from ai21.http_client.async_http_client import AsyncAI21HTTPClient
@@ -8,3 +9,4 @@ def __init__(self, client: AsyncAI21HTTPClient):
         super().__init__(client)
 
         self.conversational_rag = AsyncStudioConversationalRag(client)
+        self.assistants = AsyncStudioAssistant(client)

ai21/clients/studio/resources/beta/beta.py

@@ -1,3 +1,4 @@
+from ai21.clients.studio.resources.assistant.studio_assistant import StudioAssistant
 from ai21.clients.studio.resources.studio_conversational_rag import StudioConversationalRag
 from ai21.clients.studio.resources.studio_resource import StudioResource
 from ai21.http_client.http_client import AI21HTTPClient
@@ -8,3 +9,4 @@ def __init__(self, client: AI21HTTPClient):
         super().__init__(client)
 
         self.conversational_rag = StudioConversationalRag(client)
+        self.assistants = StudioAssistant(client)

ai21/clients/studio/resources/studio_resource.py

@@ -80,6 +80,12 @@ def _get(
         response = self._client.execute_http_request(method="GET", path=path, params=params or {})
         return _cast_response(response=response, response_cls=response_cls)
 
+    def _patch(
+        self, path: str, response_cls: Optional[ResponseT] = None, body: Dict[str, Any] = None
+    ) -> ResponseT | StreamT:
+        response = self._client.execute_http_request(method="PATCH", path=path, body=body or {})
+        return _cast_response(response=response, response_cls=response_cls)
+
     def _put(
         self, path: str, response_cls: Optional[ResponseT] = None, body: Dict[str, Any] = None
     ) -> ResponseT | StreamT:
@@ -131,6 +137,12 @@ async def _get(
         response = await self._client.execute_http_request(method="GET", path=path, params=params or {})
         return _cast_response(response=response, response_cls=response_cls)
 
+    async def _patch(
+        self, path: str, response_cls: Optional[ResponseT] = None, body: Dict[str, Any] = None
+    ) -> ResponseT | AsyncStreamT:
+        response = await self._client.execute_http_request(method="PATCH", path=path, body=body or {})
+        return _cast_response(response=response, response_cls=response_cls)
+
     async def _put(
         self, path: str, response_cls: Optional[ResponseT] = None, body: Dict[str, Any] = None
     ) -> ResponseT | AsyncStreamT:

ai21/models/responses/assistant_response.py

@@ -0,0 +1,37 @@
+from datetime import datetime
+from typing import Optional, List, Literal
+
+from typing_extensions import TypedDict
+
+from ai21.models.ai21_base_model import AI21BaseModel
+
+
+Optimization = Literal["cost", "latency"]
+Tool = Literal["rag", "internet_research", "plan_approval"]
+
+
+class ToolResources(TypedDict, total=False):
+    rag: Optional[dict]
+    internet_research: Optional[dict]
+    plan_approval: Optional[dict]
+
+
+class AssistantResponse(AI21BaseModel):
+    id: str
+    created_at: datetime
+    updated_at: datetime
+    object: str
+    name: str
+    description: Optional[str] = None
+    optimization: str
+    organization_id: str
+    user_id: str
+    avatar: Optional[str] = None
+    is_archived: bool = False
+    models: Optional[List[str]] = None
+    tools: Optional[List[str]] = None
+    tool_resources: Optional[ToolResources] = None
+
+
+class ListAssistantResponse(AI21BaseModel):
+    results: List[AssistantResponse]