litellm-mirror/litellm/tests/test_budget_manager.py

#### What this tests ####
#    This tests calling batch_completions by running 100 messages together

import sys, os, json
import traceback
import pytest 

sys.path.insert(
    0, os.path.abspath("../..")
)  # Adds the parent directory to the system path
import litellm 
litellm.set_verbose = True
from litellm import completion, BudgetManager

budget_manager = BudgetManager(project_name="test_project", client_type="hosted")

## Scenario 1: User budget enough to make call
def test_user_budget_enough():
    try: 
        user = "1234"
        # create a budget for a user
        budget_manager.create_budget(total_budget=10, user=user, duration="daily")

        # check if a given call can be made
        data = {
            "model": "gpt-3.5-turbo",
            "messages": [{"role": "user", "content": "Hey, how's it going?"}]
        }
        if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user):
            response = completion(**data)
            print(budget_manager.update_cost(completion_obj=response, user=user))
        else:
            response = "Sorry - no budget!"

        print(f"response: {response}")
    except Exception as e:
        pytest.fail(f"An error occurred - {str(e)}")

## Scenario 2: User budget not enough to make call
def test_user_budget_not_enough():
    try: 
        user = "12345"
        # create a budget for a user
        budget_manager.create_budget(total_budget=0, user=user, duration="daily")

        # check if a given call can be made
        data = {
            "model": "gpt-3.5-turbo",
            "messages": [{"role": "user", "content": "Hey, how's it going?"}]
        }
        model = data["model"]
        messages = data["messages"]
        if budget_manager.get_current_cost(user=user) < budget_manager.get_total_budget(user=user):
            response = completion(**data)
            print(budget_manager.update_cost(completion_obj=response, user=user))
        else:
            response = "Sorry - no budget!"

        print(f"response: {response}")
    except:
        pytest.fail(f"An error occurred")

## Scenario 3: Saving budget to client 
def test_save_user_budget():
    try:
        response = budget_manager.save_data()
        if response["status"] == "error":
            raise Exception(f"An error occurred - {json.dumps(response)}")
        print(response)
    except:
        pytest.fail(f"An error occurred")

## Scenario 4: Getting list of users 
def test_get_users():
    try:
        response = budget_manager.get_users()
        print(response)
    except:
        pytest.fail(f"An error occurred") 


## Scenario 5: Reset budget at the end of duration 
def test_reset_on_duration():
    try:
        # First, set a short duration budget for a user
        user = "123456"
        budget_manager.create_budget(total_budget=10, user=user, duration="daily")

        # Use some of the budget
        data = {
            "model": "gpt-3.5-turbo",
            "messages": [{"role": "user", "content": "Hello!"}]
        }
        if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user=user):
            response = litellm.completion(**data)
            print(budget_manager.update_cost(completion_obj=response, user=user))

        assert budget_manager.get_current_cost(user) > 0, f"Test setup failed: Budget did not decrease after completion"

        # Now, we need to simulate the passing of time. Since we don't want our tests to actually take days, we're going
        # to cheat a little -- we'll manually adjust the "created_at" time so it seems like a day has passed.
        # In a real-world testing scenario, we might instead use something like the `freezegun` library to mock the system time. 
        one_day_in_seconds = 24 * 60 * 60
        budget_manager.user_dict[user]["last_updated_at"] -= one_day_in_seconds

        # Now the duration should have expired, so our budget should reset
        budget_manager.update_budget_all_users()

        # Make sure the budget was actually reset
        assert budget_manager.get_current_cost(user) == 0, "Budget didn't reset after duration expired" 
    except Exception as e:
        pytest.fail(f"An error occurred - {str(e)}")

## Scenario 6: passing in text: 
def test_input_text_on_completion():
    try:
        user = "12345"
        budget_manager.create_budget(total_budget=10, user=user, duration="daily")

        input_text = "hello world"
        output_text = "it's a sunny day in san francisco"
        model = "gpt-3.5-turbo"

        budget_manager.update_cost(user=user, model=model, input_text=input_text, output_text=output_text)
        print(budget_manager.get_current_cost(user))
    except Exception as e:
        pytest.fail(f"An error occurred - {str(e)}")

test_input_text_on_completion()