llama_index_openai_agent_tracing_tutorial.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "J2wF3jp6Y-XV"
      },
      "source": [
        "<center>\n",
        "    <p style=\"text-align:center\">\n",
        "        <img alt=\"phoenix logo\" src=\"https://storage.googleapis.com/arize-assets/phoenix/assets/phoenix-logo-light.svg\" width=\"200\"/>\n",
        "        <br>\n",
        "        <a href=\"https://docs.arize.com/phoenix/\">Docs</a>\n",
        "        |\n",
        "        <a href=\"https://github.com/Arize-ai/phoenix\">GitHub</a>\n",
        "        |\n",
        "        <a href=\"https://join.slack.com/t/arize-ai/shared_invite/zt-1px8dcmlf-fmThhDFD_V_48oU7ALan4Q\">Community</a>\n",
        "    </p>\n",
        "</center>\n",
        "<h1 align=\"center\">Tracing and Evaluating a LlamaIndex OpenAI Agent Application</h1>\n",
        "\n",
        "With the new OpenAI API that supports function calling, it’s never been easier to build your own agent.\n",
        "\n",
        "In this notebook tutorial, we showcase how to write your own OpenAI agent in under 50 lines of code and use Phoenix to inspect the internals of the Agent. It is minimal, yet feature complete (with ability to carry on a conversation and use tools)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "aeDtvP2RY-XZ"
      },
      "source": [
        "Install LlamaIndex and other dependencies."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YgR53aepY-XZ"
      },
      "outputs": [],
      "source": [
        "!pip install -qq arize-phoenix llama-index openai"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "collapsed": false,
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "MUAKee0GY-Xa"
      },
      "source": [
        "Import libraries."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "FV7j681sY-Xb"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "from getpass import getpass\n",
        "\n",
        "import openai\n",
        "import pandas as pd\n",
        "import phoenix as px\n",
        "from llama_index.agent import OpenAIAgent\n",
        "from llama_index.callbacks import CallbackManager\n",
        "from llama_index.llms import OpenAI\n",
        "from llama_index.prompts.system import SHAKESPEARE_WRITING_ASSISTANT\n",
        "from llama_index.tools import FunctionTool\n",
        "from phoenix.trace.llama_index import (\n",
        "    OpenInferenceTraceCallbackHandler,\n",
        ")\n",
        "\n",
        "pd.set_option(\"display.max_colwidth\", 1000)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "UcC0QFn8Y-Xb"
      },
      "source": [
        "You can run Phoenix in the background to collect trace data emitted by any LlamaIndex application that has been instrumented with the `OpenInferenceTraceCallbackHandler`.\n",
        "\n",
        "Launch Phoenix and follow the instructions in the cell output to open the Phoenix UI (the UI should be empty because we have yet to run a LlamaIndex application)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0rHmf6gMY-Xb"
      },
      "outputs": [],
      "source": [
        "session = px.launch_app()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "xFiU6G6dY-Xc"
      },
      "source": [
        "Let’s start by importing some building blocks and defining tools that our agent can use"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ADDyCKHLY-Xc"
      },
      "outputs": [],
      "source": [
        "def multiply(a: int, b: int) -> int:\n",
        "    \"\"\"Multiple two integers and returns the result integer\"\"\"\n",
        "    return a * b\n",
        "\n",
        "\n",
        "multiply_tool = FunctionTool.from_defaults(fn=multiply)\n",
        "\n",
        "\n",
        "def add(a: int, b: int) -> int:\n",
        "    \"\"\"Add two integers and returns the result integer\"\"\"\n",
        "    return a + b\n",
        "\n",
        "\n",
        "add_tool = FunctionTool.from_defaults(fn=add)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "collapsed": false,
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "t1bgzD2uY-Xd"
      },
      "source": [
        "Provide your API keys to access Open AI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nLE1jmfxY-Xd"
      },
      "outputs": [],
      "source": [
        "if not (openai_api_key := os.getenv(\"OPENAI_API_KEY\")):\n",
        "    openai_api_key = getpass(\"🔑 Enter your OpenAI API key: \")\n",
        "openai.api_key = openai_api_key\n",
        "os.environ[\"OPENAI_API_KEY\"] = openai_api_key"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "-tTRnRzvY-Xd"
      },
      "source": [
        "Now, we define our agent that’s capable of holding a conversation and calling tools.\n",
        "\n",
        "The meat of the agent logic is in the chat method. At a high-level, there are 3 steps:\n",
        "\n",
        "- Call OpenAI to decide which tool (if any) to call and with what arguments.\n",
        "\n",
        "- Call the tool with the arguments to obtain an output\n",
        "\n",
        "- Call OpenAI to synthesize a response from the conversation context and the tool output.\n",
        "\n",
        "The reset method resets the conversation context, so we can start another conversation.\n",
        "\n",
        "For fun, let's make the agent chat in the style of Shakespeare."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lWCP-ShKY-Xd"
      },
      "outputs": [],
      "source": [
        "llm = OpenAI(model=\"gpt-3.5-turbo-0613\")\n",
        "callback_handler = OpenInferenceTraceCallbackHandler()\n",
        "callback_manager = CallbackManager(handlers=[callback_handler])\n",
        "agent = OpenAIAgent.from_tools(\n",
        "    [multiply_tool, add_tool],\n",
        "    llm=llm,\n",
        "    callback_manager=callback_manager,\n",
        "    system_prompt=SHAKESPEARE_WRITING_ASSISTANT,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "6xGrJa3cY-Xe"
      },
      "source": [
        "Let's now chat with our agent!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Im_tHTDXY-Xe"
      },
      "outputs": [],
      "source": [
        "response = agent.query(\"What is (121 * 3) + 42?\")\n",
        "print(response)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "poOzlPrZY-Xe"
      },
      "source": [
        "Let's chat with our agent a few more times. This time with some follow-up questions."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aelp5OBsY-Xe"
      },
      "outputs": [],
      "source": [
        "queries = [\n",
        "    \"What is (121 * 3) + 42?\",\n",
        "    \"what is 3 * 3?\",\n",
        "    \"what is 4 * 4?\",\n",
        "    \"what is 75 * (3 + 4)?\",\n",
        "    \"what is 23 times 87\",\n",
        "]\n",
        "\n",
        "for query in queries:\n",
        "    print(f\"> {query}\")\n",
        "    response = agent.query(query)\n",
        "    print(response)\n",
        "    agent.reset()\n",
        "    print(\"---\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "q6DAtkqeY-Xe"
      },
      "source": [
        "Open the `session.url` in your browser to take a look at the traces in Phoenix. Note that LLM spans contain the OpenAI function calls, and that we can inspect what tool the LLM picked based on the queries.\n",
        "\n",
        "To learn more about function calling, check out the [OpenAI API docs](https://openai.com/blog/function-calling-and-other-api-updates).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iKibGkI7Y-Xe"
      },
      "outputs": [],
      "source": [
        "print(f\"Open the Phoenix UI if you haven't already: {session.url}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "pycharm": {
          "name": "#%% md\n"
        },
        "id": "yvZevFhpY-Xf"
      },
      "source": [
        "We can also inspect the agent's chat history as a dataframe."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TVjSh9R-Y-Xf"
      },
      "outputs": [],
      "source": [
        "ds = px.TraceDataset.from_spans(list(callback_handler.get_spans()))\n",
        "ds.dataframe.head()"
      ]
    }
  ],
  "metadata": {
    "language_info": {
      "name": "python"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}