Feat/add retrieval apis

examples/assistant/chat_with_assistant.ipynb

@@ -42,7 +42,7 @@
    "execution_count": null,
    "outputs": [],
    "source": [
-    "from taskingai.tool import Action, ActionAuthentication, ActionAuthenticationType\n",
+    "from taskingai.models import Action, ActionAuthentication, ActionAuthenticationType\n",
     "from typing import List\n",
     "\n",
     "# create an assistant action\n",
@@ -109,7 +109,7 @@
    "execution_count": null,
    "outputs": [],
    "source": [
-    "from taskingai.assistant import Assistant, Chat, AssistantTool, AssistantToolType\n",
+    "from taskingai.models import Assistant, Chat, ToolRef, ToolType\n",
     "from taskingai.assistant.memory import AssistantMessageWindowMemory\n",
     "\n",
     "# choose an available chat_completion model from your project\n",
@@ -127,12 +127,18 @@
     "        \"You know the meaning of various numbers.\",\n",
     "        \"No matter what the user's language is, you will use the {{langugae}} to explain.\"\n",
     "    ],\n",
-    "    tools=[AssistantTool(\n",
-    "        type=AssistantToolType.ACTION,\n",
-    "        id=action.action_id,\n",
-    "    )],\n",
+    "    tools=[\n",
+    "        ToolRef(\n",
+    "            type=ToolType.ACTION,\n",
+    "            id=action.action_id,\n",
+    "        ), \n",
+    "        ToolRef(\n",
+    "            type=ToolType.ACTION,\n",
+    "            id=action.action_id,\n",
+    "        )\n",
+    "    ],\n",
     "    retrievals=[],\n",
-    "    metadata={\"foo\": \"bar\"},\n",
+    "    metadata={\"k\": \"v\"},\n",
     ")\n",
     "print(f\"created assistant: {assistant}\\n\")"
    ],
@@ -208,12 +214,12 @@
     "user_input = input(\"User Input: \")\n",
     "while user_input.strip() and user_input != \"q\":\n",
     "    # create user message\n",
-    "    taskingai.assistant.create_message(\n",
+    "    user_message = taskingai.assistant.create_message(\n",
     "        assistant_id=assistant.assistant_id,\n",
     "        chat_id=chat.chat_id,\n",
     "        text=user_input,\n",
     "    )\n",
-    "    print(f\"User: {user_input}\")\n",
+    "    print(f\"User: {user_input} ({user_message.message_id})\")\n",
     "    \n",
     "    # generate assistant response\n",
     "    assistant_message_response = taskingai.assistant.generate_message(\n",
@@ -230,8 +236,8 @@
     "        if isinstance(item, MessageChunk):\n",
     "            print(item.delta, end=\"\", flush=True)\n",
     "        elif isinstance(item, Message):\n",
-    "            print(f\"\\nmessage_id: {item.message_id}\")\n",
-    "            \n",
+    "            print(f\" ({item.message_id})\")\n",
+    "    \n",
     "    time.sleep(2)\n",
     "    # quit by input 'q\n",
     "    user_input = input(\"User: \")"
@@ -274,16 +280,6 @@
     "collapsed": false
    },
    "id": "ed39836bbfdc7a4e"
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "outputs": [],
-   "source": [],
-   "metadata": {
-    "collapsed": false
-   },
-   "id": "9cfed1128acbdd30"
   }
  ],
  "metadata": {

examples/crud/assistant_crud.ipynb

@@ -25,7 +25,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from taskingai.assistant import Assistant, Chat\n",
+    "from taskingai.models import Assistant, Chat\n",
     "from taskingai.assistant.memory import AssistantNaiveMemory\n",
     "\n",
     "# choose an available chat_completion model from your project\n",
@@ -166,7 +166,7 @@
     "    assistant_id=assistant.assistant_id,\n",
     "    chat_id=chat_id,\n",
     ")\n",
-    "print(f\"got chat: {chat}\\n\")"
+    "print(f\"chat: {chat}\\n\")"
    ],
    "metadata": {
     "collapsed": false

examples/crud/retrieval_crud.ipynb

@@ -28,10 +28,10 @@
    "execution_count": null,
    "outputs": [],
    "source": [
-    "from taskingai.retrieval import Collection, Record, TokenTextSplitter\n",
+    "from taskingai.retrieval import Collection, Record, Chunk, TokenTextSplitter\n",
     "\n",
     "# choose an available text_embedding model from your project\n",
-    "embedding_model_id = \"YOUR_MODEL_ID\""
+    "embedding_model_id = \"YOUR_EMBEDDING_MODEL_ID\""
    ],
    "metadata": {
     "collapsed": false
@@ -93,7 +93,7 @@
     "    collection_id=collection_id\n",
     ")\n",
     "\n",
-    "print(f\"got collection: {collection}\\n\")"
+    "print(f\"collection: {collection}\\n\")"
    ],
    "metadata": {
     "collapsed": false
@@ -212,11 +212,11 @@
    "execution_count": null,
    "outputs": [],
    "source": [
-    "# update record\n",
+    "# update record - metadata\n",
     "record = taskingai.retrieval.update_record(\n",
     "    collection_id=collection.collection_id,\n",
     "    record_id=record.record_id,\n",
-    "    metadata={\"foo\": \"bar\"}, # currently only metadata update is supported\n",
+    "    metadata={\"foo\": \"bar\"},\n",
     ")\n",
     "print(f\"updated record: {record}\")"
    ],
@@ -225,6 +225,25 @@
    },
    "id": "65d833b22e1e657"
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# update record - content\n",
+    "record = taskingai.retrieval.update_record(\n",
+    "    collection_id=collection.collection_id,\n",
+    "    record_id=record.record_id,\n",
+    "    content=\"New content\",\n",
+    "    text_splitter=TokenTextSplitter(chunk_size=100, chunk_overlap=20),\n",
+    ")\n",
+    "print(f\"updated record: {record}\")"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "4369989d2bd1a777"
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -257,6 +276,154 @@
     "collapsed": false
    },
    "id": "accf6d883fcffaa8"
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## Chunk Object"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "b0e4c12fb7509fea"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# create a new text record\n",
+    "chunk: Chunk = taskingai.retrieval.create_chunk(\n",
+    "    collection_id=collection.collection_id,\n",
+    "    content=\"The dog is a domesticated descendant of the wolf. Also called the domestic dog, it is derived from extinct gray wolves, and the gray wolf is the dog's closest living relative. The dog was the first species to be domesticated by humans.\",\n",
+    ")\n",
+    "print(f\"created chunk: {chunk.chunk_id} for collection: {collection.collection_id}\\n\")"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "a395337f136500fc"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# update chunk metadata\n",
+    "chunk = taskingai.retrieval.update_chunk(\n",
+    "    collection_id=collection.collection_id,\n",
+    "    chunk_id=chunk.chunk_id,\n",
+    "    metadata={\"k\": \"v\"},\n",
+    ")\n",
+    "print(f\"updated chunk: {chunk}\")"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "309e1771251bb079"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# update chunk content\n",
+    "chunk = taskingai.retrieval.update_chunk(\n",
+    "    collection_id=collection.collection_id,\n",
+    "    chunk_id=chunk.chunk_id,\n",
+    "    content=\"New content\",\n",
+    ")\n",
+    "print(f\"updated chunk: {chunk}\")"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "a9d68db12329b558"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# get chunk\n",
+    "chunk = taskingai.retrieval.get_chunk(\n",
+    "    collection_id=collection.collection_id,\n",
+    "    chunk_id=chunk.chunk_id\n",
+    ")\n",
+    "print(f\"got chunk: {chunk}\\n\")"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "d3899097cd6d0cf2"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# delete chunk\n",
+    "taskingai.retrieval.delete_chunk(\n",
+    "    collection_id=collection.collection_id,\n",
+    "    chunk_id=chunk.chunk_id,\n",
+    ")\n",
+    "print(f\"deleted chunk {chunk.chunk_id} from collection {collection.collection_id}\\n\")"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "27e643ad8e8636ed"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# create a new text record and a new chunk\n",
+    "taskingai.retrieval.create_record(\n",
+    "    collection_id=collection.collection_id,\n",
+    "    content=\"Machine learning is a subfield of artificial intelligence (AI) that involves the development of algorithms that allow computers to learn from and make decisions or predictions based on data. The term \\\"machine learning\\\" was coined by Arthur Samuel in 1959. In other words, machine learning enables a system to automatically learn and improve from experience without being explicitly programmed. This is achieved by feeding the system massive amounts of data, which it uses to learn patterns and make inferences. There are three main types of machine learning: 1. Supervised Learning: This is where the model is given labeled training data and the goal of learning is to generalize from the training data to unseen situations in a principled way. 2. Unsupervised Learning: This involves training on a dataset without explicit labels. The goal might be to discover inherent groupings or patterns within the data. 3. Reinforcement Learning: In this type, an agent learns to perform actions based on reward/penalty feedback to achieve a goal. It's commonly used in robotics, gaming, and navigation. Deep learning, a subset of machine learning, uses neural networks with many layers (\\\"deep\\\" structures) and has been responsible for many recent breakthroughs in AI, including speech recognition, image recognition, and natural language processing. It's important to note that machine learning is a rapidly developing field, with new techniques and applications emerging regularly.\",\n",
+    "    text_splitter=TokenTextSplitter(chunk_size=200, chunk_overlap=20)\n",
+    ")\n",
+    "\n",
+    "taskingai.retrieval.create_chunk(\n",
+    "    collection_id=collection.collection_id,\n",
+    "    content=\"The dog is a domesticated descendant of the wolf. Also called the domestic dog, it is derived from extinct gray wolves, and the gray wolf is the dog's closest living relative. The dog was the first species to be domesticated by humans.\",\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "a74dd7615ec28528"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# list chunks\n",
+    "chunks = taskingai.retrieval.list_chunks(collection_id=collection.collection_id)\n",
+    "for chunk in chunks:\n",
+    "    print(chunk)\n",
+    "    print(\"-\" * 50)"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "55e9645ac41f8ca"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [
+    "# delete collection\n",
+    "taskingai.retrieval.delete_collection(collection_id=collection.collection_id)"
+   ],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "b97aaa156f586e34"
   }
  ],
  "metadata": {

examples/crud/tool_crud.ipynb

@@ -197,6 +197,16 @@
     "collapsed": false
    },
    "id": "5a1a36d15055918f"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "5588d5e7457225be"
   }
  ],
  "metadata": {

examples/retrieval/semantic_search.ipynb

@@ -41,7 +41,7 @@
    "outputs": [],
    "source": [
     "# choose an available text_embedding model from your project\n",
-    "embedding_model_id = \"YOUR_MODEL_ID\""
+    "embedding_model_id = \"YOUR_EMBEDDING_MODEL_ID\""
    ],
    "metadata": {
     "collapsed": false
@@ -224,6 +224,16 @@
     "collapsed": false
    },
    "id": "fc9c1fa12d893dd1"
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "outputs": [],
+   "source": [],
+   "metadata": {
+    "collapsed": false
+   },
+   "id": "e0eb39fcac309768"
   }
  ],
  "metadata": {