tutorials/tracing/langchain_google_palm_tracing_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "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",
    "<h1 align=\"center\">Tracing a LangChain and Google PaLM Application</h1>\n",
    "\n",
    "LLM orchestration frameworks such as LangChain provide abstractions that enable users to build powerful applications in a few lines of code. However, the same abstractions can make it difficult to understand what is going on under the hood and to pinpoint the cause of issues.\n",
    "\n",
    "Phoenix makes your LLM applications *observable* by visualizing the underlying structure of each call to your chain and surfacing problematic \"spans\" of execution based on latency, token count, or other evaluation metrics.\n",
    "\n",
    "In this tutorial, you will:\n",
    "- Build a simple LLM application using LangChain and Google PaLM,\n",
    "- Record trace data in OpenInference format,\n",
    "- Inspect the traces and spans of your application to identify uncaught exceptions and sources of latency and cost.\n",
    "\n",
    "ℹ️️ This notebook requires access to the [Google PaLM API](https://developers.generativeai.google/) and a Google API key."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install Dependencies and Import Libraries"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install Phoenix and LangChain."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install arize-phoenix google-generativeai langchain"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "from getpass import getpass\n",
    "from urllib.request import urlopen\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import phoenix as px\n",
    "from langchain.chains import RetrievalQA\n",
    "from langchain.chat_models import ChatGooglePalm\n",
    "from langchain.embeddings import GooglePalmEmbeddings\n",
    "from langchain.retrievers import KNNRetriever\n",
    "from phoenix.trace.langchain import LangChainInstrumentor\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Launch Phoenix\n",
    "\n",
    "You can run Phoenix in the background to collect trace data emitted by any LlamaIndex application that has been instrumented with the `OpenInferenceTracer`.\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 our LangChain application)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session = px.launch_app()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Configure Your Google API Key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set your Google API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if \"GOOGLE_API_KEY\" not in os.environ:\n",
    "    os.environ[\"GOOGLE_API_KEY\"] = getpass(\"🔑 Enter your Google API Key: \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Instrument LangChain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LangChainInstrumentor().instrument()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Build Your LLM Application\n",
    "\n",
    "Define a `RetrievalQA` chain leveraging \"embedding-gecko-001\" and \"chat-bison-001\" from Google PaLM. The knowledge base of this chain is built over the Arize documentation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = GooglePalmEmbeddings(model_name=\"models/embedding-gecko-001\")\n",
    "database_df = pd.read_parquet(\n",
    "    \"http://storage.googleapis.com/arize-phoenix-assets/datasets/unstructured/llm/context-retrieval/langchain-vertexai/database.parquet\"\n",
    ")\n",
    "knn_retriever = KNNRetriever(\n",
    "    index=np.stack(database_df[\"text_vector\"]),\n",
    "    texts=database_df[\"text\"].tolist(),\n",
    "    embeddings=embeddings,\n",
    ")\n",
    "llm = ChatGooglePalm(model_name=\"models/chat-bison-001\")\n",
    "chain = RetrievalQA.from_chain_type(\n",
    "    llm=llm,\n",
    "    retriever=knn_retriever,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Run the Chain\n",
    "\n",
    "Download a small dataset of user queries to ask your application."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "url = \"http://storage.googleapis.com/arize-phoenix-assets/datasets/unstructured/llm/context-retrieval/arize_docs_queries.jsonl\"\n",
    "queries = []\n",
    "with urlopen(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[:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Run your chain against ten queries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for query in tqdm(queries[:10]):\n",
    "    chain.invoke(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check out Phoenix to view your collected spans and traces!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Open the Phoenix UI if you haven't already: {session.url}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Export Your Trace Data\n",
    "\n",
    "You can export your trace data as a pandas dataframe for further analysis and evaluation.\n",
    "\n",
    "In this case, we will export our retriever spans and view them in a pandas dataframe."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trace_df = px.Client().get_spans_dataframe('span_kind == \"RETRIEVER\"')\n",
    "trace_df"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}