Skip to content

Commit 2491b43

Browse files
jameszyaoSimsonW
jameszyao
authored and
SimsonW
committed
feat: add chat & message
1 parent 3aaaeb6 commit 2491b43

File tree

11 files changed

+650
-83
lines changed

11 files changed

+650
-83
lines changed

‎taskingai/assistant/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
from .assistant import *
2+
from .chat import *

‎taskingai/assistant/assistant.py‎

Lines changed: 110 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,26 +1,66 @@
1-
import time
2-
from typing import NamedTuple, Optional, List
1+
from typing import Optional, List, Dict
32

4-
from taskingai.config import Config
5-
from taskingai.client.utils import get_user_agent
6-
from taskingai.client.api.assistant_api import AssistantApi
7-
from taskingai.client.api_client import ApiClient
3+
from taskingai.client.utils import get_assistant_api_instance
84
from taskingai.client.models import Assistant, AssistantRetrieval, AssistantTool
95
from taskingai.client.models import AssistantCreateRequest, AssistantCreateResponse,\
106
AssistantUpdateRequest, AssistantUpdateResponse,\
11-
AssistantGetResponse, AssistantListResponse, AssistantDeleteResponse
7+
AssistantGetResponse, AssistantListResponse
128

139
__all__ = [
10+
"get_assistant",
11+
"list_assistants",
1412
"create_assistant",
13+
"update_assistant",
14+
"delete_assistant",
1515
]
1616

17-
def _get_api_instance():
18-
client_config = Config.OPENAPI_CONFIG
19-
client_config.api_key = client_config.api_key or {}
20-
api_client = ApiClient(configuration=client_config)
21-
api_client.user_agent = get_user_agent()
22-
api_instance = AssistantApi(api_client)
23-
return api_instance
17+
18+
def list_assistants(
19+
order: str = "desc",
20+
limit: int = 20,
21+
offset: Optional[int] = None,
22+
after: Optional[str] = None,
23+
before: Optional[str] = None,
24+
) -> List[Assistant]:
25+
26+
"""
27+
List assistants.
28+
29+
:param order: The order of the assistants. It can be "asc" or "desc".
30+
:param limit: The maximum number of assistants to return.
31+
:param offset: The offset of the assistants.
32+
:param after: The cursor to get the next page of assistants.
33+
:param before: The cursor to get the previous page of assistants.
34+
:return: The list of assistants.
35+
"""
36+
37+
api_instance = get_assistant_api_instance()
38+
# only add non-None parameters
39+
params = {
40+
"order": order,
41+
"limit": limit,
42+
"offset": offset,
43+
"after": after,
44+
"before": before,
45+
}
46+
params = {k: v for k, v in params.items() if v is not None}
47+
response: AssistantListResponse = api_instance.list_assistants(**params)
48+
assistants: List[Assistant] = [Assistant(**item) for item in response.data]
49+
return assistants
50+
51+
52+
53+
def get_assistant(assistant_id: str) -> Assistant:
54+
"""
55+
Get an assistant.
56+
57+
:param assistant_id: The ID of the assistant.
58+
"""
59+
60+
api_instance = get_assistant_api_instance()
61+
response: AssistantGetResponse = api_instance.get_assistant(assistant_id=assistant_id)
62+
assistant: Assistant = Assistant(**response.data)
63+
return assistant
2464

2565

