# #### 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 Exception as e: # pytest.fail(f"An error occurred: {str(e)}") # test_save_user_budget() # ## 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()