sglang/examples/usage/trace_and_evaluate_rag_using_parea.ipynb

RAG Powered by SGLang & Chroma Evaluated using Parea

In this notebook, we will build a simple RAG pipeline using SGLang to execute our LLM calls, Chroma as vector database for retrieval and Parea for tracing and evaluation. We will then evaluate the performance of our RAG pipeline. The dataset we will use was created by Virat and contains 100 questions, contexts and answers from the Airbnb 2023 10k filing.

The RAG pipeline consists of two steps:

1. Retrieval: Given a question, we retrieve the relevant context from all provided contexts.
2. Generation: Given the question and the retrieved context, we generate an answer.

ℹ️ This notebook requires an OpenAI API key.

ℹ️ This notebook requires a Parea API key, which can be created here.

Setting up the environment

We will first install the necessary packages: sglang, parea and chromadb.

In [ ]:

# note, if you use a Mac M1 chip, you might need to install grpcio 1.59.0 first such that installing chromadb works
# !pip install grpcio==1.59.0

!pip install sglang[openai] parea chromadb

Create a Parea API key as outlined here and save it in a .env file as PAREA_API_KEY=your-api-key.

Indexing the data

Now it's time to download the data & index it! For that, we create a collection called contexts in Chroma and add the contexts as documents.

In [ ]:

import json
import os

import chromadb

path_qca = "airbnb-2023-10k-qca.json"

if not os.path.exists(path_qca):
    !wget https://virattt.github.io/datasets/abnb-2023-10k.json -O airbnb-2023-10k-qca.json

with open(path_qca, 'r') as f:
    question_context_answers = json.load(f)

chroma_client = chromadb.PersistentClient()
collection = chroma_client.get_or_create_collection(name="contexts")
if collection.count() == 0:
    collection.add(
        documents=[qca["context"] for qca in question_context_answers],
        ids=[str(i) for i in range(len(question_context_answers))]
    )

Defining the RAG pipeline

We will start with importing the necessary packages, setting up tracing of OpenAI calls via Parea and setting OpenAI as the default backend for SGLang.

In [2]:

import os
import time

from dotenv import load_dotenv

from sglang import function, user, assistant, gen, set_default_backend, OpenAI
from sglang.lang.interpreter import ProgramState
from parea import Parea, trace


load_dotenv()

os.environ['TOKENIZERS_PARALLELISM'] = "false"

p = Parea(api_key=os.getenv("PAREA_API_KEY"), project_name="rag_sglang")
p.integrate_with_sglang()

set_default_backend(OpenAI("gpt-3.5-turbo"))

Now we can define our retrieval step shown below. Notice, the trace decorator which will automatically trace inputs, output, latency, etc. of that call.

In [ ]:

@trace
def retrieval(question: str) -> list[str]:
    return collection.query(
        query_texts=[question],
        n_results=1
    )['documents'][0]

Next we will define the generation step which uses SGLang to execute the LLM call.

In [ ]:

@function
def generation_sglang(s, question: str, *context: str):
    context = "\n".join(context)
    s += user(f'Given this question:\n{question}\n\nAnd this context:\n{context}\n\nAnswer the question.')
    s += assistant(gen("answer"))


@trace
def generation(question: str, *context):
    state: ProgramState = generation_sglang.run(question, *context)
    while not state.stream_executor.is_finished:
        time.sleep(1)
    return state.stream_executor.variables["answer"]

Finally, we can tie it together and execute a sample query.

In [3]:

@trace
def rag_pipeline(question: str) -> str:
    contexts = retrieval(question)
    return generation(question, *contexts)


rag_pipeline("When did the World Health Organization formally declare an end to the COVID-19 global health emergency?")

Out [3]:

'The World Health Organization formally declared an end to the COVID-19 global health emergency'

Debug Trace

The output is unfortunately wrong! Using the traced pipeline, we can see that

• the context is relevant to the question and contains the correct information
• but, the generation step is cut off as max tokens is set to 16

When opening the generation step in the playground and rerunning the prompt with max. tokens set to 1000, the correct answer is produced.

Evaluating RAG Pipelines

Before we apply above's fix, let's dive into evaluating RAG pipelines.

RAG pipelines consist of a retrieval step to fetch relevant information and a generation step to generate a response to a users question. A RAG pipeline can fail at either step. E.g. the retrieval step can fail to find relevant information which makes generating the correct impossible. Another failure mode is that the generation step doesn't leverage the retrieved information correctly. We will apply the following evaluation metrics to understand different failure modes:

• context_relevancy: measures how relevant the context is given the question
• percent_target_supported_by_context: measures how much of the target answer is supported by the context; this will give an upper ceiling of how well the generation step can perform
• answer_context_faithfulness: measures how much the generated answer utilizes the context
• answer_matches_target: measures how well the generated answer matches the target answer judged by a LLM and gives a sense of accuracy of our entire pipeline

