feat: add create assistant

Author: jameszyao

taskingai/__init__.py

@@ -1 +1,2 @@
 from .config import *
+from . import assistant

taskingai/__version__

@@ -0,0 +1 @@
+0.0.2

taskingai/assistant/__init__.py

@@ -0,0 +1 @@
+from .assistant import *

taskingai/assistant/assistant.py

@@ -0,0 +1,59 @@
+import time
+from typing import NamedTuple, Optional, List
+
+from taskingai.config import Config
+from taskingai.client.utils import get_user_agent
+from taskingai.client.api.assistant_api import AssistantApi
+from taskingai.client.api_client import ApiClient
+from taskingai.client.models import Assistant, AssistantRetrieval, AssistantTool
+from taskingai.client.models import AssistantCreateRequest, AssistantCreateResponse,\
+    AssistantUpdateRequest, AssistantUpdateResponse,\
+    AssistantGetResponse, AssistantListResponse, AssistantDeleteResponse
+
+__all__ = [
+    "create_assistant",
+]
+
+def _get_api_instance():
+    client_config = Config.OPENAPI_CONFIG
+    client_config.api_key = client_config.api_key or {}
+    api_client = ApiClient(configuration=client_config)
+    api_client.user_agent = get_user_agent()
+    api_instance = AssistantApi(api_client)
+    return api_instance
+
+
+def create_assistant(
+    model_id: str,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    system_prompt_template: Optional[List[str]] = None,
+    tools: Optional[List[AssistantTool]] = None,
+    retrievals: Optional[List[AssistantRetrieval]] = None,
+    metadata: Optional[dict] = None,
+) -> Assistant:
+    """Create an assistant.
+    
+    :param model_id: The ID of an available chat completion model in your project.
+    :param name: The assistant name.
+    :param description: The assistant description.
+    :param system_prompt_template: A list of system prompt chunks where prompt variables are wrapped by curly brackets, e.g. {{variable}}.
+    :param tools: The assistant tools.
+    :param retrievals: The assistant retrievals.
+    :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.
+    :return: The assistant object.
+    """
+
+    api_instance = _get_api_instance()
+    body = AssistantCreateRequest(
+        model_id=model_id,
+        name=name,
+        description=description,
+        system_prompt_template=system_prompt_template,
+        tools=tools,
+        retrievals=retrievals,
+        metadata=metadata,
+    )
+    response: AssistantCreateResponse = api_instance.create_assistant(body=body)
+    assistant: Assistant = Assistant(**response.data)
+    return assistant

taskingai/client/configuration.py

@@ -49,7 +49,7 @@ def __init__(self):
 
         # Authentication Settings
         # dict to store API key(s)
-        self.api_key = {}
+        self.api_key = ""
         env_api_key = os.getenv('TASKINGAI_API_KEY')
         if env_api_key:
             self.api_key = env_api_key

taskingai/client/constants.py

@@ -1,7 +1,18 @@
+from enum import Enum
+
 PARENT_LOGGER_NAME = 'taskingai'
 DEFAULT_PARENT_LOGGER_LEVEL = 'ERROR'
 
 REQUEST_ID: str = "request_id"
 CLIENT_VERSION_HEADER = 'X-TaskingAI-Client-Version'
 
 
+class ToolType(str, Enum):
+    action = "action"
+    function = "function"
+
+
+class RetrievalType(str, Enum):
+    collection = "collection"
+
+

taskingai/client/error_handling.py

@@ -0,0 +1,28 @@
+import inspect
+from functools import wraps
+
+from urllib3.exceptions import MaxRetryError, ProtocolError
+
+from taskingai import Config
+from .exceptions import TaskingAIProtocolError
+
+
+def validate_and_convert_errors(func):
+    @wraps(func)
+    def inner_func(*args, **kwargs):
+        Config.validate()  # raises exceptions in case of invalid config
+        try:
+            return func(*args, **kwargs)
+        except MaxRetryError as e:
+            if isinstance(e.reason, ProtocolError):
+                raise TaskingAIProtocolError(
+                    f'Failed to connect to {e.url}. Please check your network connection and try again later.') from e
+            else:
+                raise
+        except ProtocolError as e:
+            raise TaskingAIProtocolError(f'Failed to connect. Please check your network connection and try again later.') from e
+
+    # Override signature
+    sig = inspect.signature(func)
+    inner_func.__signature__ = sig
+    return inner_func

taskingai/client/models/assistant.py

