Skip to content

Commit 3aaaeb6

Browse files
jameszyaoSimsonW
authored andcommitted
feat: add create assistant
1 parent 503db80 commit 3aaaeb6

File tree

10 files changed

+124
-40
lines changed

10 files changed

+124
-40
lines changed

taskingai/__init__.py

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

taskingai/__version__

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
0.0.2

taskingai/assistant/__init__.py

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

taskingai/assistant/assistant.py

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
import time
2+
from typing import NamedTuple, Optional, List
3+
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
8+
from taskingai.client.models import Assistant, AssistantRetrieval, AssistantTool
9+
from taskingai.client.models import AssistantCreateRequest, AssistantCreateResponse,\
10+
AssistantUpdateRequest, AssistantUpdateResponse,\
11+
AssistantGetResponse, AssistantListResponse, AssistantDeleteResponse
12+
13+
__all__ = [
14+
"create_assistant",
15+
]
16+
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
24+
25+
26+
def create_assistant(
27+
model_id: str,
28+
name: Optional[str] = None,
29+
description: Optional[str] = None,
30+
system_prompt_template: Optional[List[str]] = None,
31+
tools: Optional[List[AssistantTool]] = None,
32+
retrievals: Optional[List[AssistantRetrieval]] = None,
33+
metadata: Optional[dict] = None,
34+
) -> Assistant:
35+
"""Create an assistant.
36+
37+
:param model_id: The ID of an available chat completion model in your project.
38+
:param name: The assistant name.
39+
:param description: The assistant description.
40+
:param system_prompt_template: A list of system prompt chunks where prompt variables are wrapped by curly brackets, e.g. {{variable}}.
41+
:param tools: The assistant tools.
42+
:param retrievals: The assistant retrievals.
43+
: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.
44+
:return: The assistant object.
45+
"""
46+
47+
api_instance = _get_api_instance()
48+
body = AssistantCreateRequest(
49+
model_id=model_id,
50+
name=name,
51+
description=description,
52+
system_prompt_template=system_prompt_template,
53+
tools=tools,
54+
retrievals=retrievals,
55+
metadata=metadata,
56+
)
57+
response: AssistantCreateResponse = api_instance.create_assistant(body=body)
58+
assistant: Assistant = Assistant(**response.data)
59+
return assistant

taskingai/client/configuration.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ def __init__(self):
4949

5050
# Authentication Settings
5151
# dict to store API key(s)
52-
self.api_key = {}
52+
self.api_key = ""
5353
env_api_key = os.getenv('TASKINGAI_API_KEY')
5454
if env_api_key:
5555
self.api_key = env_api_key

taskingai/client/constants.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,18 @@
1+
from enum import Enum
2+
13
PARENT_LOGGER_NAME = 'taskingai'
24
DEFAULT_PARENT_LOGGER_LEVEL = 'ERROR'
35

46
REQUEST_ID: str = "request_id"
57
CLIENT_VERSION_HEADER = 'X-TaskingAI-Client-Version'
68

79

10+
class ToolType(str, Enum):
11+
action = "action"
12+
function = "function"
13+
14+
15+
class RetrievalType(str, Enum):
16+
collection = "collection"
17+
18+

taskingai/client/error_handling.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
import inspect
2+
from functools import wraps
3+
4+
from urllib3.exceptions import MaxRetryError, ProtocolError
5+
6+
from taskingai import Config
7+
from .exceptions import TaskingAIProtocolError
8+
9+
10+
def validate_and_convert_errors(func):
11+
@wraps(func)
12+
def inner_func(*args, **kwargs):
13+
Config.validate() # raises exceptions in case of invalid config
14+
try:
15+
return func(*args, **kwargs)
16+
except MaxRetryError as e:
17+
if isinstance(e.reason, ProtocolError):
18+
raise TaskingAIProtocolError(
19+
f'Failed to connect to {e.url}. Please check your network connection and try again later.') from e
20+
else:
21+
raise
22+
except ProtocolError as e:
23+
raise TaskingAIProtocolError(f'Failed to connect. Please check your network connection and try again later.') from e
24+
25+
# Override signature
26+
sig = inspect.signature(func)
27+
inner_func.__signature__ = sig
28+
return inner_func

taskingai/client/models/assistant.py

Lines changed: 0 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -168,9 +168,6 @@ def name(self, name):
168168
:param name: The name of this Assistant. # noqa: E501
169169
:type: object
170170
"""
171-
if name is None:
172-
raise ValueError("Invalid value for `name`, must not be `None`") # noqa: E501
173-
174171
self._name = name
175172

176173
@property
@@ -193,9 +190,6 @@ def description(self, description):
193190
:param description: The description of this Assistant. # noqa: E501
194191
:type: object
195192
"""
196-
if description is None:
197-
raise ValueError("Invalid value for `description`, must not be `None`") # noqa: E501
198-
199193
self._description = description
200194

201195
@property
@@ -218,9 +212,6 @@ def system_prompt_template(self, system_prompt_template):
218212
:param system_prompt_template: The system_prompt_template of this Assistant. # noqa: E501
219213
:type: object
220214
"""
221-
if system_prompt_template is None:
222-
raise ValueError("Invalid value for `system_prompt_template`, must not be `None`") # noqa: E501
223-
224215
self._system_prompt_template = system_prompt_template
225216

226217
@property
@@ -243,9 +234,6 @@ def tools(self, tools):
243234
:param tools: The tools of this Assistant. # noqa: E501
244235
:type: object
245236
"""
246-
if tools is None:
247-
raise ValueError("Invalid value for `tools`, must not be `None`") # noqa: E501
248-
249237
self._tools = tools
250238

251239
@property
@@ -268,9 +256,6 @@ def retrievals(self, retrievals):
268256
:param retrievals: The retrievals of this Assistant. # noqa: E501
269257
:type: object
270258
"""
271-
if retrievals is None:
272-
raise ValueError("Invalid value for `retrievals`, must not be `None`") # noqa: E501
273-
274259
self._retrievals = retrievals
275260

276261
@property
@@ -293,9 +278,6 @@ def metadata(self, metadata):
293278
:param metadata: The metadata of this Assistant. # noqa: E501
294279
:type: object
295280
"""
296-
if metadata is None:
297-
raise ValueError("Invalid value for `metadata`, must not be `None`") # noqa: E501
298-
299281
self._metadata = metadata
300282

301283
@property
@@ -318,9 +300,6 @@ def created_timestamp(self, created_timestamp):
318300
:param created_timestamp: The created_timestamp of this Assistant. # noqa: E501
319301
:type: object
320302
"""
321-
if created_timestamp is None:
322-
raise ValueError("Invalid value for `created_timestamp`, must not be `None`") # noqa: E501
323-
324303
self._created_timestamp = created_timestamp
325304

326305
def to_dict(self):

taskingai/client/utils.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,23 @@
11
import inspect
22
import logging
3+
from pathlib import Path
4+
import urllib3
35

46
def check_kwargs(caller, given):
57
argspec = inspect.getfullargspec(caller)
68
diff = set(given).difference(argspec.args)
79
if diff:
8-
logging.exception(caller.__name__ + ' had unexpected keyword argument(s): ' + ', '.join(diff), exc_info=False)
10+
logging.exception(caller.__name__ + ' had unexpected keyword argument(s): ' + ', '.join(diff), exc_info=False)
11+
12+
13+
def get_version():
14+
return Path(__file__).parent.parent.joinpath('__version__').read_text().strip()
15+
16+
17+
def get_user_agent():
18+
client_id = f'python-client-{get_version()}'
19+
user_agent_details = {'urllib3': urllib3.__version__}
20+
user_agent = '{} ({})'.format(client_id, ', '.join([f'{k}:{v}' for k, v in user_agent_details.items()]))
21+
return user_agent
22+
23+

taskingai/config.py

Lines changed: 6 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -45,28 +45,17 @@ def validate(self):
4545
def reset(self, **kwargs):
4646
config = ConfigBase()
4747

48-
# Get the environment first. Make sure that it is not overwritten in subsequent config objects.
48+
# Get from kwargs first, then environment variables, then default
4949
api_key = (
50-
kwargs.pop("api_key", None)
51-
or os.getenv("TASKINGAI_API_KEY")
50+
kwargs.pop("api_key", None) or os.getenv("TASKINGAI_API_KEY")
5251
)
53-
config = config._replace(api_key=api_key)
54-
55-
# Set default config
56-
default_config = ConfigBase(
57-
host="https://api.tasking.ai"
58-
)
59-
config = config._replace(**self._preprocess_and_validate_config(default_config._asdict()))
60-
61-
# Set environment config
62-
env_config = ConfigBase(
63-
api_key=os.getenv("TASKINGAI_API_KEY"),
64-
host=os.getenv("TASKINGAI_HOST"),
52+
host = (
53+
kwargs.pop("host", None) or os.getenv("TASKINGAI_HOST") or "https://api.tasking.ai"
6554
)
66-
config = config._replace(**self._preprocess_and_validate_config(env_config._asdict()))
55+
config = config._replace(api_key=api_key, host=host)
6756

6857
# Set explicit config
69-
config = config._replace(**self._preprocess_and_validate_config(kwargs))
58+
# config = config._replace(**self._preprocess_and_validate_config(kwargs))
7059

7160
self._config = config
7261

0 commit comments

Comments
 (0)