To use these evaluation metrics, we can import them from parea.evals.rag and parea.evals.general and apply them to a function by specifying in the trace decorator which evaluation metrics to use. The @trace decorator will automatically log the results of the evaluation metrics to the Parea dashboard.

Applying them to the retrieval step:

In [ ]:

from parea.evals.rag import context_query_relevancy_factory, percent_target_supported_by_context_factory


context_relevancy_eval = context_query_relevancy_factory()
percent_target_supported_by_context = percent_target_supported_by_context_factory()


@trace(eval_funcs=[context_relevancy_eval, percent_target_supported_by_context])
def retrieval(question: str) -> list[str]:
    return collection.query(
        query_texts=[question],
        n_results=1
    )['documents'][0]

Now we can apply answer_context_faithfulness and answer_matches_target to the generation step.

In [ ]:

from parea.evals.general import answer_matches_target_llm_grader_factory
from parea.evals.rag import answer_context_faithfulness_statement_level_factory


answer_context_faithfulness = answer_context_faithfulness_statement_level_factory()
answer_matches_target_llm_grader = answer_matches_target_llm_grader_factory()

@function
def generation_sglang(s, question: str, *context: str):
    context = "\n".join(context)
    s += user(f'Given this question:\n{question}\n\nAnd this context:\n{context}\n\nAnswer the question.')
    s += assistant(gen("answer", max_tokens=1_000))


@trace(eval_funcs=[answer_context_faithfulness, answer_matches_target_llm_grader])
def generation(question: str, *context):
    state: ProgramState = generation_sglang.run(question, *context)
    while not state.stream_executor.is_finished:
        time.sleep(1)
    return state.stream_executor.variables["answer"]

Finally, we tie them together & execute the original sample query.

In [4]:

@trace
def rag_pipeline(question: str) -> str:
    contexts = retrieval(question)
    return generation(question, *contexts)


rag_pipeline("When did the World Health Organization formally declare an end to the COVID-19 global health emergency?")

Out [4]:

'The World Health Organization formally declared an end to the COVID-19 global health emergency in May 2023.'

Great, the answer is correct! Can you spot the line where we fixed the output truncation issue?

The evaluation scores appear in the bottom right of the logs (screenshot below). Note, that there is no score for answer_matches_target_llm_grader and percent_target_supported_by_context as these evals are automatically skipped if the target answer is not provided.

Running an experiment

Now we are (almost) ready to evaluate the performance of our RAG pipeline on the entire dataset. First, we will need to apply the nest_asyncio package to avoid issues with the Jupyter notebook event loop.

In [6]:

!pip install nest-asyncio
import nest_asyncio
nest_asyncio.apply()

Requirement already satisfied: nest-asyncio in /Users/joschkabraun/miniconda3/envs/sglang/lib/python3.10/site-packages (1.6.0)

Running the actual experiment is straight-forward. For that we use p.experiment to initialize the experiment with a name, the data (list of key-value pairs fed into our entry function) and the entry function. We then call run on the experiment to execute it. Note, that target is a reserved key in the data dictionary and will be used as the target answer for evaluation.

In [7]:

e = p.experiment(
    'RAG',
    data=[
        {
            "question": qca["question"],
            "target": qca["answer"],
        }
        for qca in question_context_answers
    ],
    func=rag_pipeline
).run()

Run name set to: sneak-weal, since a name was not provided.

100%|██████████| 100/100 [00:27<00:00,  3.63it/s]
Waiting for evaluations to finish: 100%|██████████| 19/19 [00:10<00:00,  1.89it/s]

Experiment RAG Run sneak-weal stats:
{
  "latency": "2.69",
  "input_tokens": "61.26",
  "output_tokens": "75.88",
  "total_tokens": "137.14",
  "cost": "0.00",
  "answer_context_faithfulness_statement_level": "0.26",
  "answer_matches_target_llm_grader": "0.22",
  "context_query_relevancy": "0.27",
  "percent_target_supported_by_context": "0.40"
}


View experiment & traces at: https://app.parea.ai/experiments/RAG/30f0244a-d56c-44ff-bdfb-8f47626304b6

Analyzing the results

When opening above experiment, we will see an overview of the experiment as shown below. The upper half shows a summary of the statistics on the left and charts to investigate the distribution and relationships of scores on the right. The lower half is a table with the individual traces which we can use to debug individual samples.

When looking at the statistics, we can see that the accuracy of our RAG pipeline is 22% as measured by answer_matches_target_llm_grader. Though when checking the quality of our retrieval step (context_query_relevancy), we can see that our retrival step is fetching relevant information in only 27% of all samples. As shown in the GIF, we investigate the relationship between the two and see the two scores have 95% agreement. This confirms that the retrieval step is a major bottleneck for our RAG pipeline. So, now it's your turn to improve the retrieval step!

Note, above link isn't publicly accessible but the experiment can be accessed through here.

In [4]: