From 19a73f911dbea4c2aa044400c827437de13e38e1 Mon Sep 17 00:00:00 2001
From: simson <simson@wanmoai.com>
Date: Wed, 17 Apr 2024 21:40:36 +0800
Subject: [PATCH] fix: update new interface usage example

---
 examples/assistant/chat_with_assistant.ipynb | 52 ++++++++++----------
 taskingai/assistant/assistant.py             |  4 ++
 2 files changed, 30 insertions(+), 26 deletions(-)

diff --git a/examples/assistant/chat_with_assistant.ipynb b/examples/assistant/chat_with_assistant.ipynb
index e66fcd1..00393f3 100644
--- a/examples/assistant/chat_with_assistant.ipynb
+++ b/examples/assistant/chat_with_assistant.ipynb
@@ -2,17 +2,17 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": null,
    "id": "initial_id",
    "metadata": {
     "collapsed": true
    },
-   "outputs": [],
    "source": [
     "import time\n",
     "import taskingai\n",
     "# Load TaskingAI API Key from environment variable"
-   ]
+   ],
+   "outputs": [],
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
@@ -39,10 +39,8 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
    "source": [
-    "from taskingai.models import Action, ActionAuthentication, ActionAuthenticationType\n",
+    "from taskingai.tool import Action, ActionAuthentication, ActionAuthenticationType\n",
     "from typing import List\n",
     "\n",
     "# create an assistant action\n",
@@ -102,14 +100,14 @@
    "metadata": {
     "collapsed": false
    },
-   "id": "3b2fda39ba58c5e9"
+   "id": "3b2fda39ba58c5e9",
+   "outputs": [],
+   "execution_count": null
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
    "source": [
-    "from taskingai.models import Assistant, Chat, ToolRef, ToolType\n",
+    "from taskingai.assistant import Assistant, Chat, ToolRef, ToolType\n",
     "from taskingai.assistant.memory import AssistantMessageWindowMemory\n",
     "\n",
     "# choose an available chat_completion model from your project\n",
@@ -145,7 +143,9 @@
    "metadata": {
     "collapsed": false
    },
-   "id": "3b3df0f232021283"
+   "id": "3b3df0f232021283",
+   "outputs": [],
+   "execution_count": null
   },
   {
    "cell_type": "markdown",
@@ -159,8 +159,6 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
    "source": [
     "chat: Chat = taskingai.assistant.create_chat(\n",
     "    assistant_id=assistant.assistant_id,\n",
@@ -170,12 +168,12 @@
    "metadata": {
     "collapsed": false
    },
-   "id": "f1e2f0b2af8b1d8d"
+   "id": "f1e2f0b2af8b1d8d",
+   "outputs": [],
+   "execution_count": null
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
    "source": [
     "from taskingai.assistant import Message, MessageChunk\n",
     "user_input = input(\"User Input: \")\n",
@@ -204,12 +202,12 @@
    "metadata": {
     "collapsed": false
    },
-   "id": "b26e30b79b71697a"
+   "id": "b26e30b79b71697a",
+   "outputs": [],
+   "execution_count": null
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
    "source": [
     "user_input = input(\"User Input: \")\n",
     "while user_input.strip() and user_input != \"q\":\n",
@@ -245,12 +243,12 @@
    "metadata": {
     "collapsed": false
    },
-   "id": "c7d73e0b138e3eba"
+   "id": "c7d73e0b138e3eba",
+   "outputs": [],
+   "execution_count": null
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
    "source": [
     "# list messages\n",
     "messages = taskingai.assistant.list_messages(\n",
@@ -264,12 +262,12 @@
    "metadata": {
     "collapsed": false
    },
-   "id": "e94e3adb0d15373b"
+   "id": "e94e3adb0d15373b",
+   "outputs": [],
+   "execution_count": null
   },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
    "source": [
     "# delete assistant\n",
     "taskingai.assistant.delete_assistant(\n",
@@ -279,7 +277,9 @@
    "metadata": {
     "collapsed": false
    },
-   "id": "ed39836bbfdc7a4e"
+   "id": "ed39836bbfdc7a4e",
+   "outputs": [],
+   "execution_count": null
   }
  ],
  "metadata": {
diff --git a/taskingai/assistant/assistant.py b/taskingai/assistant/assistant.py
index 88cff13..ea20ff7 100644
--- a/taskingai/assistant/assistant.py
+++ b/taskingai/assistant/assistant.py
@@ -7,6 +7,10 @@
     "AssistantTool",
     "AssistantRetrieval",
     "AssistantToolType",
+    "ToolRef",
+    "ToolType",
+    "RetrievalRef",
+    "RetrievalType",
     "AssistantRetrievalType",
     "get_assistant",
     "list_assistants",