refactor API (#56)

Quansight · Oct 9, 2023 · c2963c3 · c2963c3
1 parent 03721c7
commit c2963c3
Show file tree

Hide file tree

Showing 12 changed files with 512 additions and 386 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -6,7 +6,7 @@ on:
       - ".github/**"
       - "ragna/**"
       - "tests/**"
-      - "environment.yaml"
+      - "environment.yml"
       - "pyproject.toml"
 
 jobs:

diff --git a/examples/rest_api/rest_api.ipynb b/examples/rest_api/rest_api.ipynb
diff --git a/examples/s3_documents/s3_documents.ipynb b/examples/s3_documents/s3_documents.ipynb
@@ -419,7 +419,25 @@
    "execution_count": 5,
    "id": "ae31bfb0-01f6-47c8-9733-ba1204b153d8",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:     127.0.0.1:56670 - \"GET /health HTTP/1.1\" 200 OK\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:     Started server process [40064]\n",
+      "INFO:     Waiting for application startup.\n",
+      "INFO:     Application startup complete.\n",
+      "INFO:     Uvicorn running on http://127.0.0.1:31476 (Press CTRL+C to quit)\n"
+     ]
+    }
+   ],
    "source": [
     "import contextlib\n",
     "import subprocess\n",
@@ -473,10 +491,11 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'data': {'token': 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyIjoiUmFnbmEiLCJpZCI6ImQxYTgyYmM1LTFlODItNDllZi1hODI4LTNlMTIzZjE1OTIwZiIsImV4cCI6MTY5NTg0NDQ4NC43MTYyMDc1fQ._Q8sUpMZg0NeJ9MOJwZ6qZ5QfBm_jD1Ziwy7PK9Vmsw'},\n",
-      " 'document': {'id': 'd1a82bc5-1e82-49ef-a828-3e123f15920f',\n",
+      "INFO:     127.0.0.1:56670 - \"GET /document?user=Ragna&name=document0.txt HTTP/1.1\" 200 OK\n",
+      "{'data': {'token': 'eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJ1c2VyIjoiUmFnbmEiLCJpZCI6ImViY2E1NTZhLTQxZTgtNDQ4Ni1iOGI3LThkMmM3OGFlOTc1OCIsImV4cCI6MTY5NjQ1NjA2NS41ODAyMDUyfQ.MwFbgwv2nTcetRl9osgjnsnplvNC8fIxHBSBadFX7tg'},\n",
+      " 'document': {'id': 'ebca556a-41e8-4486-b8b7-8d2c78ae9758',\n",
       "              'name': 'document0.txt'},\n",
-      " 'url': 'http://127.0.0.1:31476/document/upload'}\n"
+      " 'url': 'http://127.0.0.1:31476/document'}\n"
      ]
     }
    ],
@@ -485,9 +504,7 @@
     "\n",
     "path = paths[0]\n",
     "\n",
-    "response = await client.get(\n",
-    "    f\"{URL}/document/new\", params={\"user\": USER, \"name\": path.name}\n",
-    ")\n",
+    "response = await client.get(f\"{URL}/document\", params={\"user\": USER, \"name\": path.name})\n",
     "document_info = response.json()\n",
     "document = document_info[\"document\"]\n",
     "pprint(document_info)"
@@ -498,7 +515,15 @@
    "execution_count": 8,
    "id": "fdc5c6f4-bcf8-4902-b441-768b6e3a66bc",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:     127.0.0.1:56670 - \"POST /document HTTP/1.1\" 200 OK\n"
+     ]
+    }
+   ],
    "source": [
     "response = await client.post(\n",
     "    document_info[\"url\"],\n",
@@ -514,12 +539,22 @@
    "id": "f0d769ff-d721-460a-8888-5b033eb4a909",
    "metadata": {},
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:     127.0.0.1:56670 - \"GET /document?user=Ragna&name=document1.txt HTTP/1.1\" 200 OK\n",
+      "INFO:     127.0.0.1:56670 - \"POST /document HTTP/1.1\" 200 OK\n",
+      "INFO:     127.0.0.1:56670 - \"GET /document?user=Ragna&name=document2.txt HTTP/1.1\" 200 OK\n",
+      "INFO:     127.0.0.1:56670 - \"POST /document HTTP/1.1\" 200 OK\n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
-       "[{'id': 'd1a82bc5-1e82-49ef-a828-3e123f15920f', 'name': 'document0.txt'},\n",
-       " {'id': '4258a48a-1e79-4f8f-a393-1d91bdf989aa', 'name': 'document1.txt'},\n",
-       " {'id': 'a30ef295-d420-4076-82a0-145b4315f2c4', 'name': 'document2.txt'}]"
+       "[{'id': 'ebca556a-41e8-4486-b8b7-8d2c78ae9758', 'name': 'document0.txt'},\n",
+       " {'id': 'a723700a-8401-4b09-b727-ead6797f9725', 'name': 'document1.txt'},\n",
+       " {'id': '9028a6c3-2a07-436e-a953-7b92a393113b', 'name': 'document2.txt'}]"
       ]
      },
      "execution_count": 9,
