How to Compare LLMs with an A/B Test

Create a virtual environment, e.g. if you use virtualenv then you can use

pyenv virtualenv  eppo-ai-sdk-demo

and activate it

pyenv activate eppo-ai-sdk-demo

Note, use python versions to list available versions and select the version you want to use (e.g. 3.10.6, assuming you have that version installed).

Next, install the libraries we will be using

pip install -r requirements.txt

We create a simple Hello World API as follows in src/app.py:

from fastapi import FastAPI

app = FastAPI()

@app.get("/")
def root():
    return {"Hello": "World"}

Now, start the app using

uvicorn src.app:app --reload

Next up, let's create an endpoint that can answer questions. First, set up the API endpoint by adding the following endpoint:

@app.get("/qa")
def qa(question: str):
     return {"question": question, "answer": "42"}

Now use the interactive docs or navigate to http://127.0.0.1:8000/qa?question=What%20is%20the%20answer%20to%20the%20ultimate%20question%20of%20life%3F to verify everything works as expected.

Let's store the API key locally and make sure it does not get checked into github by accident. Copy the .env.dist file to .env from the top-level directory like so:

cp src/.env.dist src/.env

Now open src/.env and add your OpenAI API key.

Next, we install the openai and python-dotenv libraries (the latter is used to read the API key you just stored).

Add the following to the top of your app.py file:

import os
import openai
from dotenv import load_dotenv

load_dotenv()

openai.api_key = os.getenv("OPENAI_API_KEY")

print(openai.api_key[:6]) # optional to test you indeed loaded your API key successfully

For the fun part, we create a function that uses OpenAI's chat completion to answer our question by adding the following code

def openai_chat_completion(question, model="gpt-3.5-turbo"):
     completion = openai.ChatCompletion.create(
          model=model,
          messages=[
               {"role": "system", "content": "You are a funny assistant."},
               {"role": "user", "content": question},
          ],
     )
     return completion.choices[0].message["content"]

and update our qa endpoint to use this function:

@app.get("/qa")
def qa(question: str):
     answer = openai_chat_completion(question)
     return {"question": question, "answer": answer}

First, create an Eppo API key in the test environment. The name is unimporant, consider calling it AI testing key, but make sure you copy paste the actual key into the your src/.env file. For the correct format, compare your .env file with the .env.dist file.

With the API key setup, we can now create a feature flag and evaluate it in our app. Create a new feature flag called AI demo model version and we can use the automatically generated key ai-demo-model-version. Let's create three variations:

First up, install the Eppo Python SDK using pip install eppo-server-sdk. We should now be able to import the SDK into our script by adding

import eppo_client
from eppo_client.config import Config
from eppo_client.assignment_logger import AssignmentLogger

Next, we set up our logger and initialize the SDK client, which fetches the allocation configuration we just set up in the Eppo UI

class PrintAssignment(AssignmentLogger):
     def log_assignment(self, assignment):
          print(assignment)
          
client_config = Config(
     api_key=os.getenv("EPPO_API_KEY"), assignment_logger=PrintAssignment()
)
eppo = eppo_client.get_instance()

@app.get("/qa")
def qa(user: str, question: str):
     variant = eppo.get_assignment(user, "ai-demo-model-version")
     answer = openai_chat_completion(question)
     return {"question": question, "answer": answer, "variant": variant}

Make sure to try the same user a couple of times to verify that indeed the exact same variant shows up every time. In my case, Alice and Bob consistently see the GPT4 variant, while Charlie sees GPT3.5, but you will see different results. Furthermore, you should see the assignments log to your terminal as well:

{'experiment': 'ai-demo-model-version', 'variation': 'gpt-4', 'subject': 'Alice', 'timestamp': '2023-07-14T18:21:18.363098', 'subjectAttributes': {}}

We saved the best for last: using our variant to dynamically pick model version. This is now a simple change to our qa endpoint.

When "gpt" is in the variant name, we can supply the model variant directly to the openai_chat_completion function. Otherwise, simply return "42". Note that it is good practice to code defensively here. While the Eppo SDK is extremely robust, by checking whether the variant actually exists we make sure the code runs correctly even when there is an issue fetching the assignment from our SDK.

@app.get("/qa")
def qa(user: str, question: str):
     variant = eppo.get_assignment(user, "ai-demo-model-version")
     if variant and "gpt" in variant:
          answer = openai_chat_completion(question, model=variant)
     else: 
          answer = "42"    
     return {"question": question, "answer": answer, "variant": variant}