feat: ✨ add support for Plan and Route resources

ai21/clients/common/beta/assistant/assistants.py

@@ -3,14 +3,18 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List
 
+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.types import NotGiven, NOT_GIVEN
 from ai21.utils.typing import remove_not_given
 
 
-class Assistants(ABC):
+class BaseAssistants(ABC):
     _module_name = "assistants"
+    plans: BasePlans
+    routes: BaseRoutes
 
     @abstractmethod
     def create(

ai21/clients/common/beta/assistant/messages.py

@@ -6,7 +6,7 @@
 from ai21.models.responses.message_response import MessageResponse, ListMessageResponse
 
 
-class Messages(ABC):
+class BaseMessages(ABC):
     _module_name = "messages"
 
     @abstractmethod

ai21/clients/common/beta/assistant/plans.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict
+
+from ai21.models.responses.plan_response import PlanResponse, ListPlanResponse
+from ai21.utils.typing import remove_not_given
+
+
+class BasePlans(ABC):
+    _module_name = "plans"
+
+    @abstractmethod
+    def create(
+        self,
+        *,
+        assistant_id: str,
+        code: str,
+        **kwargs,
+    ) -> PlanResponse:
+        pass
+
+    def _create_body(
+        self,
+        *,
+        code: str,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        return remove_not_given(
+            {
+                "code": code,
+                **kwargs,
+            }
+        )
+
+    @abstractmethod
+    def list(
+        self,
+        *,
+        assistant_id: str,
+    ) -> ListPlanResponse:
+        pass
+
+    @abstractmethod
+    def retrieve(
+        self,
+        *,
+        assistant_id: str,
+        plan_id: str,
+    ) -> PlanResponse:
+        pass
+
+    @abstractmethod
+    def modify(
+        self,
+        *,
+        assistant_id: str,
+        plan_id: str,
+        code: str,
+    ) -> PlanResponse:
+        pass

ai21/clients/common/beta/assistant/routes.py

@@ -0,0 +1,77 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+
+from ai21.models.responses.route_response import RouteResponse, ListRouteResponse
+from ai21.types import NotGiven, NOT_GIVEN
+from ai21.utils.typing import remove_not_given
+
+
+class BaseRoutes(ABC):
+    _module_name = "routes"
+
+    @abstractmethod
+    def create(
+        self,
+        *,
+        assistant_id: str,
+        plan_id: str,
+        name: str,
+        description: str,
+        examples: List[str],
+        **kwargs,
+    ) -> RouteResponse:
+        pass
+
+    def _create_body(
+        self,
+        *,
+        plan_id: str | NotGiven = NOT_GIVEN,
+        name: str | NotGiven = NOT_GIVEN,
+        description: str | NotGiven = NOT_GIVEN,
+        examples: List[str] | NotGiven = NOT_GIVEN,
+        **kwargs,
+    ) -> Dict[str, Any]:
+        return remove_not_given(
+            {
+                "plan_id": plan_id,
+                "name": name,
+                "description": description,
+                "examples": examples,
+                **kwargs,
+            }
+        )
+
+    @abstractmethod
+    def retrieve(
+        self,
+        *,
+        assistant_id: str,
+        route_id: str,
+    ) -> RouteResponse:
+        pass
+
+    @abstractmethod
+    def list(
+        self,
+        *,
+        assistant_id: str,
+        name: str | NotGiven = NotGiven,
+    ) -> ListRouteResponse:
+        pass
+
+    @abstractmethod
+    def modify(
+        self,
+        *,
+        assistant_id: str,
+        route_id: str,
+        description: str | NotGiven = NOT_GIVEN,
+        examples: List[str] | NotGiven = NOT_GIVEN,
+    ) -> RouteResponse:
+        pass
+
+    @abstractmethod
+    def delete(self, *, assistant_id: str, route_id: str):
+        pass

ai21/clients/common/beta/assistant/runs.py

@@ -10,7 +10,7 @@
 from ai21.utils.typing import remove_not_given
 
 
-class Runs(ABC):
+class BaseRuns(ABC):
     _module_name = "runs"
 
     @abstractmethod

ai21/clients/common/beta/assistant/threads.py

@@ -3,14 +3,14 @@
 from abc import ABC, abstractmethod
 from typing import List
 
-from ai21.clients.common.beta.assistant.messages import Messages
+from ai21.clients.common.beta.assistant.messages import BaseMessages
 from ai21.models.assistant.message import Message
 from ai21.models.responses.thread_response import ThreadResponse
 
 
-class Threads(ABC):
+class BaseThreads(ABC):
     _module_name = "threads"
-    messages: Messages
+    messages: BaseMessages
 
     @abstractmethod
     def create(

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

@@ -2,17 +2,27 @@
 
 from typing import List
 
-from ai21.clients.common.beta.assistant.assistants import Assistants
+from ai21.clients.common.beta.assistant.assistants import BaseAssistants
+from ai21.clients.studio.resources.beta.assistant.assistants_plans import AssistantPlans, AsyncAssistantPlans
+from ai21.clients.studio.resources.beta.assistant.assistant_routes import AssistantRoutes, AsyncAssistantRoutes
 from ai21.clients.studio.resources.studio_resource import (
     AsyncStudioResource,
     StudioResource,
 )
+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.types import NotGiven, NOT_GIVEN
 
 
-class Assistant(StudioResource, Assistants):
+class Assistants(StudioResource, BaseAssistants):
+    def __init__(self, client: AI21HTTPClient):
+        super().__init__(client)
+
+        self.plans = AssistantPlans(client)
+        self.routes = AssistantRoutes(client)
+
     def create(
         self,
         name: str,
@@ -67,7 +77,13 @@ def modify(
         return self._patch(path=f"/{self._module_name}/{assistant_id}", body=body, response_cls=AssistantResponse)
 
 
-class AsyncAssistant(AsyncStudioResource, Assistants):
+class AsyncAssistants(AsyncStudioResource, BaseAssistants):
+    def __init__(self, client: AsyncAI21HTTPClient):
+        super().__init__(client)
+
+        self.plans = AsyncAssistantPlans(client)
+        self.routes = AsyncAssistantRoutes(client)
+
     async def create(
         self,
         name: str,

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

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+from typing import List
+
+from ai21.clients.common.beta.assistant.routes import BaseRoutes
+from ai21.clients.studio.resources.studio_resource import (
+    AsyncStudioResource,
+    StudioResource,
+)
+from ai21.models.responses.route_response import RouteResponse, ListRouteResponse
+from ai21.types import NotGiven
+
+
+class AssistantRoutes(StudioResource, BaseRoutes):
+    def create(
+        self,
+        *,
+        assistant_id: str,
+        plan_id: str,
+        name: str,
+        description: str,
+        examples: List[str],
+        **kwargs,
+    ) -> RouteResponse:
+        body = self._create_body(
+            plan_id=plan_id,
+            name=name,
+            description=description,
+            examples=examples,
+            **kwargs,
+        )
+
+        return self._post(path=f"/assistants/{assistant_id}/{self._module_name}", body=body, response_cls=RouteResponse)
+
+    def list(
+        self,
+        *,
+        assistant_id: str,
+        name: str | NotGiven = NotGiven,
+    ) -> ListRouteResponse:
+        params = self._create_body(name=name)
+
+        return self._get(
+            path=f"/assistants/{assistant_id}/{self._module_name}", params=params, response_cls=ListRouteResponse
+        )
+
+    def retrieve(
+        self,
+        *,
+        assistant_id: str,
+        route_id: str,
+    ) -> RouteResponse:
+        return self._get(path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}", response_cls=RouteResponse)
+
+    def modify(
+        self,
+        *,
+        assistant_id: str,
+        route_id: str,
+        description: str | NotGiven = NotGiven,
+        examples: List[str] | NotGiven = NotGiven,
+    ) -> RouteResponse:
+        body = self._create_body(
+            description=description,
+            examples=examples,
+        )
+
+        return self._patch(
+            path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}", body=body, response_cls=RouteResponse
+        )
+
+    def delete(
+        self,
+        *,
+        assistant_id: str,
+        route_id: str,
+    ):
+        return self._delete(path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}")
+
+
+class AsyncAssistantRoutes(AsyncStudioResource, BaseRoutes):
+    async def create(
+        self,
+        *,
+        assistant_id: str,
+        plan_id: str,
+        name: str,
+        description: str,
+        examples: List[str],
+        **kwargs,
+    ) -> RouteResponse:
+        body = self._create_body(
+            plan_id=plan_id,
+            name=name,
+            description=description,
+            examples=examples,
+            **kwargs,
+        )
+
+        return await self._post(
+            path=f"/assistants/{assistant_id}/{self._module_name}", body=body, response_cls=RouteResponse
+        )
+
+    async def list(
+        self,
+        *,
+        assistant_id: str,
+        name: str | NotGiven = NotGiven,
+    ) -> ListRouteResponse:
+        params = self._create_body(name=name)
+
+        return await self._get(
+            path=f"/assistants/{assistant_id}/{self._module_name}", params=params, response_cls=ListRouteResponse
+        )
+
+    async def retrieve(
+        self,
+        *,
+        assistant_id: str,
+        route_id: str,
+    ) -> RouteResponse:
+        return await self._get(
+            path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}", response_cls=RouteResponse
+        )
+
+    async def modify(
+        self,
+        *,
+        assistant_id: str,
+        route_id: str,
+        description: str | NotGiven = NotGiven,
+        examples: List[str] | NotGiven = NotGiven,
+    ) -> RouteResponse:
+        body = self._create_body(
+            description=description,
+            examples=examples,
+        )
+
+        return await self._patch(
+            path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}", body=body, response_cls=RouteResponse
+        )
+
+    async def delete(
+        self,
+        *,
+        assistant_id: str,
+        route_id: str,
+    ):
+        return await self._delete(path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}")