feat: add tools and retrievals

Author: jameszyao

taskingai/assistant/assistant.py

@@ -7,6 +7,7 @@
     AssistantGetResponse, AssistantListResponse
 
 __all__ = [
+    "Assistant",
     "get_assistant",
     "list_assistants",
     "create_assistant",
@@ -33,7 +34,7 @@ def list_assistants(
     :param before: The cursor to get the previous page of assistants.
     :return: The list of assistants.
     """
-
+    # todo: verify only one of offset, after and before is not None
     api_instance = get_assistant_api_instance()
     # only add non-None parameters
     params = {
@@ -70,7 +71,7 @@ def create_assistant(
     system_prompt_template: Optional[List[str]] = None,
     tools: Optional[List[AssistantTool]] = None,
     retrievals: Optional[List[AssistantRetrieval]] = None,
-    metadata: Optional[Dict] = None,
+    metadata: Optional[Dict[str, str]] = None,
 ) -> Assistant:
     """
     Create an assistant.
@@ -82,7 +83,7 @@ def create_assistant(
     :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.
+    :return: The created assistant object.
     """
 
     api_instance = get_assistant_api_instance()
@@ -108,7 +109,7 @@ def update_assistant(
     system_prompt_template: Optional[List[str]] = None,
     tools: Optional[List[AssistantTool]] = None,
     retrievals: Optional[List[AssistantRetrieval]] = None,
-    metadata: Optional[Dict] = None,
+    metadata: Optional[Dict[str, str]] = None,
 ) -> Assistant:
     """
     Update an assistant.
@@ -121,7 +122,7 @@ def update_assistant(
     :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.
+    :return: The updated assistant object.
     """
 
     api_instance = get_assistant_api_instance()

taskingai/assistant/chat.py

@@ -7,6 +7,7 @@
     ChatGetResponse, ChatListResponse
 
 __all__ = [
+    "Chat",
     "get_chat",
     "list_chats",
     "create_chat",
@@ -24,7 +25,7 @@ def list_chats(
 ) -> List[Chat]:
     """
     List chats.
-
+    :param assistant_id: The ID of the assistant.
     :param order: The order of the chats. It can be "asc" or "desc".
     :param limit: The maximum number of chats to return.
     :param after: The cursor to get the next page of chats.
@@ -68,14 +69,14 @@ def get_chat(assistant_id: str, chat_id: str) -> Chat:
 
 def create_chat(
     assistant_id: str,
-    metadata: Optional[Dict] = None,
+    metadata: Optional[Dict[str, str]] = None,
 ) -> Chat:
     """
     Create a chat.
 
     :param assistant_id: The ID of the assistant.
     :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.
-    :return: The assistant object.
+    :return: The created chat object.
     """
 
     api_instance = get_assistant_api_instance()
@@ -93,15 +94,15 @@ def create_chat(
 def update_chat(
     assistant_id: str,
     chat_id: str,
-    metadata: Dict,
+    metadata: Dict[str, str],
 ) -> Chat:
     """
     Update a chat.
 
     :param assistant_id: The ID of the assistant.
     :param chat_id: The ID of the chat.
     :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.
+    :return: The updated chat object.
     """
 
     api_instance = get_assistant_api_instance()

taskingai/assistant/message.py

@@ -7,6 +7,7 @@
     MessageGetResponse, MessageListResponse, MessageGenerateRequest
 
 __all__ = [
+    "Message",
     "get_message",
     "list_messages",
     "create_user_message",
@@ -26,6 +27,8 @@ def list_messages(
     """
     List messages.
 
+    :param assistant_id: The ID of the assistant.
+    :param chat_id: The ID of the chat.
     :param order: The order of the messages. It can be "asc" or "desc".
     :param limit: The maximum number of messages to return.
     :param after: The cursor to get the next page of messages.
@@ -78,7 +81,7 @@ def create_user_message(
     assistant_id: str,
     chat_id: str,
     text: str,
-    metadata: Optional[Dict] = None,
+    metadata: Optional[Dict[str, str]] = None,
 ) -> Message:
     """
     Create a message.
@@ -87,7 +90,7 @@ def create_user_message(
     :param chat_id: The ID of the chat.
     :param text: The text content of the message.
     :param metadata: The message 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.
+    :return: The created message object.
     """
 
     api_instance = get_assistant_api_instance()
@@ -108,15 +111,15 @@ def create_user_message(
 def update_message(
     assistant_id: str,
     message_id: str,
-    metadata: Dict,
+    metadata: Dict[str, str],
 ) -> Message:
     """
     Update a message.
 
     :param assistant_id: The ID of the assistant.
     :param message_id: The ID of the message.
     :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.
+    :return: The updated message object.
     """
 
     api_instance = get_assistant_api_instance()
@@ -145,7 +148,7 @@ def generate_assistant_message(
     :param chat_id: The ID of the chat.
     :param text: The text content of the message.
     :param metadata: The message 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.
+    :return: The generated message object.
     """
 
     api_instance = get_assistant_api_instance()

taskingai/client/__init__.py

@@ -20,8 +20,7 @@
 from taskingai.client.configuration import Configuration
 # import models into sdk package
 from taskingai.client.models.action import Action
-from taskingai.client.models.action_authentication_input import ActionAuthenticationInput
-from taskingai.client.models.action_authentication_output import ActionAuthenticationOutput
+from taskingai.client.models.action_authentication import ActionAuthentication
 from taskingai.client.models.action_authentication_type import ActionAuthenticationType
 from taskingai.client.models.action_create_request import ActionCreateRequest
 from taskingai.client.models.action_create_response import ActionCreateResponse
@@ -62,8 +61,7 @@
 from taskingai.client.models.chunk_query_request import ChunkQueryRequest
 from taskingai.client.models.chunk_query_response import ChunkQueryResponse
 from taskingai.client.models.collection import Collection
-from taskingai.client.models.collection_config_input import CollectionConfigInput
-from taskingai.client.models.collection_config_output import CollectionConfigOutput
+from taskingai.client.models.collection_config import CollectionConfig
 from taskingai.client.models.collection_create_request import CollectionCreateRequest
 from taskingai.client.models.collection_create_response import CollectionCreateResponse
 from taskingai.client.models.collection_get_response import CollectionGetResponse

taskingai/client/api/retrieval_api.py

@@ -133,7 +133,7 @@ def create_record(self, body, collection_id, **kwargs):  # noqa: E501
         Create a new record in a collection.  # noqa: E501
         This method makes a synchronous HTTP request by default. To make an
         asynchronous HTTP request, please pass async_req=True
-        >>> thread = api.create_record(body, collection_id, async_req=True)
+        >>> thread = api.create_text_record(body, collection_id, async_req=True)
         >>> result = thread.get()
 
         :param async_req bool

taskingai/client/models/__init__.py

@@ -11,8 +11,7 @@
 
 # import models into model package
 from taskingai.client.models.action import Action
-from taskingai.client.models.action_authentication_input import ActionAuthenticationInput
-from taskingai.client.models.action_authentication_output import ActionAuthenticationOutput
+from taskingai.client.models.action_authentication import ActionAuthentication
 from taskingai.client.models.action_authentication_type import ActionAuthenticationType
 from taskingai.client.models.action_create_request import ActionCreateRequest
 from taskingai.client.models.action_create_response import ActionCreateResponse
@@ -53,8 +52,7 @@
 from taskingai.client.models.chunk_query_request import ChunkQueryRequest
 from taskingai.client.models.chunk_query_response import ChunkQueryResponse
 from taskingai.client.models.collection import Collection
-from taskingai.client.models.collection_config_input import CollectionConfigInput
-from taskingai.client.models.collection_config_output import CollectionConfigOutput
+from taskingai.client.models.collection_config import CollectionConfig
 from taskingai.client.models.collection_create_request import CollectionCreateRequest
 from taskingai.client.models.collection_create_response import CollectionCreateResponse
 from taskingai.client.models.collection_get_response import CollectionGetResponse

…nt/models/action_authentication_input.py …i/client/models/action_authentication.pytaskingai/client/models/action_authentication_input.py renamed to taskingai/client/models/action_authentication.py taskingai/client/models/action_authentication_input.py renamed to taskingai/client/models/action_authentication.py

@@ -11,7 +11,7 @@
 
 import six
 
-class ActionAuthenticationInput(object):
+class ActionAuthentication(object):
     """NOTE: This class is auto generated by the swagger code generator program.
 
     Do not edit the class manually.
@@ -36,7 +36,7 @@ class ActionAuthenticationInput(object):
     }
 
     def __init__(self, type=None, secret=None, content=None):  # noqa: E501
-        """ActionAuthenticationInput - a model defined in Swagger"""  # noqa: E501
+        """ActionAuthentication - a model defined in Swagger"""  # noqa: E501
         self._type = None
         self._secret = None
         self._content = None
@@ -49,20 +49,20 @@ def __init__(self, type=None, secret=None, content=None):  # noqa: E501
 
     @property
     def type(self):
-        """Gets the type of this ActionAuthenticationInput.  # noqa: E501
+        """Gets the type of this ActionAuthentication.  # noqa: E501
 
 
-        :return: The type of this ActionAuthenticationInput.  # noqa: E501
+        :return: The type of this ActionAuthentication.  # noqa: E501
         :rtype: ActionAuthenticationType
         """
         return self._type
 
     @type.setter
     def type(self, type):
-        """Sets the type of this ActionAuthenticationInput.
+        """Sets the type of this ActionAuthentication.
 
 
-        :param type: The type of this ActionAuthenticationInput.  # noqa: E501
+        :param type: The type of this ActionAuthentication.  # noqa: E501
         :type: ActionAuthenticationType
         """
         if type is None:
@@ -72,41 +72,41 @@ def type(self, type):
 
     @property
     def secret(self):
-        """Gets the secret of this ActionAuthenticationInput.  # noqa: E501
+        """Gets the secret of this ActionAuthentication.  # noqa: E501
 
 
-        :return: The secret of this ActionAuthenticationInput.  # noqa: E501
+        :return: The secret of this ActionAuthentication.  # noqa: E501
         :rtype: object
         """
         return self._secret
 
     @secret.setter
     def secret(self, secret):
-        """Sets the secret of this ActionAuthenticationInput.
+        """Sets the secret of this ActionAuthentication.
 
 
-        :param secret: The secret of this ActionAuthenticationInput.  # noqa: E501
+        :param secret: The secret of this ActionAuthentication.  # noqa: E501
         :type: object
         """
 
         self._secret = secret
 
     @property
     def content(self):
-        """Gets the content of this ActionAuthenticationInput.  # noqa: E501
+        """Gets the content of this ActionAuthentication.  # noqa: E501
 
 
-        :return: The content of this ActionAuthenticationInput.  # noqa: E501
+        :return: The content of this ActionAuthentication.  # noqa: E501
         :rtype: object
         """
         return self._content
 
     @content.setter
     def content(self, content):
-        """Sets the content of this ActionAuthenticationInput.
+        """Sets the content of this ActionAuthentication.
 
 
-        :param content: The content of this ActionAuthenticationInput.  # noqa: E501
+        :param content: The content of this ActionAuthentication.  # noqa: E501
         :type: object
         """
 
@@ -133,7 +133,7 @@ def to_dict(self):
                 ))
             else:
                 result[attr] = value
-        if issubclass(ActionAuthenticationInput, dict):
+        if issubclass(ActionAuthentication, dict):
             for key, value in self.items():
                 result[key] = value
 
@@ -149,7 +149,7 @@ def __repr__(self):
 
     def __eq__(self, other):
         """Returns true if both objects are equal"""
-        if not isinstance(other, ActionAuthenticationInput):
+        if not isinstance(other, ActionAuthentication):
             return False
 
         return self.__dict__ == other.__dict__