feat: Use pydantic model to replace original ones

Author: jameszyao

examples/assistant/chat_with_assistant.ipynb

@@ -2,14 +2,11 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": null,
    "id": "initial_id",
    "metadata": {
     "collapsed": true,
-    "ExecuteTime": {
-     "end_time": "2023-11-28T08:47:28.230279Z",
-     "start_time": "2023-11-28T08:47:28.205410Z"
-    }
+    "is_executing": true
    },
    "outputs": [],
    "source": [
@@ -133,7 +130,7 @@
     "actions: List[Action] = taskingai.tool.bulk_create_actions(\n",
     "    schema=NUMBERS_API_SCHEMA,\n",
     "    authentication=ActionAuthentication(\n",
-    "        type=ActionAuthenticationType.NONE,\n",
+    "        type=ActionAuthenticationType.none,\n",
     "    )\n",
     ")\n",
     "action = actions[0]\n",
@@ -280,7 +277,7 @@
     "            \"language\": \"Spanish\"\n",
     "        }\n",
     "    )\n",
-    "    print(f\"Assistant: {assistant_message.content['text']}\")\n",
+    "    print(f\"Assistant: {assistant_message.content.text}\")\n",
     "    time.sleep(2)\n",
     "    user_input = input(\"User: \")"
    ],
@@ -319,7 +316,7 @@
     "    chat_id=chat.chat_id,\n",
     ")\n",
     "for message in messages:\n",
