Skip to content

Commit f1f5cbc

Browse files
benshukbenshuk
benshuk
and
benshuk
authored
feat: ✨ add support for Plan and Route resources (#238)
* feat: ✨ add support for Plan resource * feat: ✨ add support for Route resource * fix: 🏷️ List instead of list * refactor: 🚚 rename classes and files --------- Co-authored-by: benshuk <bens@ai21.com>
1 parent d4beb22 commit f1f5cbc

File tree

16 files changed

+478
-32
lines changed

16 files changed

+478
-32
lines changed

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,18 @@
33
from abc import ABC, abstractmethod
44
from typing import Any, Dict, List
55

6+
from ai21.clients.common.beta.assistant.plans import BasePlans
7+
from ai21.clients.common.beta.assistant.routes import BaseRoutes
68
from ai21.models.assistant.assistant import Optimization, Tool, ToolResources
79
from ai21.models.responses.assistant_response import AssistantResponse, ListAssistant
810
from ai21.types import NotGiven, NOT_GIVEN
911
from ai21.utils.typing import remove_not_given
1012

1113

12-
class Assistants(ABC):
14+
class BaseAssistants(ABC):
1315
_module_name = "assistants"
16+
plans: BasePlans
17+
routes: BaseRoutes
1418

1519
@abstractmethod
1620
def create(

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
from ai21.models.responses.message_response import MessageResponse, ListMessageResponse
77

88

9-
class Messages(ABC):
9+
class BaseMessages(ABC):
1010
_module_name = "messages"
1111

1212
@abstractmethod

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

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
1+
from __future__ import annotations
2+
3+
from abc import ABC, abstractmethod
4+
from typing import Any, Dict
5+
6+
from ai21.models.responses.plan_response import PlanResponse, ListPlanResponse
7+
from ai21.utils.typing import remove_not_given
8+
9+
10+
class BasePlans(ABC):
11+
_module_name = "plans"
12+
13+
@abstractmethod
14+
def create(
15+
self,
16+
*,
17+
assistant_id: str,
18+
code: str,
19+
**kwargs,
20+
) -> PlanResponse:
21+
pass
22+
23+
def _create_body(
24+
self,
25+
*,
26+
code: str,
27+
**kwargs,
28+
) -> Dict[str, Any]:
29+
return remove_not_given(
30+
{
31+
"code": code,
32+
**kwargs,
33+
}
34+
)
35+
36+
@abstractmethod
37+
def list(
38+
self,
39+
*,
40+
assistant_id: str,
41+
) -> ListPlanResponse:
42+
pass
43+
44+
@abstractmethod
45+
def retrieve(
46+
self,
47+
*,
48+
assistant_id: str,
49+
plan_id: str,
50+
) -> PlanResponse:
51+
pass
52+
53+
@abstractmethod
54+
def modify(
55+
self,
56+
*,
57+
assistant_id: str,
58+
plan_id: str,
59+
code: str,
60+
) -> PlanResponse:
61+
pass

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

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
from __future__ import annotations
2+
3+
from abc import ABC, abstractmethod
4+
from typing import Any, Dict, List
5+
6+
from ai21.models.responses.route_response import RouteResponse, ListRouteResponse
7+
from ai21.types import NotGiven, NOT_GIVEN
8+
from ai21.utils.typing import remove_not_given
9+
10+
11+
class BaseRoutes(ABC):
12+
_module_name = "routes"
13+
14+
@abstractmethod
15+
def create(
16+
self,
17+
*,
18+
assistant_id: str,
19+
plan_id: str,
20+
name: str,
21+
description: str,
22+
examples: List[str],
23+
**kwargs,
24+
) -> RouteResponse:
25+
pass
26+
27+
def _create_body(
28+
self,
29+
*,
30+
plan_id: str | NotGiven = NOT_GIVEN,
31+
name: str | NotGiven = NOT_GIVEN,
32+
description: str | NotGiven = NOT_GIVEN,
33+
examples: List[str] | NotGiven = NOT_GIVEN,
34+
**kwargs,
35+
) -> Dict[str, Any]:
36+
return remove_not_given(
37+
{
38+
"plan_id": plan_id,
39+
"name": name,
40+
"description": description,
41+
"examples": examples,
42+
**kwargs,
43+
}
44+
)
45+
46+
@abstractmethod
47+
def retrieve(
48+
self,
49+
*,
50+
assistant_id: str,
51+
route_id: str,
52+
) -> RouteResponse:
53+
pass
54+
55+
@abstractmethod
56+
def list(
57+
self,
58+
*,
59+
assistant_id: str,
60+
name: str | NotGiven = NotGiven,
61+
) -> ListRouteResponse:
62+
pass
63+
64+
@abstractmethod
65+
def modify(
66+
self,
67+
*,
68+
assistant_id: str,
69+
route_id: str,
70+
description: str | NotGiven = NOT_GIVEN,
71+
examples: List[str] | NotGiven = NOT_GIVEN,
72+
) -> RouteResponse:
73+
pass
74+
75+
@abstractmethod
76+
def delete(self, *, assistant_id: str, route_id: str):
77+
pass

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
from ai21.utils.typing import remove_not_given
1111

1212

13-
class Runs(ABC):
13+
class BaseRuns(ABC):
1414
_module_name = "runs"
1515

1616
@abstractmethod

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,14 +3,14 @@
33
from abc import ABC, abstractmethod
44
from typing import List
55

6-
from ai21.clients.common.beta.assistant.messages import Messages
6+
from ai21.clients.common.beta.assistant.messages import BaseMessages
77
from ai21.models.assistant.message import Message
88
from ai21.models.responses.thread_response import ThreadResponse
99

1010

11-
class Threads(ABC):
11+
class BaseThreads(ABC):
1212
_module_name = "threads"
13-
messages: Messages
13+
messages: BaseMessages
1414

1515
@abstractmethod
1616
def create(

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

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -2,17 +2,27 @@
22

33
from typing import List
44

5-
from ai21.clients.common.beta.assistant.assistants import Assistants
5+
from ai21.clients.common.beta.assistant.assistants import BaseAssistants
6+
from ai21.clients.studio.resources.beta.assistant.assistants_plans import AssistantPlans, AsyncAssistantPlans
7+
from ai21.clients.studio.resources.beta.assistant.assistant_routes import AssistantRoutes, AsyncAssistantRoutes
68
from ai21.clients.studio.resources.studio_resource import (
79
AsyncStudioResource,
810
StudioResource,
911
)
12+
from ai21.http_client.async_http_client import AsyncAI21HTTPClient
13+
from ai21.http_client.http_client import AI21HTTPClient
1014
from ai21.models.assistant.assistant import Tool, ToolResources
1115
from ai21.models.responses.assistant_response import AssistantResponse, ListAssistant
1216
from ai21.types import NotGiven, NOT_GIVEN
1317

1418

15-
class Assistant(StudioResource, Assistants):
19+
class Assistants(StudioResource, BaseAssistants):
20+
def __init__(self, client: AI21HTTPClient):
21+
super().__init__(client)
22+
23+
self.plans = AssistantPlans(client)
24+
self.routes = AssistantRoutes(client)
25+
1626
def create(
1727
self,
1828
name: str,
@@ -67,7 +77,13 @@ def modify(
6777
return self._patch(path=f"/{self._module_name}/{assistant_id}", body=body, response_cls=AssistantResponse)
6878

6979

70-
class AsyncAssistant(AsyncStudioResource, Assistants):
80+
class AsyncAssistants(AsyncStudioResource, BaseAssistants):
81+
def __init__(self, client: AsyncAI21HTTPClient):
82+
super().__init__(client)
83+
84+
self.plans = AsyncAssistantPlans(client)
85+
self.routes = AsyncAssistantRoutes(client)
86+
7187
async def create(
7288
self,
7389
name: str,

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

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
from __future__ import annotations
2+
3+
from typing import List
4+
5+
from ai21.clients.common.beta.assistant.routes import BaseRoutes
6+
from ai21.clients.studio.resources.studio_resource import (
7+
AsyncStudioResource,
8+
StudioResource,
9+
)
10+
from ai21.models.responses.route_response import RouteResponse, ListRouteResponse
11+
from ai21.types import NotGiven
12+
13+
14+
class AssistantRoutes(StudioResource, BaseRoutes):
15+
def create(
16+
self,
17+
*,
18+
assistant_id: str,
19+
plan_id: str,
20+
name: str,
21+
description: str,
22+
examples: List[str],
23+
**kwargs,
24+
) -> RouteResponse:
25+
body = self._create_body(
26+
plan_id=plan_id,
27+
name=name,
28+
description=description,
29+
examples=examples,
30+
**kwargs,
31+
)
32+
33+
return self._post(path=f"/assistants/{assistant_id}/{self._module_name}", body=body, response_cls=RouteResponse)
34+
35+
def list(
36+
self,
37+
*,
38+
assistant_id: str,
39+
name: str | NotGiven = NotGiven,
40+
) -> ListRouteResponse:
41+
params = self._create_body(name=name)
42+
43+
return self._get(
44+
path=f"/assistants/{assistant_id}/{self._module_name}", params=params, response_cls=ListRouteResponse
45+
)
46+
47+
def retrieve(
48+
self,
49+
*,
50+
assistant_id: str,
51+
route_id: str,
52+
) -> RouteResponse:
53+
return self._get(path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}", response_cls=RouteResponse)
54+
55+
def modify(
56+
self,
57+
*,
58+
assistant_id: str,
59+
route_id: str,
60+
description: str | NotGiven = NotGiven,
61+
examples: List[str] | NotGiven = NotGiven,
62+
) -> RouteResponse:
63+
body = self._create_body(
64+
description=description,
65+
examples=examples,
66+
)
67+
68+
return self._patch(
69+
path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}", body=body, response_cls=RouteResponse
70+
)
71+
72+
def delete(
73+
self,
74+
*,
75+
assistant_id: str,
76+
route_id: str,
77+
):
78+
return self._delete(path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}")
79+
80+
81+
class AsyncAssistantRoutes(AsyncStudioResource, BaseRoutes):
82+
async def create(
83+
self,
84+
*,
85+
assistant_id: str,
86+
plan_id: str,
87+
name: str,
88+
description: str,
89+
examples: List[str],
90+
**kwargs,
91+
) -> RouteResponse:
92+
body = self._create_body(
93+
plan_id=plan_id,
94+
name=name,
95+
description=description,
96+
examples=examples,
97+
**kwargs,
98+
)
99+
100+
return await self._post(
101+
path=f"/assistants/{assistant_id}/{self._module_name}", body=body, response_cls=RouteResponse
102+
)
103+
104+
async def list(
105+
self,
106+
*,
107+
assistant_id: str,
108+
name: str | NotGiven = NotGiven,
109+
) -> ListRouteResponse:
110+
params = self._create_body(name=name)
111+
112+
return await self._get(
113+
path=f"/assistants/{assistant_id}/{self._module_name}", params=params, response_cls=ListRouteResponse
114+
)
115+
116+
async def retrieve(
117+
self,
118+
*,
119+
assistant_id: str,
120+
route_id: str,
121+
) -> RouteResponse:
122+
return await self._get(
123+
path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}", response_cls=RouteResponse
124+
)
125+
126+
async def modify(
127+
self,
128+
*,
129+
assistant_id: str,
130+
route_id: str,
131+
description: str | NotGiven = NotGiven,
132+
examples: List[str] | NotGiven = NotGiven,
133+
) -> RouteResponse:
134+
body = self._create_body(
135+
description=description,
136+
examples=examples,
137+
)
138+
139+
return await self._patch(
140+
path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}", body=body, response_cls=RouteResponse
141+
)
142+
143+
async def delete(
144+
self,
145+
*,
146+
assistant_id: str,
147+
route_id: str,
148+
):
149+
return await self._delete(path=f"/assistants/{assistant_id}/{self._module_name}/{route_id}")

0 commit comments

Comments
 (0)