openai-cookbook/examples/third_party/Openai_monitoring_with_wandb_weave.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "id": "30ccfdbc",
      "metadata": {
        "id": "30ccfdbc"
      },
      "source": [
        "# OpenAI API Monitoring with W&B Weave\n",
        "\n",
        "<img src=\"http://wandb.me/logo-im-png\" width=\"400\" alt=\"Weights & Biases\" />\n",
        "\n",
        "<!--- @wandbcode{weave_openai_client_qs} -->\n",
        "\n",
        "<a target=\"_blank\" href=\"https://colab.research.google.com/github/wandb/weave/blob/master/examples/prompts/llm_monitoring/openai_client_quickstart.ipynb\">\n",
        "  <img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/>\n",
        "</a>\n",
        "\n",
        "**Note:** you will need an [OpenAI API key](https://platform.openai.com/account/api-keys) to run this colab.\n",
        "\n",
        "Use the W&B OpenAI integration to monitor OpenAI API calls and understand how your projects and teams are leveraging LLMs.\n",
        "In this example, we'll generate templated Weave Boards: LLM usage monitoring dashboards which you can explore and customize from the UI.\n",
        "\n",
        "* automatically track LLM usage and aggregate useful metrics like  cost, latency and throughput across your projects/teams\n",
        "* dynamically query and derive insights from the logs of all your OpenAI API calls\n",
        "* iterate visually to slice, aggregate, and explore your data; customize panels to focus on interesting patterns; share progress more easily with your team through an interactive dashboard\n",
        "\n",
        "<img src=\"https://raw.githubusercontent.com/wandb/weave/master/docs/assets/full_board_view.png\">\n",
        "\n",
        "[Play with a live version of this Weave Board →](http://wandb.me/llm-monitoring-board)\n",
        "\n",
        "#### New to Weights & Biases? [-> Sign up for an account here <-](https://wandb.ai/site)\n",
        "\n",
        "# Step 0: Setup\n",
        "\n",
        "Install dependencies, login to W&B so you can save and share your work, and authenticate with OpenAI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8107732a-fb90-45f8-8377-6381bd28475d",
      "metadata": {
        "id": "8107732a-fb90-45f8-8377-6381bd28475d"
      },
      "outputs": [],
      "source": [
        "# if not already installed\n",
        "!pip install -qqq weave openai tiktoken wandb"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "16686825-ce56-4d56-9207-ef4d2357eaf0",
      "metadata": {
        "id": "16686825-ce56-4d56-9207-ef4d2357eaf0"
      },
      "outputs": [],
      "source": [
        "import wandb\n",
        "wandb.login()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7273a9b0-f41e-4490-b0e9-4ab2cdeb80ec",
      "metadata": {
        "id": "7273a9b0-f41e-4490-b0e9-4ab2cdeb80ec"
      },
      "outputs": [],
      "source": [
        "import weave\n",
        "import os\n",
        "WANDB_BASE_URL = \"https://api.wandb.ai\"\n",
        "os.environ[\"WANDB_BASE_URL\"] = WANDB_BASE_URL"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "de5c26c3-393d-45a0-84a8-b277c42e4724",
      "metadata": {
        "id": "de5c26c3-393d-45a0-84a8-b277c42e4724"
      },
      "outputs": [],
      "source": [
        "# authenticate with OpenAI\n",
        "from getpass import getpass\n",
        "\n",
        "if os.getenv(\"OPENAI_API_KEY\") is None:\n",
        "  os.environ[\"OPENAI_API_KEY\"] = getpass(\"Paste your OpenAI key from: https://platform.openai.com/account/api-keys\\n\")\n",
        "assert os.getenv(\"OPENAI_API_KEY\", \"\").startswith(\"sk-\"), \"This doesn't look like a valid OpenAI API key\"\n",
        "print(\"OpenAI API key configured\")"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "id": "32e91f9c-0640-4346-8c4a-e5f35256afe6",
      "metadata": {
        "id": "32e91f9c-0640-4346-8c4a-e5f35256afe6"
      },
      "source": [
        "# Step 1: Configure data streaming and storage in W&B\n",
        "\n",
        "Set WB_ENTITY to your wandb username or team name. Log in to W&B and navigate to Home Page at [wandb.ai/home](https://wandb.ai/home) to see valid options under your \"Profile\" and \"Teams\" in the left sidebar."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "347bacb8-1351-4517-8116-d11212aef57b",
      "metadata": {
        "id": "347bacb8-1351-4517-8116-d11212aef57b"
      },
      "outputs": [],
      "source": [
        "WB_ENTITY = \"\" # set to your wandb username or team name\n",
        "WB_PROJECT = \"weave\" # top-level directory for this work\n",
        "STREAM_NAME = \"openai_logs\" # record table which stores the logs of OpenAI API calls as they stream in"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "id": "26402897-4f73-4dd6-9314-2a1c60638e4d",
      "metadata": {
        "id": "26402897-4f73-4dd6-9314-2a1c60638e4d"
      },
      "source": [
        "# Step 2: Call init_monitor()\n",
        "\n",
        "To start monitoring OpenAI API usage, call `init_monitor(<stream>)`, where `<stream>` has the form `<wandb_team_or_user>/<wandb_project>/<stream_name>`. The stream records and stores all the OpenAI API calls.\n",
        "\n",
        "Running this cell will print out a link to view the current project in the Weave UI."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ba2b2070",
      "metadata": {
        "id": "ba2b2070"
      },
      "outputs": [],
      "source": [
        "from weave.monitoring import openai, init_monitor\n",
        "m = init_monitor(f\"{WB_ENTITY}/{WB_PROJECT}/{STREAM_NAME}\")\n",
        "\n",
        "# specifying a single model for simplicity\n",
        "OPENAI_MODEL = 'gpt-3.5-turbo'\n",
        "\n",
        "# prefill with some sample logs\n",
        "r = openai.ChatCompletion.create(model=OPENAI_MODEL, messages=[{\"role\": \"user\", \"content\": \"hello world!\"}])\n",
        "r = openai.ChatCompletion.create(model=OPENAI_MODEL, messages=[{\"role\": \"user\", \"content\": \"what is 2+2?\"}])"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "id": "fb51abb4",
      "metadata": {
        "id": "fb51abb4"
      },
      "source": [
        "# Step 3: Preview monitoring dashboard\n",
        "\n",
        "Click on the link above to preview the data stream, then click \"OpenAI Monitor Board\" in the right sidebar to create a Weave Board for this data stream.\n",
        "\n",
        "<img src=\"https://raw.githubusercontent.com/wandb/weave/master/docs/assets/short_board_attempt.gif\" width=75%>\n",
        "\n",
        "# Step 4: Explore & understand your LLM usage\n",
        "\n",
        "To save your work, rename the board by clicking on the autogenerated name at the top of the page. To share your board, click \\\"Publish\\\" in the top right.\n",
        "\n",
        "<img src=\"https://raw.githubusercontent.com/wandb/weave/master/docs/assets/publish_board_short.gif\" width=75%>\n",
        "\n",
        "To visualize your work in real-time as you iterate, you can:\n",
        "* keep the Board open in a separate tab and refresh to view the latest data\n",
        "* rename the Board for easier reference at any point and \\\"Publish\\\" that version to share a link with others\n",
        "* find previously saved Boards by navigating to the relevant W&B entity and W&B project name from weave.wandb.ai\n",
        "* or open a new instance of a Board template to start fresh with all the data accumulated so far\n",
        "\n",
        "\n",
        "Next we'll illustrate a few ways you could track OpenAI API calls. There are many more possibilities depending on your use case, and we can't wait to see what you create from these starter templates.\n",
        "\n",
        "# Examples\n",
        "\n",
        "## Example 0: Log a prompt and its completion\n",
        "\n",
        "Monitor a ChatCompletion request and print the corresponding response, extracting only the text of the completion."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e711d521",
      "metadata": {
        "id": "e711d521"
      },
      "outputs": [],
      "source": [
        "response = openai.ChatCompletion.create(model=OPENAI_MODEL, messages=[\n",
        "        {\"role\": \"user\", \"content\": f\"What is the meaning of life, the universe, and everything?\"},\n",
        "    ])\n",
        "print(response['choices'][0]['message']['content'])"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "id": "4ace8f71-3e50-444b-89e0-a8a619a9e0a8",
      "metadata": {
        "id": "4ace8f71-3e50-444b-89e0-a8a619a9e0a8"
      },
      "source": [
        "## Example 1: Track relevant parameters as attributes\n",
        "\n",
        "Factor out parameters of interest and track them as attributes on the logged record.\n",
        "Here we track the \"system prompt\" separately from the \"prompt template\" and the \"equation\" parameter. This time we'll print the full structured response from the ChatCompletion call."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "7930e774",
      "metadata": {
        "id": "7930e774"
      },
      "outputs": [],
      "source": [
        "system_prompt = \"you always write in bullet points\"\n",
        "prompt_template = 'solve the following equation step by step: {equation}'\n",
        "params = {'equation': '4 * (3 - 1)'}\n",
        "openai.ChatCompletion.create(model=OPENAI_MODEL,\n",
        "                             messages=[\n",
        "                                    {\"role\": \"system\", \"content\": system_prompt},\n",
        "                                    {\"role\": \"user\", \"content\": prompt_template.format(**params)},\n",
        "                                ],\n",
        "                             # you can add additional attributes to the logged record\n",
        "                             # see the monitor_api notebook for more examples\n",
        "                             monitor_attributes={\n",
        "                                 'system_prompt': system_prompt,\n",
        "                                 'prompt_template': prompt_template,\n",
        "                                 'params': params\n",
        "                             })"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "id": "450dc210-3526-4f6e-a05a-df15ebf8d398",
      "metadata": {
        "id": "450dc210-3526-4f6e-a05a-df15ebf8d398"
      },
      "source": [
        "## Example 2: Log an ongoing stream of messages\n",
        "\n",
        "Monitor a stream of messages and log the result as a single record. Note: tokens are not counted in this format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b8b423df",
      "metadata": {
        "id": "b8b423df"
      },
      "outputs": [],
      "source": [
        "from weave.monitoring.openai import message_from_stream\n",
        "r = openai.ChatCompletion.create(model=OPENAI_MODEL, messages=[\n",
        "        {\"role\": \"system\", \"content\": \"You are a robot and only speak in robot, like beep bloop bop.\"},\n",
        "        {\"role\": \"user\", \"content\": \"Tell me a 50-word story.\"},\n",
        "    ], stream=True)\n",
        "for s in message_from_stream(r):\n",
        "    print(s, end='')"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "id": "81945b3c-310e-404b-aea7-a47c6760db2e",
      "metadata": {
        "id": "81945b3c-310e-404b-aea7-a47c6760db2e"
      },
      "source": [
        "## Example 3: Structure prompt engineering experiments\n",
        "\n",
        "Here we compare a few toy options for the system prompt, user question, and intended audience. Try your own experiments and see if any interesting insights emerge as you explore in the Board and group by different parameters."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bfde1485-8d5f-4b25-a878-0c9de719dc6f",
      "metadata": {
        "id": "bfde1485-8d5f-4b25-a878-0c9de719dc6f"
      },
      "outputs": [],
      "source": [
        "def explain_math(system_prompt, prompt_template, params):\n",
        "    openai.ChatCompletion.create(model=OPENAI_MODEL,\n",
        "                             messages=[\n",
        "                                    {\"role\": \"system\", \"content\": system_prompt},\n",
        "                                    {\"role\": \"user\", \"content\": prompt_template.format(**params)},\n",
        "                                ],\n",
        "                             # you can add additional attributes to the logged record\n",
        "                             # see the monitor_api notebook for more examples\n",
        "                             monitor_attributes={\n",
        "                                 'system_prompt': system_prompt,\n",
        "                                 'prompt_template': prompt_template,\n",
        "                                 'params': params\n",
        "                             })"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "48d73400-835e-4e8d-b985-a812cf11b99d",
      "metadata": {
        "id": "48d73400-835e-4e8d-b985-a812cf11b99d"
      },
      "outputs": [],
      "source": [
        "# feel free to substitute your own prompts :)\n",
        "system_prompts = [\"you're extremely flowery and poetic\", \"you're very direct and precise\", \"balance brevity with insight\"]\n",
        "prompt_template = 'explain the solution of the following to a {audience}: {equation}'\n",
        "equations = ['x^2 + 4x + 9 = 0', '15 * (2 - 6) / 4']\n",
        "audience = [\"new student\", \"math genius\"]\n",
        "\n",
        "for system_prompt in system_prompts:\n",
        "    for equation in equations:\n",
        "        for person in audience:\n",
        "            params = {\"equation\" : equation, \"audience\" : person}\n",
        "            explain_math(system_prompt, prompt_template, params)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}