-    "    print(f\"{message.role}: {message.content['text']}\")"
+    "    print(f\"{message.role}: {message.content.text}\")"
    ],
    "metadata": {
     "collapsed": false,

examples/inference/chat_completion.ipynb

@@ -242,7 +242,7 @@
    ],
    "source": [
     "# add function message\n",
-    "chat_completion = taskingai.inference.chat_completion(\n",
+    "chat_completion_result = taskingai.inference.chat_completion(\n",
     "    model_id=model_id,\n",
     "    messages=[\n",
     "        SystemMessage(\"You are a professional assistant.\"),\n",

examples/retrieval/semantic_search.ipynb

@@ -193,7 +193,7 @@
   {
    "cell_type": "markdown",
    "source": [
-    "## Query Relavent Chunks"
+    "## Query Relevant Chunks"
    ],
    "metadata": {
     "collapsed": false

taskingai/assistant/assistant.py

@@ -1,8 +1,7 @@
 from typing import Optional, List, Dict
 
 from taskingai.client.utils import get_api_instance, ModuleType
-from taskingai.client.constants import AssistantToolType, AssistantRetrievalType
-from taskingai.client.models import Assistant, AssistantRetrieval, AssistantTool
+from taskingai.client.models import Assistant, AssistantRetrieval, AssistantTool, AssistantToolType, AssistantRetrievalType
 from taskingai.client.models import AssistantCreateRequest, AssistantCreateResponse,\
     AssistantUpdateRequest, AssistantUpdateResponse,\
     AssistantGetResponse, AssistantListResponse

taskingai/assistant/message.py

@@ -199,7 +199,7 @@ async def a_create_user_message(
 
     api_instance = get_api_instance(ModuleType.assistant, async_client=True)
     body = MessageCreateRequest(
-        role=MessageRole.USER,
+        role=MessageRole.user,
         content=MessageContent(text=text),
         metadata=metadata,
     )

taskingai/client/api_client.py

@@ -20,6 +20,7 @@
 from taskingai.client.configuration import Configuration
 from taskingai.client import rest
 import taskingai.client as client
+from pydantic import BaseModel
 
 
 class BaseApiClient(object):
@@ -105,6 +106,8 @@ def sanitize_for_serialization(self, obj):
                          for sub_obj in obj)
         elif isinstance(obj, (datetime.datetime, datetime.date)):
             return obj.isoformat()
+        elif isinstance(obj, BaseModel):
+            return obj.model_dump()
 
         if isinstance(obj, dict):
             obj_dict = obj

taskingai/client/constants.py

@@ -10,15 +10,6 @@ class ModuleType(str, Enum):
     inference = "inference"
 
 
-class AssistantToolType(str, Enum):
-    action = "action"
-    function = "function"
-
-
-class AssistantRetrievalType(str, Enum):
-    collection = "collection"
-
-
 
 
 

taskingai/client/models/__init__.py

@@ -10,33 +10,16 @@
 from __future__ import absolute_import
 
 # import models into model package
-from taskingai.client.models.entity.action import Action
-from taskingai.client.models.entity.action_authentication import ActionAuthentication
-from taskingai.client.models.entity.action_authentication_type import ActionAuthenticationType
-from taskingai.client.models.entity.assistant import Assistant
-from taskingai.client.models.entity.assistant_retrieval import AssistantRetrieval
-from taskingai.client.models.entity.assistant_tool import AssistantTool
-from taskingai.client.models.entity.chat import Chat
-from taskingai.client.models.entity.chat_completion import ChatCompletion
-from taskingai.client.models.entity.chat_completion_assistant_message import ChatCompletionAssistantMessage
-from taskingai.client.models.entity.chat_completion_function import ChatCompletionFunction
-from taskingai.client.models.entity.chat_completion_function_call import ChatCompletionFunctionCall
-from taskingai.client.models.entity.chat_completion_function_message import ChatCompletionFunctionMessage
-from taskingai.client.models.entity.chat_completion_role import ChatCompletionRole
-from taskingai.client.models.entity.chat_completion_system_message import ChatCompletionSystemMessage
-from taskingai.client.models.entity.chat_completion_user_message import ChatCompletionUserMessage
-from taskingai.client.models.entity.chunk import Chunk
-from taskingai.client.models.entity.collection import Collection
-from taskingai.client.models.entity.collection_config import CollectionConfig
-from taskingai.client.models.entity.function import Function
-from taskingai.client.models.entity.message import Message
-from taskingai.client.models.entity.message_content import MessageContent
-from taskingai.client.models.entity.message_role import MessageRole
-from taskingai.client.models.entity.parameters_input import ParametersInput
-from taskingai.client.models.entity.parameters_output import ParametersOutput
-from taskingai.client.models.entity.record import Record
-from taskingai.client.models.entity.text_embedding_output import TextEmbeddingOutput
-from taskingai.client.models.entity.message_generation_response_option import MessageGenerationResponseOption
+from taskingai.client.models.entity.assistant.assistant import *
+from taskingai.client.models.entity.assistant.chat import *
+from taskingai.client.models.entity.assistant.message import *
+from taskingai.client.models.entity.tool.action import *
+from taskingai.client.models.entity.tool.function import *
+from taskingai.client.models.entity.retrieval.collection import *
+from taskingai.client.models.entity.retrieval.record import *
+from taskingai.client.models.entity.retrieval.chunk import *
+from taskingai.client.models.entity.inference.text_embedding import *
+from taskingai.client.models.entity.inference.chat_completion import *
 from taskingai.client.models.rest.action_bulk_create_request import ActionBulkCreateRequest
 from taskingai.client.models.rest.action_bulk_create_response import ActionBulkCreateResponse
 from taskingai.client.models.rest.action_delete_response import ActionDeleteResponse

taskingai/client/models/entity/_base.py

@@ -0,0 +1,22 @@
+# coding: utf-8
+
+"""
+    TaskingAI API
+
+    OpenAPI spec version: 0.1.0
+"""
+
+from pydantic import BaseModel
+
+
+class TaskingaiBaseModel(BaseModel):
+
+    def to_dict(self):
+        """Returns the model properties as a dict"""
+        return self.model_dump()
+
+    def to_str(self):
+        """Returns the string representation of the model"""
+        return self.model_dump_json()
+
+