diff --git a/python/llm/rag/rag-cookbook.ipynb b/python/llm/rag/rag-cookbook.ipynb
index 0c9813e..a99fe2f 100644
--- a/python/llm/rag/rag-cookbook.ipynb
+++ b/python/llm/rag/rag-cookbook.ipynb
@@ -5,59 +5,31 @@
"metadata": {
"id": "SUknhuHKyc-E"
},
- "source": [
- "\n",
- " \n",
- " ![\"arize](\"https://storage.googleapis.com/arize-assets/arize-logo-white.jpg\")
\n",
- "
\n",
- " Docs\n",
- " |\n",
- " GitHub\n",
- " |\n",
- " Community\n",
- "
\n",
- "\n",
- "\n",
- "Using Arize with RAG
\n",
- "\n",
- "This guide shows you how to create a retrieval augmented generation chatbot and evaluate performance with Arize. RAG is typically to respond to queries using a specified set of documents instead of using the LLM's own training data, reducing hallucination and incorrect generations.\n",
- "\n",
- "We'll go through the following steps:\n",
- "\n",
- "* Create a RAG chatbot using LlamaIndex\n",
- "\n",
- "* Trace the retrieval and llm calls using Arize\n",
- "\n",
- "* Create a dataset to benchmark performance\n",
- "\n",
- "* Evaluate performance using LLM as a judge"
- ]
+ "source": "\n \n \n
\n Docs\n |\n GitHub\n |\n Community\n
\n\n\nEvaluating RAG Retrieval Quality and Correctness
\n\nThis is the companion notebook for the [Evaluating RAG Retrieval Quality and Correctness](https://arize.com/docs/ax/cookbooks/evaluate/evaluating-rag) guide. It shows you how to build a retrieval-augmented generation (RAG) chatbot, trace it in Arize AX, then evaluate, diagnose, and improve its retrieval quality.\n\nYou'll work through the following steps:\n\n* **Step 1: Trace your RAG app.** Instrument a LlamaIndex RAG app and send traces to Arize AX.\n* **Step 2: Evaluate retrieval quality and correctness.** Score retrieval relevance with LLM-as-a-judge and log the results back onto your spans.\n* **Step 3: Diagnose retrieval failures with embeddings.** Trace bad responses to their root cause. This step is a UI walkthrough, so it lives in the guide rather than in this notebook.\n* **Step 4: Improve retrieval with experiments.** Compare retrieval settings such as chunk size, chunk overlap, and the number of documents retrieved (`k`)."
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": "## Background: how RAG retrieval works and where it fails\n\nRAG connects your own data to an LLM. A user asks a question, an embedding is generated from the query, the most relevant context in your knowledge base is retrieved, and that context is added to the prompt sent to the LLM.\n\nWhen a RAG application returns a bad answer, it usually traces back to one of three retrieval failure modes:\n\n1. **Bad response.** The final answer is wrong or unhelpful, often surfaced by negative user feedback or low eval scores. This is usually a symptom of one of the failures below.\n1. **Missing context.** The retriever couldn't find any documents close enough to the query, which means users are asking about topics missing from your corpus.\n1. **Most similar is not most relevant.** A document had the closest embedding to the query but wasn't actually the most relevant one to answer it.\n\nThis notebook traces and evaluates a RAG app so you can catch these failures, and Step 4 shows how to run experiments to improve retrieval."
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": "## Before you start\n\nTo run this notebook, you'll need:\n\n* An Arize AX account. [Sign up here](https://app.arize.com/auth/join) and get your [space ID and API key](https://arize.com/docs/ax/llm-tracing/quickstart-llm#get-your-api-keys).\n* An OpenAI API key.\n\nStep 1 installs the required packages and prompts you for these keys."
},
{
"cell_type": "markdown",
"metadata": {
"id": "FfImo32BJYkr"
},
- "source": [
- "# Create a RAG chatbot using LlamaIndex\n",
- "\n",
- "Let's start with all of our boilerplate setup:\n",
- "\n",
- "1. Install packages for tracing and retrieval\n",
- "2. Setup our API keys\n",
- "3. Setup Phoenix for tracing\n",
- "4. Create our LlamaIndex query engine\n",
- "5. See your results in Phoenix"
- ]
+ "source": "# Step 1: Trace your RAG app\n\nArize AX auto-instruments many RAG stacks, including LlamaIndex and LangChain, and supports manual instrumentation for custom ones. This notebook uses LlamaIndex.\n\nIn this step you'll:\n\n1. Install the packages for tracing and retrieval\n1. Set your API keys\n1. Configure Arize AX for tracing\n1. Build a LlamaIndex query engine and send your first query\n1. View the resulting traces in Arize AX"
},
{
"cell_type": "markdown",
"metadata": {
"id": "DcHymV1dh_SS"
},
- "source": [
- "### Install packages for tracing and retrieval"
- ]
+ "source": "### Install packages for tracing and retrieval\n\nRun the cell below to install LlamaIndex and the Arize tracing and evaluation packages."
},
{
"cell_type": "code",
@@ -69,7 +41,7 @@
"!pip install -qq llama-index openai llama-index-core\n",
"\n",
"# Install arize packages for tracing and evaluation\n",
- "!pip install -qq arize-phoenix-evals arize-otel openinference-instrumentation-llama-index \"arize[Datasets]\""
+ "!pip install -qq arize-phoenix-evals arize-otel openinference-instrumentation-llama-index \"arize>=8.0.0\""
]
},
{
@@ -77,9 +49,7 @@
"metadata": {
"id": "jQnyEnJisyn3"
},
- "source": [
- "### Setup our API Keys"
- ]
+ "source": "### Set your API keys\n\nRun the cell below and enter your Arize Space ID, Arize API key, and OpenAI API key when prompted. Every later step reads these credentials, so run it once at the start."
},
{
"cell_type": "code",
@@ -105,11 +75,7 @@
"metadata": {
"id": "kfid5cE99yN5"
},
- "source": [
- "### Setup Arize for Tracing\n",
- "\n",
- "To follow with this tutorial, you'll need to sign up for Arize and get your API key. You can see the [guide here](https://docs.arize.com/arize/llm-tracing/quickstart-llm)."
- ]
+ "source": "### Configure Arize AX for tracing\n\nSet up tracing with [arize-otel](https://pypi.org/project/arize-otel/), a convenience package that configures OpenTelemetry alongside OpenInference auto-instrumentation, which maps LLM metadata to a standardized set of trace and span attributes. This sends every LlamaIndex call to Arize AX. For help finding your keys, see the [LLM tracing quickstart](https://arize.com/docs/ax/llm-tracing/quickstart-llm#get-your-api-keys)."
},
{
"cell_type": "code",
@@ -135,9 +101,7 @@
"metadata": {
"id": "9Ewpx7Dgebym"
},
- "source": [
- "### Create our LlamaIndex query engine"
- ]
+ "source": "### Build the LlamaIndex query engine\n\nFirst, download the knowledge base the chatbot will answer from. This uses an essay by Paul Graham, the sample document from [LlamaIndex's starter tutorial](https://docs.llamaindex.ai/en/stable/getting_started/starter_example/). Run the cell below to download it into a local `data/` folder."
},
{
"cell_type": "code",
@@ -146,7 +110,7 @@
"outputs": [],
"source": [
"!mkdir data\n",
- "!wget \"https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt\" -O data/paul_graham_essay.txt"
+ "!wget \"https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt\" -O data/paul_graham_essay.txt"
]
},
{
@@ -154,66 +118,59 @@
"execution_count": null,
"metadata": {},
"outputs": [],
- "source": [
- "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
- "from pprint import pprint\n",
- "from llama_index.llms.openai import OpenAI\n",
- "\n",
- "# load documents\n",
- "documents = SimpleDirectoryReader(\"data\").load_data()\n",
- "index = VectorStoreIndex.from_documents(documents)\n",
- "query_engine = index.as_query_engine(llm=OpenAI(model=\"gpt-4o-mini\"))\n",
- "response = query_engine.query(\"What did Paul Graham work on?\")\n",
- "pprint(response)"
- ]
+ "source": "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\nfrom pprint import pprint\nfrom llama_index.llms.openai import OpenAI\n\n# Load the sample document, build a vector index over it, and expose the index\n# as a query engine. This query engine is the RAG app you'll trace and evaluate.\ndocuments = SimpleDirectoryReader(\"data\").load_data()\nindex = VectorStoreIndex.from_documents(documents)\nquery_engine = index.as_query_engine(llm=OpenAI(model=\"gpt-4.1-mini\"))\n\n# Send a first query. This creates a trace in Arize AX.\nresponse = query_engine.query(\"What did Paul Graham work on?\")\npprint(response)"
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
- "source": [
- "for node in response.source_nodes:\n",
- " text_fmt = node.node.get_content().strip().replace(\"\\n\", \" \")[:200] + \"...\"\n",
- " print(text_fmt)\n",
- " print(node.score)\n",
- " print(\"--------\")"
- ]
+ "source": "# Inspect the chunks the retriever returned for this query, with their similarity scores.\nfor node in response.source_nodes:\n text_fmt = node.node.get_content().strip().replace(\"\\n\", \" \")[:200] + \"...\"\n print(text_fmt)\n print(node.score)\n print(\"--------\")"
},
{
"cell_type": "markdown",
"metadata": {
"id": "yUyvcly1iNrv"
},
- "source": [
- "### See your results in the Arize UI\n",
- "Once you've run a single query, you can see the trace in the Arize UI with each step taken by the retriever, the embedding, and the llm query.\n",
- "\n",
- "Click through the queries to better understand how the query engine is performing. Arize can be used to understand and troubleshoot your RAG pipeline by surfacing:\n",
- " - Application latency\n",
- " - Token usage\n",
- " - Runtime exceptions\n",
- " - Retrieved documents\n",
- " - Embeddings\n",
- " - LLM parameters\n",
- " - Prompt templates\n",
- " - Tool descriptions\n",
- " - LLM function calls\n",
- " - And more!\n",
- "\n",
- "![](\"https://storage.googleapis.com/arize-assets/tutorials/images/llamaindex-arize-starter.png\")
"
- ]
+ "source": "### See your results in Arize AX\nOnce you've run a single query, you can see the trace in Arize AX with each step taken by the retriever, the embedding, and the LLM query.\n\nClick through the queries to better understand how the query engine is performing. Arize AX helps you understand and troubleshoot your RAG pipeline by surfacing:\n - Application latency\n - Token usage\n - Runtime exceptions\n - Retrieved documents\n - Embeddings\n - LLM parameters\n - Prompt templates\n - Tool descriptions\n - LLM function calls\n - And more!\n\n"
+ },
+ {
+ "cell_type": "markdown",
+ "source": "# Step 2: Evaluate retrieval quality and correctness\n\nNow that traces are flowing into Arize AX, evaluate them so you don't have to inspect every trace by hand. In this step you'll score **retrieval relevance**: is the retrieved context relevant to the question? Retrieval relevance is judged on the **retriever spans**, which carry both the query and the retrieved documents. You'll then log the scores back onto those spans.\n\nFirst, run the cell below to send a handful of queries so you have traces to evaluate.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "source": "questions = [\n \"What did Paul Graham work on before college?\",\n \"What did Paul Graham study in college?\",\n \"What did Paul Graham do at Y Combinator?\",\n \"Why did Paul Graham start painting?\",\n \"What programming language did Paul Graham help create?\",\n]\nfor q in questions:\n query_engine.query(q)\nprint(\"Sent queries. Traces are on their way to Arize\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "source": "import time\nfrom datetime import datetime, timedelta\nfrom arize import ArizeClient\n\nclient = ArizeClient(api_key=API_KEY)\n\n# Spans take a few moments to become queryable after ingestion.\n# If retriever_df comes back empty, wait and re-run this cell.\ntime.sleep(30)\n\nend_time = datetime.now()\nstart_time = end_time - timedelta(hours=1)\n\nspans_df = client.spans.export_to_df(\n space_id=SPACE_ID,\n project_name=\"rag-cookbook\",\n start_time=start_time,\n end_time=end_time,\n)\n\n# Retriever spans carry both the query (input) and the retrieved documents\nretriever_df = spans_df[\n spans_df[\"attributes.openinference.span.kind\"] == \"RETRIEVER\"\n].copy()\nprint(f\"Exported {len(spans_df)} spans, {len(retriever_df)} retriever spans\")\nretriever_df.head()",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "source": "Now score the retrieved context with an LLM-as-a-judge **retrieval relevance** evaluator and log the results back onto the spans. We use the current [Phoenix Evals](https://arize.com/docs/phoenix/evaluation) API (`LLM` + `create_classifier` + `evaluate_dataframe`), then `to_annotation_dataframe` and `update_evaluations` to attach the scores to each span.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "source": "from phoenix.evals import LLM, create_classifier, evaluate_dataframe\nfrom phoenix.evals.utils import to_annotation_dataframe\n\n# The template placeholders {input} and {reference} are filled from same-named\n# columns, so map the span attributes onto those column names first.\nretriever_df[\"input\"] = retriever_df[\"attributes.input.value\"]\nretriever_df[\"reference\"] = retriever_df[\"attributes.retrieval.documents\"]\n\nllm = LLM(provider=\"openai\", model=\"gpt-4.1\")\n\nRELEVANCE_TEMPLATE = \"\"\"You are comparing a reference text to a question and trying to determine\nif the reference text contains information relevant to answering the question. Here is the data:\n [BEGIN DATA]\n ************\n [Question]: {input}\n ************\n [Reference text]: {reference}\n [END DATA]\nIs the Reference text relevant to answering the Question? Answer \"relevant\" or \"unrelated\".\"\"\"\n\nrelevance_evaluator = create_classifier(\n name=\"retrieval_relevance\",\n llm=llm,\n prompt_template=RELEVANCE_TEMPLATE,\n choices={\"relevant\": 1.0, \"unrelated\": 0.0},\n)\n\nresults_df = evaluate_dataframe(retriever_df, [relevance_evaluator])\n\n# to_annotation_dataframe explodes the scores into label/score/explanation and\n# carries context.span_id through. update_evaluations expects the eval..* convention.\nannotation_df = to_annotation_dataframe(results_df, [\"retrieval_relevance\"])\neval_df = annotation_df.rename(\n columns={\n \"label\": \"eval.retrieval_relevance.label\",\n \"score\": \"eval.retrieval_relevance.score\",\n \"explanation\": \"eval.retrieval_relevance.explanation\",\n }\n)[\n [\n \"context.span_id\",\n \"eval.retrieval_relevance.label\",\n \"eval.retrieval_relevance.score\",\n \"eval.retrieval_relevance.explanation\",\n ]\n]\n\nclient.spans.update_evaluations(\n space_id=SPACE_ID,\n project_name=\"rag-cookbook\",\n dataframe=eval_df,\n force_http=True,\n)\nprint(\"Logged retrieval_relevance evals back to spans\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
},
{
"cell_type": "markdown",
"metadata": {
"id": "k0Qvn8tAs9vL"
},
- "source": [
- "# Create synthetic dataset of questions\n",
- "\n",
- "Using the template below, we're going to generate a dataframe of 25 questions we can use to test our customer support agent."
- ]
+ "source": "# Step 4: Improve retrieval with experiments\n\nOnce you can measure retrieval quality, experiment to improve it.\n\nStep 3, diagnosing retrieval failures with embeddings, is a UI walkthrough. See the [guide](https://arize.com/docs/ax/cookbooks/evaluate/evaluating-rag) for that step.\n\nFirst, generate a synthetic benchmark set of questions from the essay. You'll reuse this set to compare retrieval settings: chunk size, chunk overlap, and the number of documents retrieved (`k`)."
},
{
"cell_type": "code",
@@ -247,47 +204,28 @@
"execution_count": null,
"metadata": {},
"outputs": [],
- "source": [
- "import nest_asyncio\n",
- "import pandas as pd\n",
- "\n",
- "nest_asyncio.apply()\n",
- "from phoenix.evals import OpenAIModel\n",
- "\n",
- "pd.set_option(\"display.max_colwidth\", 500)\n",
- "\n",
- "model = OpenAIModel(model=\"gpt-4o\", max_tokens=1300)"
- ]
+ "source": "import pandas as pd\nfrom phoenix.evals import LLM\n\npd.set_option(\"display.max_colwidth\", 500)\n\ngen_llm = LLM(provider=\"openai\", model=\"gpt-4.1\")"
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
- "source": [
- "resp = model(GEN_TEMPLATE)"
- ]
+ "source": "# Generate the benchmark questions with gpt-4.1.\nresp = gen_llm.generate_text(GEN_TEMPLATE)"
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
- "source": [
- "split_response = resp.strip().split(\"\\n\\n\")\n",
- "\n",
- "questions_df = pd.DataFrame(split_response, columns=[\"input\"])\n",
- "print(questions_df.head(3))"
- ]
+ "source": "questions = [q.strip() for q in resp.strip().split(\"\\n\") if q.strip()]\n\nquestions_df = pd.DataFrame(questions, columns=[\"input\"])\nquestions_df.head(3)"
},
{
"cell_type": "markdown",
"metadata": {
"id": "oGIbV49kHp4H"
},
- "source": [
- "Now let's run it and manually inspect the traces! "
- ]
+ "source": "Before running experiments, run the baseline query engine over the benchmark set to sanity-check its answers and retrieved context. Each row gets an `output` (the answer) and a `reference` (the retrieved context)."
},
{
"cell_type": "code",
@@ -321,11 +259,7 @@
"metadata": {
"id": "beUkwcCgLaEa"
},
- "source": [
- "# Evaluating your RAG app\n",
- "\n",
- "Now that we have a set of test cases, we can create evaluators to measure performance. This way, we don't have to manually inspect every single trace to see if the LLM is doing the right thing."
- ]
+ "source": "Now that we have a benchmark set and baseline RAG outputs, create evaluators to measure performance so we don't have to inspect every trace by hand. We'll reuse these evaluators to score each experiment run."
},
{
"cell_type": "code",
@@ -373,76 +307,45 @@
"metadata": {
"id": "1aivaxTCRQFl"
},
- "source": [
- "We will be creating an LLM as a judge using the prompt templates above by taking the spans recorded by Phoenix, and then giving them labels using the `llm_classify` function. This function uses LLMs to evaluate your LLM calls and gives them labels and explanations. You can read more detail [here](https://docs.arize.com/phoenix/api/evals#phoenix.evals.llm_classify)."
- ]
+ "source": "We create two LLM-as-a-judge evaluators from the prompt templates above using the current Phoenix Evals API: `create_classifier` combines a prompt, an `LLM`, and a set of `choices` into a classifier. Calling `.evaluate()` on a row returns a `Score` with a `label`, `score`, and `explanation`. The `run_evaluators` helper scores every row of a RAG run and writes the labels and explanations into columns we can log as experiment evaluations."
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
- "source": [
- "from phoenix.evals import OpenAIModel, llm_classify\n",
- "\n",
- "RELEVANCE_RAILS = [\"relevant\", \"unrelated\"]\n",
- "CORRECTNESS_RAILS = [\"incorrect\", \"correct\"]\n",
- "\n",
- "relevance_eval_df = llm_classify(\n",
- " dataframe=response_df,\n",
- " template=RELEVANCE_EVAL_TEMPLATE,\n",
- " model=OpenAIModel(model=\"gpt-4o\"),\n",
- " rails=RELEVANCE_RAILS,\n",
- " provide_explanation=True,\n",
- " include_prompt=True,\n",
- " concurrency=4,\n",
- ")\n",
- "\n",
- "correctness_eval_df = llm_classify(\n",
- " dataframe=response_df,\n",
- " template=CORRECTNESS_EVAL_TEMPLATE,\n",
- " model=OpenAIModel(model=\"gpt-4o\"),\n",
- " rails=CORRECTNESS_RAILS,\n",
- " provide_explanation=True,\n",
- " include_prompt=True,\n",
- " concurrency=4,\n",
- ")"
- ]
+ "source": "from phoenix.evals import LLM, create_classifier\n\neval_llm = LLM(provider=\"openai\", model=\"gpt-4.1\")\n\nrelevance_experiment_evaluator = create_classifier(\n name=\"relevance\",\n llm=eval_llm,\n prompt_template=RELEVANCE_EVAL_TEMPLATE,\n choices={\"relevant\": 1.0, \"unrelated\": 0.0},\n)\ncorrectness_evaluator = create_classifier(\n name=\"correctness\",\n llm=eval_llm,\n prompt_template=CORRECTNESS_EVAL_TEMPLATE,\n choices={\"correct\": 1.0, \"incorrect\": 0.0},\n)\n\n\ndef run_evaluators(rag_df):\n for i, row in rag_df.iterrows():\n relevance = relevance_experiment_evaluator.evaluate(row.to_dict())[0]\n correctness = correctness_evaluator.evaluate(row.to_dict())[0]\n rag_df.loc[i, \"relevance\"] = relevance.label\n rag_df.loc[i, \"relevance_explanation\"] = relevance.explanation\n rag_df.loc[i, \"correctness\"] = correctness.label\n rag_df.loc[i, \"correctness_explanation\"] = correctness.explanation\n return rag_df\n\n\n# Score the baseline run from above\nresponse_df = run_evaluators(response_df)\nresponse_df.head(3)"
},
{
"cell_type": "markdown",
"metadata": {
"id": "vDV1KBdYQ_vh"
},
- "source": [
- "Let's look at and inspect the results of our evaluatiion!"
- ]
+ "source": "Inspect the relevance and correctness labels for the baseline run."
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
- "source": [
- "relevance_eval_df"
- ]
+ "source": "response_df[[\"input\", \"relevance\", \"correctness\"]]"
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
- "source": [
- "correctness_eval_df"
- ]
+ "source": "print(response_df[\"relevance\"].value_counts())\nprint(response_df[\"correctness\"].value_counts())"
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Experiment with different k-values\n",
+ "## Compare retrieval settings\n",
+ "\n",
+ "Now rebuild the query engine with different retrieval settings and compare them as experiments in Arize AX. `client.experiments.run` executes a **task** across the benchmark set, applies our **evaluators**, and logs each run so we can compare relevance and correctness side by side in the UI.\n",
"\n",
- "We can also experiment with different k-values for the retriever. This is the number of documents retrieved from the vector store. We can also experiment with different chunk sizes, chunk overlaps, and rerankers. We'll be using the ColbertReranker from LlamaIndex. You can read more about it [here](https://docs.llamaindex.ai/docs/postprocessors/colbert-reranker)."
+ "Define a helper that builds a query engine for a given configuration, then wrap it as a task. The task answers one question per dataset row and returns the answer along with the retrieved context, which is carried through only so the evaluators can score it."
]
},
{
@@ -455,26 +358,33 @@
"from llama_index.llms.openai import OpenAI\n",
"\n",
"\n",
- "def run_rag_with_settings(questions_df, k_value, chunk_size, chunk_overlap):\n",
+ "def make_task(k_value, chunk_size, chunk_overlap):\n",
" node_parser = SimpleNodeParser.from_defaults(\n",
" chunk_size=chunk_size, chunk_overlap=chunk_overlap\n",
" )\n",
" nodes = node_parser.get_nodes_from_documents(documents)\n",
" vector_index = VectorStoreIndex(nodes)\n",
- " query_engine = vector_index.as_query_engine(\n",
+ " engine = vector_index.as_query_engine(\n",
" similarity_top_k=k_value, # Default is 2\n",
" response_mode=\"compact\", # or use \"tree-summarize\"\n",
- " llm=OpenAI(model=\"gpt-4o-mini\"),\n",
+ " llm=OpenAI(model=\"gpt-4.1-mini\"),\n",
" )\n",
- " response_df = run_rag(query_engine, questions_df)\n",
- " return response_df"
+ "\n",
+ " def task(dataset_row):\n",
+ " response = engine.query(dataset_row[\"input\"])\n",
+ " return {\n",
+ " \"output\": str(response),\n",
+ " \"reference\": \"\\n\".join(n.text for n in response.source_nodes),\n",
+ " }\n",
+ "\n",
+ " return task"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "Let's setup our evaluators to see how the performance changes."
+ "Reuse the Phoenix Evals classifiers defined above as experiment evaluators. Each evaluator receives the task's `output` and the `dataset_row`, then returns an `EvaluationResult` with a label, score, and explanation."
]
},
{
@@ -483,43 +393,35 @@
"metadata": {},
"outputs": [],
"source": [
- "def run_evaluators(rag_df):\n",
- " relevance_eval_df = llm_classify(\n",
- " dataframe=rag_df,\n",
- " template=RELEVANCE_EVAL_TEMPLATE,\n",
- " model=OpenAIModel(model=\"gpt-4o\"),\n",
- " rails=RELEVANCE_RAILS,\n",
- " provide_explanation=True,\n",
- " concurrency=4,\n",
- " )\n",
- " rag_df[\"relevance\"] = relevance_eval_df[\"label\"]\n",
- " rag_df[\"relevance_explanation\"] = relevance_eval_df[\"explanation\"]\n",
+ "from arize.experiments import EvaluationResult\n",
"\n",
- " correctness_eval_df = llm_classify(\n",
- " dataframe=rag_df,\n",
- " template=CORRECTNESS_EVAL_TEMPLATE,\n",
- " model=OpenAIModel(model=\"gpt-4o\"),\n",
- " rails=CORRECTNESS_RAILS,\n",
- " provide_explanation=True,\n",
- " concurrency=4,\n",
+ "\n",
+ "def relevance(output, dataset_row):\n",
+ " score = relevance_experiment_evaluator.evaluate({\n",
+ " \"input\": dataset_row[\"input\"],\n",
+ " \"reference\": output[\"reference\"],\n",
+ " })[0]\n",
+ " return EvaluationResult(\n",
+ " score=score.score, label=score.label, explanation=score.explanation\n",
" )\n",
- " rag_df[\"correctness\"] = correctness_eval_df[\"label\"]\n",
- " rag_df[\"correctness_explanation\"] = correctness_eval_df[\"explanation\"]\n",
- " return rag_df"
- ]
- },
- {
- "cell_type": "markdown",
- "metadata": {},
- "source": [
- "Let's log these results to Arize and see how they compare."
+ "\n",
+ "\n",
+ "def correctness(output, dataset_row):\n",
+ " score = correctness_evaluator.evaluate({\n",
+ " \"input\": dataset_row[\"input\"],\n",
+ " \"reference\": output[\"reference\"],\n",
+ " \"output\": output[\"output\"],\n",
+ " })[0]\n",
+ " return EvaluationResult(\n",
+ " score=score.score, label=score.label, explanation=score.explanation\n",
+ " )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "First we'll create a dataset to store our questions."
+ "Create a dataset from the benchmark questions, then define a helper that runs one experiment per configuration. `client.experiments.run` executes the task across every example, applies the evaluators, and logs the results to Arize AX."
]
},
{
@@ -528,36 +430,24 @@
"metadata": {},
"outputs": [],
"source": [
- "from arize.experimental.datasets import ArizeDatasetsClient\n",
- "from uuid import uuid1\n",
- "from arize.experimental.datasets.experiments.types import (\n",
- " ExperimentTaskResultColumnNames,\n",
- " EvaluationResultColumnNames,\n",
- ")\n",
- "from arize.experimental.datasets.utils.constants import GENERATIVE\n",
- "import pandas as pd\n",
+ "dataset_name = \"rag-experiments\"\n",
+ "client.datasets.create(space=SPACE_ID, name=dataset_name, examples=questions_df)\n",
"\n",
- "# Set up the arize client\n",
- "arize_client = ArizeDatasetsClient(api_key=API_KEY)\n",
- "dataset = None\n",
- "dataset_name = \"rag-experiments-\" + str(uuid1())[:3]\n",
"\n",
- "dataset_id = arize_client.create_dataset(\n",
- " space_id=SPACE_ID,\n",
- " dataset_name=dataset_name,\n",
- " dataset_type=GENERATIVE,\n",
- " data=questions_df,\n",
- ")\n",
- "dataset = arize_client.get_dataset(space_id=SPACE_ID, dataset_id=dataset_id)\n",
- "print(dataset)"
+ "def run_experiment(name, k_value, chunk_size, chunk_overlap):\n",
+ " client.experiments.run(\n",
+ " name=name,\n",
+ " dataset=dataset_name,\n",
+ " space=SPACE_ID,\n",
+ " task=make_task(k_value, chunk_size, chunk_overlap),\n",
+ " evaluators=[relevance, correctness],\n",
+ " )"
]
},
{
"cell_type": "markdown",
"metadata": {},
- "source": [
- "Next we'll define which columns of our dataframe will be mapped to outputs and which will be mapped to evaluation labels and explanations.."
- ]
+ "source": "Run the experiments that vary `k`, the number of documents retrieved:"
},
{
"cell_type": "code",
@@ -565,42 +455,15 @@
"metadata": {},
"outputs": [],
"source": [
- "# Define column mappings for task\n",
- "task_cols = ExperimentTaskResultColumnNames(\n",
- " example_id=\"example_id\", result=\"output\"\n",
- ")\n",
- "# Define column mappings for evaluator\n",
- "relevance_evaluator_cols = EvaluationResultColumnNames(\n",
- " label=\"relevance\",\n",
- " explanation=\"relevance_explanation\",\n",
- ")\n",
- "correctness_evaluator_cols = EvaluationResultColumnNames(\n",
- " label=\"correctness\",\n",
- " explanation=\"correctness_explanation\",\n",
- ")\n",
- "\n",
- "\n",
- "def log_experiment_to_arize(experiment_df, experiment_name):\n",
- " experiment_df[\"example_id\"] = dataset[\"id\"]\n",
- " return arize_client.log_experiment(\n",
- " space_id=SPACE_ID,\n",
- " experiment_name=experiment_name + \"-\" + str(uuid1())[:2],\n",
- " experiment_df=experiment_df,\n",
- " task_columns=task_cols,\n",
- " evaluator_columns={\n",
- " \"correctness\": correctness_evaluator_cols,\n",
- " \"relevance\": relevance_evaluator_cols,\n",
- " },\n",
- " dataset_name=dataset_name,\n",
- " )"
+ "run_experiment(\"k2-chunk100-overlap10\", k_value=2, chunk_size=100, chunk_overlap=10)\n",
+ "run_experiment(\"k4-chunk100-overlap10\", k_value=4, chunk_size=100, chunk_overlap=10)\n",
+ "run_experiment(\"k10-chunk100-overlap10\", k_value=10, chunk_size=100, chunk_overlap=10)"
]
},
{
"cell_type": "markdown",
"metadata": {},
- "source": [
- "Now let's run it for each of our experiments."
- ]
+ "source": "Run the experiments that vary chunk size and overlap:"
},
{
"cell_type": "code",
@@ -608,49 +471,9 @@
"metadata": {},
"outputs": [],
"source": [
- "# Run Experiments for k-size\n",
- "k_2_chunk_100_overlap_10 = run_rag_with_settings(\n",
- " questions_df, k_value=2, chunk_size=100, chunk_overlap=10\n",
- ")\n",
- "k_4_chunk_100_overlap_10 = run_rag_with_settings(\n",
- " questions_df, k_value=4, chunk_size=100, chunk_overlap=10\n",
- ")\n",
- "k_10_chunk_100_overlap_10 = run_rag_with_settings(\n",
- " questions_df, k_value=10, chunk_size=100, chunk_overlap=10\n",
- ")\n",
- "k_2_chunk_100_overlap_10 = run_evaluators(k_2_chunk_100_overlap_10)\n",
- "k_4_chunk_100_overlap_10 = run_evaluators(k_4_chunk_100_overlap_10)\n",
- "k_10_chunk_100_overlap_10 = run_evaluators(k_10_chunk_100_overlap_10)\n",
- "\n",
- "log_experiment_to_arize(k_2_chunk_100_overlap_10, \"k_2_chunk_100_overlap_10\")\n",
- "log_experiment_to_arize(k_4_chunk_100_overlap_10, \"k_4_chunk_100_overlap_10\")\n",
- "log_experiment_to_arize(k_10_chunk_100_overlap_10, \"k_10_chunk_100_overlap_10\")"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "# Run experiments for chunk size\n",
- "k_2_chunk_200_overlap_10 = run_rag_with_settings(\n",
- " questions_df, k_value=2, chunk_size=200, chunk_overlap=10\n",
- ")\n",
- "k_2_chunk_500_overlap_20 = run_rag_with_settings(\n",
- " questions_df, k_value=2, chunk_size=500, chunk_overlap=20\n",
- ")\n",
- "k_2_chunk_1000_overlap_50 = run_rag_with_settings(\n",
- " questions_df, k_value=2, chunk_size=1000, chunk_overlap=50\n",
- ")\n",
- "\n",
- "k_2_chunk_200_overlap_10 = run_evaluators(k_2_chunk_200_overlap_10)\n",
- "k_2_chunk_500_overlap_20 = run_evaluators(k_2_chunk_500_overlap_20)\n",
- "k_2_chunk_1000_overlap_50 = run_evaluators(k_2_chunk_1000_overlap_50)\n",
- "\n",
- "log_experiment_to_arize(k_2_chunk_200_overlap_10, \"k_2_chunk_200_overlap_10\")\n",
- "log_experiment_to_arize(k_2_chunk_500_overlap_20, \"k_2_chunk_500_overlap_20\")\n",
- "log_experiment_to_arize(k_2_chunk_1000_overlap_50, \"k_2_chunk_1000_overlap_50\")"
+ "run_experiment(\"k2-chunk200-overlap10\", k_value=2, chunk_size=200, chunk_overlap=10)\n",
+ "run_experiment(\"k2-chunk500-overlap20\", k_value=2, chunk_size=500, chunk_overlap=20)\n",
+ "run_experiment(\"k2-chunk1000-overlap50\", k_value=2, chunk_size=1000, chunk_overlap=50)"
]
},
{
@@ -671,24 +494,38 @@
"from llama_index.core.indices.query.query_transform import HyDEQueryTransform\n",
"from llama_index.core.query_engine import TransformQueryEngine\n",
"\n",
- "# Setup HyDe\n",
- "hyde = HyDEQueryTransform(\n",
- " include_original=True, llm=OpenAI(model=\"gpt-4o-mini\")\n",
- ")\n",
- "query_engine = (\n",
- " index.as_query_engine()\n",
- ") # default k=2, chunk_size=1024, chunk_overlap=20\n",
- "hyde_query_engine = TransformQueryEngine(query_engine, hyde)\n",
- "\n",
- "# Run RAG with HyDE\n",
- "hyde_response_df = run_rag(hyde_query_engine, questions_df)\n",
- "\n",
- "# Evaluate RAG with HyDE\n",
- "hyde_response_df = run_evaluators(hyde_response_df)\n",
"\n",
- "# Log to Arize\n",
- "log_experiment_to_arize(hyde_response_df, \"hyde\")"
+ "def make_hyde_task():\n",
+ " hyde = HyDEQueryTransform(\n",
+ " include_original=True, llm=OpenAI(model=\"gpt-4.1-mini\")\n",
+ " )\n",
+ " # default k=2, chunk_size=1024, chunk_overlap=20\n",
+ " base_engine = index.as_query_engine()\n",
+ " engine = TransformQueryEngine(base_engine, hyde)\n",
+ "\n",
+ " def task(dataset_row):\n",
+ " response = engine.query(dataset_row[\"input\"])\n",
+ " return {\n",
+ " \"output\": str(response),\n",
+ " \"reference\": \"\\n\".join(n.text for n in response.source_nodes),\n",
+ " }\n",
+ "\n",
+ " return task\n",
+ "\n",
+ "\n",
+ "client.experiments.run(\n",
+ " name=\"hyde\",\n",
+ " dataset=dataset_name,\n",
+ " space=SPACE_ID,\n",
+ " task=make_hyde_task(),\n",
+ " evaluators=[relevance, correctness],\n",
+ ")"
]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": "## Summary and next steps\n\nIn this notebook, you built a RAG application and used Arize AX to:\n\n* Trace retrieval, embeddings, and LLM calls end to end\n* Score retrieval relevance with LLM-as-a-judge and log the results back onto your spans\n* Run experiments to improve retrieval across chunk size, overlap, and `k`\n\nTo go further:\n\n* Diagnose retrieval failures with embeddings in the UI, as described in Step 3 of the [guide](https://arize.com/docs/ax/cookbooks/evaluate/evaluating-rag)\n* Evaluate across more dimensions such as hallucination, citation, and user frustration\n* Go deeper on [datasets and experiments](https://arize.com/docs/ax/develop/datasets-and-experiments)"
}
],
"metadata": {