examples/notebooks/simpleaichat_async.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "tFuqw725l_4A"
      },
      "source": [
        "# Async ChatGPT with simpleaichat\n",
        "\n",
        "simpleaichat has an `AsyncAIChat` class which allows all requests to be async. This works for normal chats, streaming, and tools!\n",
        "\n",
        "With that, you can implement it into an async webapp such as FastAPI, or query multiple requests at the same time. And because `AsyncAIChat` is a session manager, you can conduct independent chat sessions!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "JuI3NyYylPK6"
      },
      "outputs": [],
      "source": [
        "!pip install -q simpleaichat\n",
        "\n",
        "from simpleaichat import AsyncAIChat\n",
        "from getpass import getpass\n",
        "import asyncio"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "OSp2OXNRmi1j"
      },
      "source": [
        "For the following cell, input your OpenAI API key when prompted. **It will not be saved to the notebook**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "as8yMGC8mo6I",
        "outputId": "8f7ea9e9-364b-4f2e-c19d-8147630dc0ac"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "OpenAI Key: ··········\n"
          ]
        }
      ],
      "source": [
        "api_key = getpass(\"OpenAI Key: \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "id": "NK-f4AJElPK7"
      },
      "outputs": [],
      "source": [
        "ai = AsyncAIChat(api_key=api_key, console=False)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "0cTgt37umrms"
      },
      "source": [
        "## Async Generation\n",
        "\n",
        "Async calls are typical async, with an `await` keyword."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "R80NS8jflPK7",
        "outputId": "e34f515b-d3c1-40e0-e319-0c0fa3e20fdf"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "The capital of California is Sacramento.\n"
          ]
        }
      ],
      "source": [
        "response = await ai(\"What is the capital of California?\")\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "_ZxNQq8alPK8",
        "outputId": "765697cb-7cb7-4855-a894-1f6261039e87"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Sacramento was founded on February 27, 1850.\n"
          ]
        }
      ],
      "source": [
        "response = await ai(\"When was it founded?\")\n",
        "print(response)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ivvJNrErn6aJ"
      },
      "source": [
        "Now, let's ask for multiple distinct states, at the same time. It will take roughly the same amount of time to complete as a single state!\n",
        "\n",
        "To do that, we create a session for each input state:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 34,
      "metadata": {
        "id": "6NeY0Ls_lPK9"
      },
      "outputs": [],
      "source": [
        "states = [\"Washington\", \"New Mexico\", \"Texas\", \"Mississippi\", \"Alaska\"]\n",
        "\n",
        "ai_2 = AsyncAIChat(api_key=api_key, console=False)\n",
        "for state in states:\n",
        "    ai_2.new_session(api_key=api_key, id=state)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "FyxsQi4RoKhF"
      },
      "source": [
        "Each call creates a task coroutine; we can store the tasks, then run them all with `asyncio.gather`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 35,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6v7V_jM_lPLA",
        "outputId": "b83589d3-76de-4c28-fad6-43e1c37a6f42"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['The capital of Washington is Olympia.',\n",
              " 'The capital of New Mexico is Santa Fe.',\n",
              " 'The capital of Texas is Austin.',\n",
              " 'The capital of Mississippi is Jackson.',\n",
              " 'The capital of Alaska is Juneau.']"
            ]
          },
          "execution_count": 35,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tasks = []\n",
        "for state in states:\n",
        "    tasks.append(ai_2(f\"What is the capital of {state}?\", id=state))\n",
        "\n",
        "results = await asyncio.gather(*tasks)\n",
        "results"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "xJTvRG2YpmPN"
      },
      "source": [
        "Now, to ask the same question to all states:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 36,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fOggtHsiptcI",
        "outputId": "2208559d-57e7-42aa-d821-41ff058c4d57"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['Olympia was founded in 1853.',\n",
              " 'Santa Fe was founded in 1610, making it the oldest state capital in the United States.',\n",
              " 'Austin was founded on December 27, 1839.',\n",
              " 'Jackson was founded on December 23, 1821.',\n",
              " 'Juneau was founded on October 18, 1880.']"
            ]
          },
          "execution_count": 36,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tasks = []\n",
        "for state in states:\n",
        "    tasks.append(ai_2(\"When was it founded?\", id=state))\n",
        "\n",
        "results = await asyncio.gather(*tasks)\n",
        "results"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ywVjhqJHoYqb"
      },
      "source": [
        "Indeed, the messages are stored correctly by session, and are still independent between sessions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 37,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "NUfgFE4TlPLB",
        "outputId": "38a1b0d7-9603-42a9-b66a-a4b5e76e957c"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[What is the capital of Washington?,\n",
              " The capital of Washington is Olympia.,\n",
              " When was it founded?,\n",
              " Olympia was founded in 1853.]"
            ]
          },
          "execution_count": 37,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "ai_2.sessions[\"Washington\"].messages"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zzbXUhgtlPLB",
        "outputId": "44fb73e6-902a-4e9c-a584-b2ea8943bbba"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[What is the capital of Texas?,\n",
              " The capital of Texas is Austin.,\n",
              " When was it founded?,\n",
              " Austin was founded on December 27, 1839.]"
            ]
          },
          "execution_count": 38,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "ai_2.sessions[\"Texas\"].messages"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "n47XlrZmojU5"
      },
      "source": [
        "## Async Streaming\n",
        "\n",
        "Now, let's do the same thing, except with streaming."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "id": "cgJaPdlvoar9"
      },
      "outputs": [],
      "source": [
        "ai = AsyncAIChat(api_key=api_key, console=False)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "kjyUp_PZotON"
      },
      "source": [
        "In this case, you need an async generator for the streaming call."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "niO2SZN4owuZ",
        "outputId": "d3c28d6d-fa59-40ca-a903-e3ce8ba21ebc"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{'delta': 'The', 'response': 'The'}\n",
            "{'delta': ' capital', 'response': 'The capital'}\n",
            "{'delta': ' of', 'response': 'The capital of'}\n",
            "{'delta': ' California', 'response': 'The capital of California'}\n",
            "{'delta': ' is', 'response': 'The capital of California is'}\n",
            "{'delta': ' Sacramento', 'response': 'The capital of California is Sacramento'}\n",
            "{'delta': '.', 'response': 'The capital of California is Sacramento.'}\n"
          ]
        }
      ],
      "source": [
        "async for chunk in await ai.stream(\"What is the capital of California?\"):\n",
        "    print(chunk)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "4um4Mv05o4OB"
      },
      "source": [
        "For multistate generation:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "metadata": {
        "id": "Y4G3VO6qo5yy"
      },
      "outputs": [],
      "source": [
        "states = [\"Washington\", \"New Mexico\", \"Texas\", \"Mississippi\", \"Alaska\"]\n",
        "\n",
        "ai_2 = AsyncAIChat(api_key=api_key, console=False)\n",
        "for state in states:\n",
        "    ai_2.new_session(api_key=api_key, id=state)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "oewS5NAVpNwT"
      },
      "source": [
        "This implementation is slightly more complicated since you need to wrap each async generator in its own async function. However, it provides the best demonstration of async, as you can clearly see that each chunk is received in a different order."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 40,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HVrZ16aXpAYm",
        "outputId": "9015c66a-fbf1-4cb6-bbbf-ef62b1a6bf79"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "{'delta': 'The', 'response': 'The'}\n",
            "{'delta': 'The', 'response': 'The'}\n",
            "{'delta': ' capital', 'response': 'The capital'}\n",
            "{'delta': ' capital', 'response': 'The capital'}\n",
            "{'delta': ' of', 'response': 'The capital of'}\n",
            "{'delta': ' of', 'response': 'The capital of'}\n",
            "{'delta': 'The', 'response': 'The'}\n",
            "{'delta': 'The', 'response': 'The'}\n",
            "{'delta': ' Washington', 'response': 'The capital of Washington'}\n",
            "{'delta': ' Mississippi', 'response': 'The capital of Mississippi'}\n",
            "{'delta': 'The', 'response': 'The'}\n",
            "{'delta': ' capital', 'response': 'The capital'}\n",
            "{'delta': ' capital', 'response': 'The capital'}\n",
            "{'delta': ' is', 'response': 'The capital of Washington is'}\n",
            "{'delta': ' is', 'response': 'The capital of Mississippi is'}\n",
            "{'delta': ' of', 'response': 'The capital of'}\n",
            "{'delta': ' capital', 'response': 'The capital'}\n",
            "{'delta': ' of', 'response': 'The capital of'}\n",
            "{'delta': ' Jackson', 'response': 'The capital of Mississippi is Jackson'}\n",
            "{'delta': ' Olympia', 'response': 'The capital of Washington is Olympia'}\n",
            "{'delta': ' Alaska', 'response': 'The capital of Alaska'}\n",
            "{'delta': ' New', 'response': 'The capital of New'}\n",
            "{'delta': '.', 'response': 'The capital of Washington is Olympia.'}\n",
            "{'delta': ' of', 'response': 'The capital of'}\n",
            "{'delta': '.', 'response': 'The capital of Mississippi is Jackson.'}\n",
            "{'delta': ' Mexico', 'response': 'The capital of New Mexico'}\n",
            "{'delta': ' is', 'response': 'The capital of Alaska is'}\n",
            "{'delta': ' Texas', 'response': 'The capital of Texas'}\n",
            "{'delta': ' is', 'response': 'The capital of New Mexico is'}\n",
            "{'delta': ' June', 'response': 'The capital of Alaska is June'}\n",
            "{'delta': ' is', 'response': 'The capital of Texas is'}\n",
            "{'delta': ' Santa', 'response': 'The capital of New Mexico is Santa'}\n",
            "{'delta': ' Austin', 'response': 'The capital of Texas is Austin'}\n",
            "{'delta': 'au', 'response': 'The capital of Alaska is Juneau'}\n",
            "{'delta': ' Fe', 'response': 'The capital of New Mexico is Santa Fe'}\n",
            "{'delta': '.', 'response': 'The capital of Texas is Austin.'}\n",
            "{'delta': '.', 'response': 'The capital of Alaska is Juneau.'}\n",
            "{'delta': '.', 'response': 'The capital of New Mexico is Santa Fe.'}\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "['The capital of Washington is Olympia.',\n",
              " 'The capital of New Mexico is Santa Fe.',\n",
              " 'The capital of Texas is Austin.',\n",
              " 'The capital of Mississippi is Jackson.',\n",
              " 'The capital of Alaska is Juneau.']"
            ]
          },
          "execution_count": 40,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "async def capital_stream(state):\n",
        "    async for chunk in await ai_2.stream(f\"What is the capital of {state}?\", id=state):\n",
        "        response = chunk\n",
        "        print(response)\n",
        "    return response[\"response\"]\n",
        "\n",
        "tasks = []\n",
        "for state in states:\n",
        "    tasks.append(capital_stream(state))\n",
        "\n",
        "results = await asyncio.gather(*tasks)\n",
        "results"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## MIT License\n",
        "\n",
        "Copyright (c) 2023 Max Woolf\n",
        "\n",
        "Permission is hereby granted, free of charge, to any person obtaining a copy\n",
        "of this software and associated documentation files (the \"Software\"), to deal\n",
        "in the Software without restriction, including without limitation the rights\n",
        "to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n",
        "copies of the Software, and to permit persons to whom the Software is\n",
        "furnished to do so, subject to the following conditions:\n",
        "\n",
        "The above copyright notice and this permission notice shall be included in all\n",
        "copies or substantial portions of the Software.\n",
        "\n",
        "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n",
        "IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n",
        "FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n",
        "AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n",
        "LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n",
        "OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n",
        "SOFTWARE.\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.12"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}