File size: 5,098 Bytes
395201c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 |
# #### 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() |