@@ -532,9 +567,7 @@
     "\n",
     "for path in paths[1:]:\n",
     "    document_info = (\n",
-    "        await client.get(\n",
-    "            f\"{URL}/document/new\", params={\"user\": USER, \"name\": path.name}\n",
-    "        )\n",
+    "        await client.get(f\"{URL}/document\", params={\"user\": USER, \"name\": path.name})\n",
     "    ).json()\n",
     "    documents.append(document_info[\"document\"])\n",
     "    await client.post(\n",
@@ -552,10 +585,17 @@
    "id": "ae364aec-e63b-4f00-8b55-217148d6df24",
    "metadata": {},
    "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:     127.0.0.1:56670 - \"POST /chats?user=Ragna HTTP/1.1\" 200 OK\n"
+     ]
+    },
     {
      "data": {
       "text/plain": [
-       "'http://127.0.0.1:31476/chat/0352c50a-cec0-4708-bb7c-4eec76e0a5f0'"
+       "'http://127.0.0.1:31476/chats/98a82c99-a221-4ff5-8b24-cc42ced36477'"
       ]
      },
      "execution_count": 10,
@@ -566,18 +606,19 @@
    "source": [
     "chat = (\n",
     "    await client.post(\n",
-    "        f\"{URL}/chat/new\",\n",
+    "        f\"{URL}/chats\",\n",
     "        params={\"user\": USER},\n",
     "        json={\n",
     "            \"name\": \"Ragna REST API example\",\n",
     "            \"document_ids\": [d[\"id\"] for d in documents],\n",
     "            \"source_storage\": \"Ragna/DemoSourceStorage\",\n",
     "            \"assistant\": \"Ragna/DemoAssistant\",\n",
+    "            \"params\": {},\n",
     "        },\n",
     "    )\n",
     ").json()\n",
     "\n",
-    "CHAT_URL = f\"{URL}/chat/{chat['id']}\"\n",
+    "CHAT_URL = f\"{URL}/chats/{chat['id']}\"\n",
     "CHAT_URL"
    ]
   },
@@ -591,14 +632,16 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "I just pretend to be an LLM. I can't actually help with your prompt:\n",
+      "INFO:     127.0.0.1:56670 - \"POST /chats/98a82c99-a221-4ff5-8b24-cc42ced36477/start?user=Ragna HTTP/1.1\" 200 OK\n",
+      "INFO:     127.0.0.1:56670 - \"POST /chats/98a82c99-a221-4ff5-8b24-cc42ced36477/answer?user=Ragna&prompt=Hello%20World%21 HTTP/1.1\" 200 OK\n",
+      "I can't really help you with your prompt:\n",
       "\n",
       "> Hello World!\n",
       "\n",
-      "I was given the following sources:\n",
+      "I can at least show you the sources that I was given:\n",
       "\n",
-      "- document1.txt: This is content of document 1 located on S3\n",
       "- document2.txt: This is content of document 2 located on S3\n",
+      "- document1.txt: This is content of document 1 located on S3\n",
       "- document0.txt: This is content of document 0 located on S3\n"
      ]
     }

diff --git a/pyproject.toml b/pyproject.toml
@@ -18,7 +18,6 @@ classifiers = [
 ]
 requires-python = ">=3.9"
 dependencies = [
-    "cloudpickle",
     "huey",
     "importlib_metadata >=4.6; python_version < '3.10'",
     "packaging",
@@ -34,6 +33,7 @@ console = [
     "rich",
 ]
 api = [
+    "aiofiles",
     "fastapi",
     "PyJWT",
     "python-multipart",
@@ -52,6 +52,7 @@ builtin-components = [
 ]
 complete = [
     "PyJWT",
+    "aiofiles",
     "chromadb>=0.4",
     "fastapi",
     "httpx",