2666
def create_assistant(
@@ -30,9 +70,10 @@ def create_assistant(
3070
system_prompt_template: Optional[List[str]] = None,
3171
tools: Optional[List[AssistantTool]] = None,
3272
retrievals: Optional[List[AssistantRetrieval]] = None,
33-
metadata: Optional[dict] = None,
73+
metadata: Optional[Dict] = None,
3474
) -> Assistant:
35-
"""Create an assistant.
75+
"""
76+
Create an assistant.
3677
3778
:param model_id: The ID of an available chat completion model in your project.
3879
:param name: The assistant name.
@@ -44,7 +85,7 @@ def create_assistant(
4485
:return: The assistant object.
4586
"""
4687

47-
api_instance = _get_api_instance()
88+
api_instance = get_assistant_api_instance()
4889
body = AssistantCreateRequest(
4990
model_id=model_id,
5091
name=name,
@@ -56,4 +97,55 @@ def create_assistant(
5697
)
5798
response: AssistantCreateResponse = api_instance.create_assistant(body=body)
5899
assistant: Assistant = Assistant(**response.data)
59-
return assistant
100+
return assistant
101+
102+
103+
def update_assistant(
104+
assistant_id: str,
105+
model_id: Optional[str] = None,
106+
name: Optional[str] = None,
107+
description: Optional[str] = None,
108+
system_prompt_template: Optional[List[str]] = None,
109+
tools: Optional[List[AssistantTool]] = None,
110+
retrievals: Optional[List[AssistantRetrieval]] = None,
111+
metadata: Optional[Dict] = None,
112+
) -> Assistant:
113+
"""
114+
Update an assistant.
115+
116+
:param assistant_id: The ID of the assistant.
117+
:param model_id: The ID of an available chat completion model in your project.
118+
:param name: The assistant name.
119+
:param description: The assistant description.
120+
:param system_prompt_template: A list of system prompt chunks where prompt variables are wrapped by curly brackets, e.g. {{variable}}.
121+
:param tools: The assistant tools.
122+
:param retrievals: The assistant retrievals.
123+
:param metadata: The assistant metadata. It can store up to 16 key-value pairs where each key's length is less than 64 and value's length is less than 512.
124+
:return: The assistant object.
125+
"""
126+
127+
api_instance = get_assistant_api_instance()
128+
body = AssistantUpdateRequest(
129+
model_id=model_id,
130+
name=name,
131+
description=description,
132+
system_prompt_template=system_prompt_template,
133+
tools=tools,
134+
retrievals=retrievals,
135+
metadata=metadata,
136+
)
137+
response: AssistantUpdateResponse = api_instance.update_assistant(assistant_id=assistant_id, body=body)
138+
assistant: Assistant = Assistant(**response.data)
139+
return assistant
140+
141+
142+
def delete_assistant(assistant_id: str) -> None:
143+
"""
144+
Delete an assistant.
145+
146+
:param assistant_id: The ID of the assistant.
147+
"""
148+
149+
api_instance = get_assistant_api_instance()
150+
api_instance.delete_assistant(assistant_id=assistant_id)
151+

‎taskingai/assistant/chat.py‎

Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
from typing import Optional, List, Dict
2+
3+
from taskingai.client.utils import get_assistant_api_instance
4+
from taskingai.client.models import Chat
5+
from taskingai.client.models import ChatCreateRequest, ChatCreateResponse, \
6+
ChatUpdateRequest, ChatUpdateResponse, \
7+
ChatGetResponse, ChatListResponse
8+
9+
__all__ = [
10+
"get_chat",
11+
"list_chats",
12+
"create_chat",
13+
"update_chat",
14+
"delete_chat",
15+
]
16+
17+
18+
def list_chats(
19+
assistant_id: str,
20+
order: str = "desc",
21+
limit: int = 20,
22+
after: Optional[str] = None,
23+
before: Optional[str] = None,
24+
) -> List[Chat]:
25+
"""
26+
List chats.
27+
28+
:param order: The order of the chats. It can be "asc" or "desc".
29+
:param limit: The maximum number of chats to return.
30+
:param after: The cursor to get the next page of chats.
31+
:param before: The cursor to get the previous page of chats.
32+
:return: The list of chats.
33+
"""
34+
35+
api_instance = get_assistant_api_instance()
36+
# only add non-None parameters
37+
params = {
38+
"order": order,
39+
"limit": limit,
40+
"after": after,
41+
"before": before,
42+
}
43+
params = {k: v for k, v in params.items() if v is not None}
44+
response: ChatListResponse = api_instance.list_chats(
45+
assistant_id=assistant_id,
46+
**params
47+
)
48+
chats: List[Chat] = [Chat(**item) for item in response.data]
49+
return chats
50+
51+
52+
def get_chat(assistant_id: str, chat_id: str) -> Chat:
53+
"""
54+
Get a chat.
55+
56+
:param assistant_id: The ID of the assistant.
57+
:param chat_id: The ID of the chat.
58+
"""
59+
60+
api_instance = get_assistant_api_instance()
61+
response: ChatGetResponse = api_instance.get_chat(
62+
assistant_id=assistant_id,
63+
chat_id=chat_id,
64+
)
65+
chat: Chat = Chat(**response.data)
66+
return chat
67+
68+
69+
def create_chat(
70+
assistant_id: str,
71+
metadata: Optional[Dict] = None,
72+
) -> Chat:
73+
"""
74+
Create a chat.
75+
76+
:param assistant_id: The ID of the assistant.
77+
:param metadata: The chat metadata. It can store up to 16 key-value pairs where each key's length is less than 64 and value's length is less than 512.
78+
:return: The assistant object.
79+
"""
80+
81+
api_instance = get_assistant_api_instance()
82+
body = ChatCreateRequest(
83+
metadata=metadata,
84+
)
85+
response: ChatCreateResponse = api_instance.create_chat(
86+
assistant_id=assistant_id,
87+
body=body
88+
)
89+
chat: Chat = Chat(**response.data)
90+
return chat
91+
92+
93+
def update_chat(
94+
assistant_id: str,
95+
chat_id: str,
96+
metadata: Dict,
97+
) -> Chat:
98+
"""
99+
Update a chat.
100+
101+
:param assistant_id: The ID of the assistant.
102+
:param chat_id: The ID of the chat.
103+
:param metadata: The assistant metadata. It can store up to 16 key-value pairs where each key's length is less than 64 and value's length is less than 512.
104+
:return: The assistant object.
105+
"""
106+
107+
api_instance = get_assistant_api_instance()
108+
body = ChatUpdateRequest(
109+
metadata=metadata,
110+
)
111+
response: ChatUpdateResponse = api_instance.update_chat(
112+
assistant_id=assistant_id,
113+
chat_id=chat_id,
114+
body=body
115+
)
116+
chat: Chat = Chat(**response.data)
117+
return chat
118+
119+
120+
def delete_chat(
121+
assistant_id: str,
122+
chat_id: str,
123+
) -> None:
124+
"""
125+
Delete a chat.
126+
127+
:param assistant_id: The ID of the assistant.
128+
:param chat_id: The ID of the chat.
129+
"""
130+
131+
api_instance = get_assistant_api_instance()
132+
api_instance.delete_chat(assistant_id=assistant_id, chat_id=chat_id)
133+

0 commit comments

Comments
 (0)