phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

remove budget manager

Browse source
This commit is contained in:
Krrish Dholakia 2023-09-30 11:42:45 -07:00
parent 8d1f5ba69d
commit 2af46e8be9
3 changed files with 110 additions and 107 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
litellm/__pycache__/utils.cpython-311.pyc
View file

Binary file not shown.

2
litellm/budget_manager.py
View file

@ -149,6 +149,8 @@ class BudgetManager:
'project_name' : self.project_name, 'project_name' : self.project_name,
"user_dict": self.user_dict "user_dict": self.user_dict
} }
print(f"data: {data}")
response = requests.post(url, headers=headers, json=data) response = requests.post(url, headers=headers, json=data)
print(f"response: {response.text}")
response = response.json() response = response.json()
return response return response

215
litellm/tests/test_budget_manager.py
View file

@ -1,129 +1,130 @@
#### What this tests #### # #### What this tests ####
# This tests calling batch_completions by running 100 messages together # # This tests calling batch_completions by running 100 messages together
import sys, os, json # import sys, os, json
import traceback # import traceback
import pytest # import pytest
sys.path.insert( # sys.path.insert(
0, os.path.abspath("../..") # 0, os.path.abspath("../..")
) # Adds the parent directory to the system path # ) # Adds the parent directory to the system path
import litellm # import litellm
litellm.set_verbose = True # litellm.set_verbose = True
from litellm import completion, BudgetManager # from litellm import completion, BudgetManager
budget_manager = BudgetManager(project_name="test_project", client_type="hosted") # budget_manager = BudgetManager(project_name="test_project", client_type="hosted")
## Scenario 1: User budget enough to make call # ## Scenario 1: User budget enough to make call
def test_user_budget_enough(): # def test_user_budget_enough():
try: # try:
user = "1234" # user = "1234"
# create a budget for a user # # create a budget for a user
budget_manager.create_budget(total_budget=10, user=user, duration="daily") # budget_manager.create_budget(total_budget=10, user=user, duration="daily")
# check if a given call can be made # # check if a given call can be made
data = { # data = {
"model": "gpt-3.5-turbo", # "model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "Hey, how's it going?"}] # "messages": [{"role": "user", "content": "Hey, how's it going?"}]
} # }
if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user): # if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user):
response = completion(**data) # response = completion(**data)
print(budget_manager.update_cost(completion_obj=response, user=user)) # print(budget_manager.update_cost(completion_obj=response, user=user))
else: # else:
response = "Sorry - no budget!" # response = "Sorry - no budget!"
print(f"response: {response}") # print(f"response: {response}")
except Exception as e: # except Exception as e:
pytest.fail(f"An error occurred - {str(e)}") # pytest.fail(f"An error occurred - {str(e)}")
## Scenario 2: User budget not enough to make call # ## Scenario 2: User budget not enough to make call
def test_user_budget_not_enough(): # def test_user_budget_not_enough():
try: # try:
user = "12345" # user = "12345"
# create a budget for a user # # create a budget for a user
budget_manager.create_budget(total_budget=0, user=user, duration="daily") # budget_manager.create_budget(total_budget=0, user=user, duration="daily")
# check if a given call can be made # # check if a given call can be made
data = { # data = {
"model": "gpt-3.5-turbo", # "model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "Hey, how's it going?"}] # "messages": [{"role": "user", "content": "Hey, how's it going?"}]
} # }
model = data["model"] # model = data["model"]
messages = data["messages"] # messages = data["messages"]
if budget_manager.get_current_cost(user=user) < budget_manager.get_total_budget(user=user): # if budget_manager.get_current_cost(user=user) < budget_manager.get_total_budget(user=user):
response = completion(**data) # response = completion(**data)
print(budget_manager.update_cost(completion_obj=response, user=user)) # print(budget_manager.update_cost(completion_obj=response, user=user))
else: # else:
response = "Sorry - no budget!" # response = "Sorry - no budget!"
print(f"response: {response}") # print(f"response: {response}")
except: # except:
pytest.fail(f"An error occurred") # pytest.fail(f"An error occurred")
## Scenario 3: Saving budget to client # ## Scenario 3: Saving budget to client
def test_save_user_budget(): # def test_save_user_budget():
try: # try:
response = budget_manager.save_data() # response = budget_manager.save_data()
if response["status"] == "error": # if response["status"] == "error":
raise Exception(f"An error occurred - {json.dumps(response)}") # raise Exception(f"An error occurred - {json.dumps(response)}")
print(response) # print(response)
except Exception as e: # except Exception as e:
pytest.fail(f"An error occurred: {str(e)}") # pytest.fail(f"An error occurred: {str(e)}")
## Scenario 4: Getting list of users # test_save_user_budget()
def test_get_users(): # ## Scenario 4: Getting list of users
try: # def test_get_users():
response = budget_manager.get_users() # try:
print(response) # response = budget_manager.get_users()
except: # print(response)
pytest.fail(f"An error occurred") # except:
# pytest.fail(f"An error occurred")
## Scenario 5: Reset budget at the end of duration # ## Scenario 5: Reset budget at the end of duration
def test_reset_on_duration(): # def test_reset_on_duration():
try: # try:
# First, set a short duration budget for a user # # First, set a short duration budget for a user
user = "123456" # user = "123456"
budget_manager.create_budget(total_budget=10, user=user, duration="daily") # budget_manager.create_budget(total_budget=10, user=user, duration="daily")
# Use some of the budget # # Use some of the budget
data = { # data = {
"model": "gpt-3.5-turbo", # "model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "Hello!"}] # "messages": [{"role": "user", "content": "Hello!"}]
} # }
if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user=user): # if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user=user):
response = litellm.completion(**data) # response = litellm.completion(**data)
print(budget_manager.update_cost(completion_obj=response, user=user)) # 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" # 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 # # 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. # # 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. # # 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 # one_day_in_seconds = 24 * 60 * 60
budget_manager.user_dict[user]["last_updated_at"] -= one_day_in_seconds # budget_manager.user_dict[user]["last_updated_at"] -= one_day_in_seconds
# Now the duration should have expired, so our budget should reset # # Now the duration should have expired, so our budget should reset
budget_manager.update_budget_all_users() # budget_manager.update_budget_all_users()
# Make sure the budget was actually reset # # Make sure the budget was actually reset
assert budget_manager.get_current_cost(user) == 0, "Budget didn't reset after duration expired" # assert budget_manager.get_current_cost(user) == 0, "Budget didn't reset after duration expired"
except Exception as e: # except Exception as e:
pytest.fail(f"An error occurred - {str(e)}") # pytest.fail(f"An error occurred - {str(e)}")
## Scenario 6: passing in text: # ## Scenario 6: passing in text:
def test_input_text_on_completion(): # def test_input_text_on_completion():
try: # try:
user = "12345" # user = "12345"
budget_manager.create_budget(total_budget=10, user=user, duration="daily") # budget_manager.create_budget(total_budget=10, user=user, duration="daily")
input_text = "hello world" # input_text = "hello world"
output_text = "it's a sunny day in san francisco" # output_text = "it's a sunny day in san francisco"
model = "gpt-3.5-turbo" # model = "gpt-3.5-turbo"
budget_manager.update_cost(user=user, model=model, input_text=input_text, output_text=output_text) # budget_manager.update_cost(user=user, model=model, input_text=input_text, output_text=output_text)
print(budget_manager.get_current_cost(user)) # print(budget_manager.get_current_cost(user))
except Exception as e: # except Exception as e:
pytest.fail(f"An error occurred - {str(e)}") # pytest.fail(f"An error occurred - {str(e)}")
test_input_text_on_completion() # test_input_text_on_completion()