proxy_benchmark.py

import argparse
import os
import time

import requests

URLS = {
    "cloudflare": "https://proxy.braintrustapi.com/v1/",
    "vercel": "https://proxy.braintrustdata.com/api/v1/",
    "aws": "https://bb6cfc32ggbhfbd53ttg4aldyy0irfvt.lambda-url.us-east-1.on.aws/proxy/v1/",
    "aws-pc": "https://3hv3t2hyckl47s52vg6yhqcb3i0hoexz.lambda-url.us-east-1.on.aws/proxy/v1/",
#    "local": "http://localhost:8001/proxy/v1/",
}


def format_time(ms):
    return f"{ms:.2f}ms"


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-n", "--number", type=int, default=10)
    args = parser.parse_args()

    s = requests.Session()

    timing = {}
    for name, url in URLS.items():
        print(f"Running {name}")
        times = []
        s = requests.Session()
        for i in range(args.number):
            start = time.time()
            resp = s.post(
                url + "chat/completions",
                headers={"Authorization": f"Bearer {os.environ['BRAINTRUST_API_KEY']}"},
                json={"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "What is 1+1"}], "seed": 2},
            )
            assert resp.ok, resp.text
            times.append(1000 * (time.time() - start))
        timing[name] = times

    for name in URLS.keys():
        print(
            f"{name}: AVG: {format_time(sum(timing[name]) / len(timing[name]))}, MIN: {format_time(min(timing[name]))}, MAX: {format_time(max(timing[name]))}"
        )
        #print(timing[name])