litellm/tests/test_end_users.py

import os
# What is this?
## Unit tests for the /end_users/* endpoints
import pytest
import asyncio
import aiohttp
import time
from litellm._uuid import uuid
from openai import AsyncOpenAI
from typing import Optional

"""
- `/end_user/new` 
- `/end_user/info` 
"""


async def generate_key(
    session,
    i,
    budget=None,
    budget_duration=None,
    models=["azure-models", "gpt-4", "dall-e-3"],
    max_parallel_requests: Optional[int] = None,
    user_id: Optional[str] = None,
    team_id: Optional[str] = None,
    calling_key=os.environ["LITELLM_MASTER_KEY"],
):
    url = "http://0.0.0.0:4000/key/generate"
    headers = {
        "Authorization": f"Bearer {calling_key}",
        "Content-Type": "application/json",
    }
    data = {
        "models": models,
        "aliases": {"mistral-7b": "gpt-3.5-turbo"},
        "duration": None,
        "max_budget": budget,
        "budget_duration": budget_duration,
        "max_parallel_requests": max_parallel_requests,
        "user_id": user_id,
        "team_id": team_id,
    }

    print(f"data: {data}")

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        print(f"Response {i} (Status code: {status}):")
        print(response_text)
        print()

        if status != 200:
            raise Exception(f"Request {i} did not return a 200 status code: {status}")

        return await response.json()


async def new_end_user(
    session,
    i,
    user_id=str(uuid.uuid4()),
    model_region=None,
    default_model=None,
    budget_id=None,
):
    url = "http://0.0.0.0:4000/end_user/new"
    headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"}
    data = {
        "user_id": user_id,
        "allowed_model_region": model_region,
        "default_model": default_model,
    }

    if budget_id is not None:
        data["budget_id"] = budget_id
    print("end user data: {}".format(data))

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        print(f"Response {i} (Status code: {status}):")
        print(response_text)
        print()

        if status != 200:
            raise Exception(f"Request {i} did not return a 200 status code: {status}")

        return await response.json()


async def new_budget(session, i, budget_id=None):
    url = "http://0.0.0.0:4000/budget/new"
    headers = {"Authorization": f"Bearer {os.environ['LITELLM_MASTER_KEY']}", "Content-Type": "application/json"}
    data = {
        "budget_id": budget_id,
        "tpm_limit": 2,
    }

    async with session.post(url, headers=headers, json=data) as response:
        status = response.status
        response_text = await response.text()

        print(f"Response {i} (Status code: {status}):")
        print(response_text)
        print()


@pytest.mark.asyncio
async def test_end_user_new():
    """
    Make 20 parallel calls to /user/new. Assert all worked.
    """
    async with aiohttp.ClientSession() as session:
        tasks = [new_end_user(session, i, str(uuid.uuid4())) for i in range(1, 11)]
        await asyncio.gather(*tasks)


@pytest.mark.asyncio
async def test_enduser_tpm_limits_non_master_key():
    """
    1. budget_id = Create Budget with tpm_limit = 10
    2. create end_user with budget_id
    3. Make /chat/completions calls
    4. Sleep 1 second
    4. Make  /chat/completions call -> expect this to fail because rate limit hit
    """
    async with aiohttp.ClientSession() as session:
        # create a budget with budget_id = "free-tier"
        budget_id = f"free-tier-{uuid.uuid4()}"
        await new_budget(session, 0, budget_id=budget_id)
        await asyncio.sleep(2)

        end_user_id = str(uuid.uuid4())

        await new_end_user(
            session=session, i=0, user_id=end_user_id, budget_id=budget_id
        )

        ## MAKE CALL ##
        key_gen = await generate_key(session=session, i=0, models=[])

        key = key_gen["key"]

    # chat completion 1
    client = AsyncOpenAI(api_key=key, base_url="http://0.0.0.0:4000", max_retries=0)

    # chat completion 2
    passed = 0
    for _ in range(10):
        try:
            result = await client.chat.completions.create(
                model="fake-openai-endpoint",
                messages=[{"role": "user", "content": "Hey!"}],
                user=end_user_id,
            )
            passed += 1
        except Exception:
            pass
    print("Passed requests=", passed)

    assert (
        passed < 5
    ), f"Sent 10 requests and end-user has tpm_limit of 2. Number requests passed: {passed}. Expected less than 5 to pass"


@pytest.mark.asyncio
async def test_enduser_tpm_limits_with_master_key():
    """
    1. budget_id = Create Budget with tpm_limit = 10
    2. create end_user with budget_id
    3. Make /chat/completions calls
    4. Sleep 1 second
    4. Make  /chat/completions call -> expect this to fail because rate limit hit
    """
    async with aiohttp.ClientSession() as session:
        # create a budget with budget_id = "free-tier"
        budget_id = f"free-tier-{uuid.uuid4()}"
        await new_budget(session, 0, budget_id=budget_id)

        end_user_id = str(uuid.uuid4())

        await new_end_user(
            session=session, i=0, user_id=end_user_id, budget_id=budget_id
        )

    # chat completion 1
    client = AsyncOpenAI(
        api_key=os.environ["LITELLM_MASTER_KEY"], base_url="http://0.0.0.0:4000", max_retries=0
    )

    # chat completion 2
    passed = 0
    for _ in range(10):
        try:
            result = await client.chat.completions.create(
                model="fake-openai-endpoint",
                messages=[{"role": "user", "content": "Hey!"}],
                user=end_user_id,
            )
            passed += 1
        except Exception:
            pass
    print("Passed requests=", passed)

    assert (
        passed < 5
    ), f"Sent 10 requests and end-user has tpm_limit of 2. Number requests passed: {passed}. Expected less than 5 to pass"