docs: agentic rag notebook (#6227)

tutorials/tracing/agentic_rag_tracing.ipynb

@@ -0,0 +1,382 @@
+{
+ "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 an Agentic RAG App</h1>\n",
+    "\n",
+    "In this tutorial, you will:\n",
+    "\n",
+    "- Build an agentic RAG app\n",
+    "- Instrument and trace the agentic RAG app with Phoenix\n",
+    "- Inspect the trace data in Phoenix\n",
+    "\n",
+    "## Background\n",
+    "\n",
+    "Agentic RAG (Retrieval Augmented Generation) combines the power of traditional RAG systems with autonomous agents that can make decisions and take actions. While traditional RAG simply retrieves relevant context and generates responses, agentic RAG adds a layer of agency - the ability to:\n",
+    "\n",
+    "- Break down complex queries into sub-tasks\n",
+    "- Choose appropriate tools and actions to complete those tasks\n",
+    "- Reason about and synthesize information from multiple sources\n",
+    "- Make decisions about what information is relevant\n",
+    "\n",
+    "For example, rather than just answering questions about data directly, an agentic RAG system might:\n",
+    "1. Analyze the user's question to determine required information\n",
+    "2. Query multiple data sources or tools as needed\n",
+    "3. Combine and reason about the retrieved information\n",
+    "4. Generate a comprehensive response\n",
+    "\n",
+    "This tutorial demonstrates building an agentic RAG system using LlamaIndex's ReAct agent framework combined with vector and SQL query tools.\n",
+    "\n",
+    "Let's get started!\n",
+    "\n",
+    "ℹ️ This notebook requires an OpenAI API key."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Import Dependencies"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "!pip install -q llama-index openai arize-phoenix-otel openinference-instrumentation-llama-index llama-index-vector-stores-chroma chromadb sqlalchemy"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from getpass import getpass\n",
+    "\n",
+    "import chromadb\n",
+    "import chromadb.utils.embedding_functions as embedding_functions\n",
+    "from llama_index.core import VectorStoreIndex\n",
+    "from llama_index.core.agent import ReActAgent\n",
+    "from llama_index.core.tools import QueryEngineTool, ToolMetadata\n",
+    "from llama_index.llms.openai import OpenAI\n",
+    "from llama_index.vector_stores.chroma import ChromaVectorStore\n",
+    "from openinference.instrumentation.llama_index import LlamaIndexInstrumentor\n",
+    "\n",
+    "from phoenix.otel import register"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Add Secret Keys"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### OpenAI API Key"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Phoenix Key\n",
+    "\n",
+    "Phoenix can be run locally or access via phoenix.arize.com for free."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "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_CLIENT_HEADERS\"] = f\"api_key={phoenix_api_key}\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Instrument Our Agent\n",
+    "\n",
+    "Now you'll connect your notebook to Phoenix. This will allow you to trace the agent's actions and see the trace data in Phoenix.\n",
+    "\n",
+    "Phoenix's automatic instrumentation will trace any calls to LlamaIndex, meaning you don't need to do any extra instrumentation work."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# configure the Phoenix tracer\n",
+    "tracer_provider = register(\n",
+    "    project_name=\"agentic-rag-demo\",\n",
+    "    endpoint=\"https://app.phoenix.arize.com/v1/traces\",  # change this endpoint if you're running Phoenix locally\n",
+    ")\n",
+    "\n",
+    "# Finish automatic instrumentation\n",
+    "LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Build Query Engine Tools using Chroma\n",
+    "\n",
+    "Now you're ready to create the two databases that your agent will use to answer questions. This will be done using Chroma, a vector database that will store the company policies and employees."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Company Policies Database"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "openai_ef = embedding_functions.OpenAIEmbeddingFunction(\n",
+    "    api_key=os.environ[\"OPENAI_API_KEY\"], model_name=\"text-embedding-3-small\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chroma_client = chromadb.Client()\n",
+    "chroma_collection = chroma_client.get_or_create_collection(\n",
+    "    \"agentic-rag-demo-company-policies\", embedding_function=openai_ef\n",
+    ")\n",
+    "\n",
+    "chroma_collection.add(\n",
+    "    ids=[\"1\", \"2\", \"3\"],\n",
+    "    documents=[\n",
+    "        \"The travel policy is: Employees must book travel through the company portal. Economy class flights and standard hotel rooms are covered. Meals during travel are reimbursed up to $75/day. All expenses require receipts.\",\n",
+    "        \"The pto policy is: Full-time employees receive 20 days of paid time off per year, accrued monthly. PTO requests must be submitted at least 2 weeks in advance through the HR portal. Unused PTO can carry over up to 5 days into the next year.\",\n",
+    "        \"The dress code is: Business casual attire is required in the office. This includes collared shirts, slacks or knee-length skirts, and closed-toe shoes. Jeans are permitted on Fridays. No athletic wear or overly casual clothing.\",\n",
+    "    ],\n",
+    ")\n",
+    "\n",
+    "vector_store = ChromaVectorStore(chroma_collection=chroma_collection)\n",
+    "\n",
+    "chroma_index = VectorStoreIndex.from_vector_store(vector_store=vector_store)\n",
+    "chroma_engine_policy = chroma_index.as_query_engine(similarity_top_k=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chroma_collection.peek()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chroma_engine_policy.query(\"What is the travel policy?\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Employees Database"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chroma_client = chromadb.Client()\n",
+    "chroma_collection = chroma_client.get_or_create_collection(\n",
+    "    \"agentic-rag-demo-company-employees\", embedding_function=openai_ef\n",
+    ")\n",
+    "\n",
+    "chroma_collection.add(\n",
+    "    ids=[\"1\", \"2\", \"3\"],\n",
+    "    documents=[\n",
+    "        \"John Smith is a Software Engineer in the Engineering department who started on 2023-01-15\",\n",
+    "        \"Sarah Johnson is a Marketing Manager in the Marketing department who started on 2022-08-01\",\n",
+    "        \"Michael Williams is a Sales Director in the Sales department who started on 2021-03-22\",\n",
+    "    ],\n",
+    ")\n",
+    "\n",
+    "vector_store = ChromaVectorStore(chroma_collection=chroma_collection)\n",
+    "\n",
+    "chroma_index = VectorStoreIndex.from_vector_store(vector_store=vector_store)\n",
+    "chroma_engine_employees = chroma_index.as_query_engine(similarity_top_k=1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "chroma_engine_employees.query(\"What is the name of the Sales Director?\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Add as Tools\n",
+    "\n",
+    "LlamaIndex's ReAct agent framework allows you to add tools to the agent. Here you'll add the two tools that will be used to answer questions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "query_engine_tools = [\n",
+    "    QueryEngineTool(\n",
+    "        query_engine=chroma_engine_employees,\n",
+    "        metadata=ToolMetadata(\n",
+    "            name=\"ChromaEmployees\",\n",
+    "            description=(\n",
+    "                \"Provides information about an employee's department, start date, and name from a relational database.\"\n",
+    "                \"Use a detailed plain text question as input to the tool.\"\n",
+    "            ),\n",
+    "        ),\n",
+    "    ),\n",
+    "    QueryEngineTool(\n",
+    "        query_engine=chroma_engine_policy,\n",
+    "        metadata=ToolMetadata(\n",
+    "            name=\"ChromaPolicy\",\n",
+    "            description=(\n",
+    "                \"Provides information about company policies and preocedures. Use this to get more detailed information about company policies.\"\n",
+    "                \"Use a detailed plain text statement about a specific policy as input to the tool.\"\n",
+    "            ),\n",
+    "        ),\n",
+    "    ),\n",
+    "]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create ReAct Agent\n",
+    "\n",
+    "LlamaIndex provides a ReAct agent framework that allows you to create an agent that can use tools to answer questions. Here you'll create an agent that can use the two tools you created earlier to answer questions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "CONTEXT = \"\"\"\n",
+    "You are a chatbot designed to answer questions about the company's employees and policies.\n",
+    "You have access to a Chroma database with information about the company's employees and their departmental information.\n",
+    "You also have access to a Chroma database with information about the company's policies. Use provided context to help answer\n",
+    "the question. Make sure that you have all the context required to answer the question and if you don't, check if there are\n",
+    "other tools that can help you answer the question. If you still can't answer the question, ask the user for more information and\n",
+    "apologize that you can't answer.\n",
+    "\"\"\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
+    "\n",
+    "agent = ReActAgent.from_tools(query_engine_tools, llm=llm, verbose=True, context=CONTEXT)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Test it out!"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "response = agent.chat(\"What department is Sarah in?\")\n",
+    "print(str(response))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "response = agent.chat(\"What is the pto policy for the ML Solutions team?\")\n",
+    "print(str(response))"
+   ]
+  }
+ ],
+ "metadata": {
+  "language_info": {
+   "name": "python"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}