@@ -168,9 +168,6 @@ def name(self, name):
         :param name: The name of this Assistant.  # noqa: E501
         :type: object
         """
-        if name is None:
-            raise ValueError("Invalid value for `name`, must not be `None`")  # noqa: E501
-
         self._name = name
 
     @property
@@ -193,9 +190,6 @@ def description(self, description):
         :param description: The description of this Assistant.  # noqa: E501
         :type: object
         """
-        if description is None:
-            raise ValueError("Invalid value for `description`, must not be `None`")  # noqa: E501
-
         self._description = description
 
     @property
@@ -218,9 +212,6 @@ def system_prompt_template(self, system_prompt_template):
         :param system_prompt_template: The system_prompt_template of this Assistant.  # noqa: E501
         :type: object
         """
-        if system_prompt_template is None:
-            raise ValueError("Invalid value for `system_prompt_template`, must not be `None`")  # noqa: E501
-
         self._system_prompt_template = system_prompt_template
 
     @property
@@ -243,9 +234,6 @@ def tools(self, tools):
         :param tools: The tools of this Assistant.  # noqa: E501
         :type: object
         """
-        if tools is None:
-            raise ValueError("Invalid value for `tools`, must not be `None`")  # noqa: E501
-
         self._tools = tools
 
     @property
@@ -268,9 +256,6 @@ def retrievals(self, retrievals):
         :param retrievals: The retrievals of this Assistant.  # noqa: E501
         :type: object
         """
-        if retrievals is None:
-            raise ValueError("Invalid value for `retrievals`, must not be `None`")  # noqa: E501
-
         self._retrievals = retrievals
 
     @property
@@ -293,9 +278,6 @@ def metadata(self, metadata):
         :param metadata: The metadata of this Assistant.  # noqa: E501
         :type: object
         """
-        if metadata is None:
-            raise ValueError("Invalid value for `metadata`, must not be `None`")  # noqa: E501
-
         self._metadata = metadata
 
     @property
@@ -318,9 +300,6 @@ def created_timestamp(self, created_timestamp):
         :param created_timestamp: The created_timestamp of this Assistant.  # noqa: E501
         :type: object
         """
-        if created_timestamp is None:
-            raise ValueError("Invalid value for `created_timestamp`, must not be `None`")  # noqa: E501
-
         self._created_timestamp = created_timestamp
 
     def to_dict(self):

taskingai/client/utils.py

@@ -1,8 +1,23 @@
 import inspect
 import logging
+from pathlib import Path
+import urllib3
 
 def check_kwargs(caller, given):
     argspec = inspect.getfullargspec(caller)
     diff = set(given).difference(argspec.args)
     if diff:
-        logging.exception(caller.__name__ + ' had unexpected keyword argument(s): ' + ', '.join(diff), exc_info=False)
+        logging.exception(caller.__name__ + ' had unexpected keyword argument(s): ' + ', '.join(diff), exc_info=False)
+
+
+def get_version():
+    return Path(__file__).parent.parent.joinpath('__version__').read_text().strip()
+
+
+def get_user_agent():
+    client_id = f'python-client-{get_version()}'
+    user_agent_details = {'urllib3': urllib3.__version__}
+    user_agent = '{} ({})'.format(client_id, ', '.join([f'{k}:{v}' for k, v in user_agent_details.items()]))
+    return user_agent
+
+

taskingai/config.py

@@ -45,28 +45,17 @@ def validate(self):
     def reset(self, **kwargs):
         config = ConfigBase()
 
-        # Get the environment first. Make sure that it is not overwritten in subsequent config objects.
+        # Get from kwargs first, then environment variables, then default
         api_key = (
-            kwargs.pop("api_key", None)
-            or os.getenv("TASKINGAI_API_KEY")
+            kwargs.pop("api_key", None) or os.getenv("TASKINGAI_API_KEY")
         )
-        config = config._replace(api_key=api_key)
-
-        # Set default config
-        default_config = ConfigBase(
-            host="https://api.tasking.ai"
-        )
-        config = config._replace(**self._preprocess_and_validate_config(default_config._asdict()))
-
-        # Set environment config
-        env_config = ConfigBase(
-            api_key=os.getenv("TASKINGAI_API_KEY"),
-            host=os.getenv("TASKINGAI_HOST"),
+        host = (
+            kwargs.pop("host", None) or os.getenv("TASKINGAI_HOST") or "https://api.tasking.ai"
         )
-        config = config._replace(**self._preprocess_and_validate_config(env_config._asdict()))
+        config = config._replace(api_key=api_key, host=host)
 
         # Set explicit config
-        config = config._replace(**self._preprocess_and_validate_config(kwargs))
+        # config = config._replace(**self._preprocess_and_validate_config(kwargs))
 
         self._config = config