{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CmwwTi0KUa1F"
   },
   "source": [
    "<center>\n",
    "    <p style=\"text-align:center\">\n",
    "        <img alt=\"phoenix logo\" src=\"https://storage.googleapis.com/arize-phoenix-assets/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",
    "\n",
    "\n",
    "\n",
    "# <center>Getting Started with Llamatrace"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "u4-cym_JUfow"
   },
   "source": [
    "This guide demonstrates how to use Llamatrace, a collaboration between Arize and LlamaIndex to deliver hosted observability and evals with native support for LlamaIndex.\n",
    "\n",
    "ℹ️ This notebook requires an OpenAI API key\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BEFoI3zIUwt1"
   },
   "source": [
    "## Step 1: Install Dependencies 📚\n",
    "Let's get the notebook setup with dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "\n",
    "pip install -q \"arize-phoenix>=4.29.0\" openai 'httpx<0.28' openinference-instrumentation-openai\n",
    "pip install -q gcsfs nest-asyncio \"openinference-instrumentation-llama-index>=3.0.0\"\n",
    "pip install -qU llama-index-callbacks-arize-phoenix\n",
    "pip install -qU llama-index llama-index-llms-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "if not (openai_api_key := os.getenv(\"OPENAI_API_KEY\")):\n",
    "    openai_api_key = getpass(\"🔑 Enter your OpenAI API key: \")\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = openai_api_key\n",
    "\n",
    "if not (phoenix_api_key := os.getenv(\"PHOENIX_API_KEY\")):\n",
    "    phoenix_api_key = getpass(\"🔑 Enter your Phoenix API key: \")\n",
    "\n",
    "os.environ[\"PHOENIX_API_KEY\"] = phoenix_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SGTKOk-oU18k"
   },
   "source": [
    "## Step 2: Setup Tracing\n",
    "Let's send a trace to Hosted Phoenix! Note the key lines below using `PHOENIX_CLIENT_HEADERS` and `app.phoenix.arize.com/v1/traces`\n",
    "LlamaIndex has a built-in integration with Phoenix, so we can use their `set_global_handler` method to send traces to Phoenix. You can also use the `register` method from `arize-phoenix-otel` to achieve the same thing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup Arize Phoenix for logging/observability\n",
    "import os\n",
    "\n",
    "from openinference.instrumentation.llama_index import LlamaIndexInstrumentor\n",
    "\n",
    "from phoenix.otel import register\n",
    "\n",
    "os.environ[\"PHOENIX_CLIENT_HEADERS\"] = f\"api_key={os.environ['PHOENIX_API_KEY']}\"\n",
    "os.environ[\"PHOENIX_COLLECTOR_ENDPOINT\"] = \"https://app.phoenix.arize.com\"\n",
    "\n",
    "# Configuration is picked up from your environment variables\n",
    "tracer_provider = register()\n",
    "\n",
    "# Instrument LlamaIndex. This allows Phoenix to collect traces from LlamaIndex queries.\n",
    "LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider, skip_dep_check=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "COGJEi8HoOaB"
   },
   "source": [
    "Setup imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from urllib.request import urlopen\n",
    "\n",
    "import nest_asyncio\n",
    "import pandas as pd\n",
    "from gcsfs import GCSFileSystem\n",
    "from llama_index.core import (\n",
    "    Settings,\n",
    "    StorageContext,\n",
    "    load_index_from_storage,\n",
    ")\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from tqdm import tqdm\n",
    "\n",
    "import phoenix as px\n",
    "\n",
    "nest_asyncio.apply()  # needed for concurrent evals in notebook environments\n",
    "pd.set_option(\"display.max_colwidth\", 1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WxNUMLgwIRIm"
   },
   "source": [
    "Run LlamaIndex Query on Arize Docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_system = GCSFileSystem(project=\"public-assets-275721\")\n",
    "index_path = \"arize-phoenix-assets/datasets/unstructured/llm/llama-index/arize-docs/index/\"\n",
    "storage_context = StorageContext.from_defaults(\n",
    "    fs=file_system,\n",
    "    persist_dir=index_path,\n",
    ")\n",
    "\n",
    "Settings.llm = OpenAI(model=\"gpt-4o\")\n",
    "Settings.embed_model = OpenAIEmbedding(model=\"text-embedding-ada-002\")\n",
    "index = load_index_from_storage(\n",
    "    storage_context,\n",
    ")\n",
    "query_engine = index.as_query_engine()\n",
    "\n",
    "queries_url = \"http://storage.googleapis.com/arize-phoenix-assets/datasets/unstructured/llm/context-retrieval/arize_docs_queries.jsonl\"\n",
    "queries = []\n",
    "with urlopen(queries_url) as response:\n",
    "    for line in response:\n",
    "        line = line.decode(\"utf-8\").strip()\n",
    "        data = json.loads(line)\n",
    "        queries.append(data[\"query\"])\n",
    "queries[:5]\n",
    "\n",
    "for query in tqdm(queries[:5]):\n",
    "    query_engine.query(query)\n",
    "\n",
    "response = query_engine.query(\"What is Arize and how can it help me as an AI Engineer?\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M9dDee9_9ywx"
   },
   "source": [
    "## Step 3: Access your Phoenix instance\n",
    "You can access your Phoenix instance to power evaluations, experiments, upload datasets, etc by using the px.Client() object.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "px_client = px.Client()\n",
    "phoenix_df = px_client.get_spans_dataframe()\n",
    "print(phoenix_df.head())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}