mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-24 18:24:20 +00:00
remove budget manager
This commit is contained in:
Krrish Dholakia
parent
8d1f5ba69d
commit
2af46e8be9
3 changed files with 110 additions and 107 deletions
Binary file not shown.
|
|
@ -149,6 +149,8 @@ class BudgetManager:
|
|||
'project_name' : self.project_name,
|
||||
"user_dict": self.user_dict
|
||||
}
|
||||
print(f"data: {data}")
|
||||
response = requests.post(url, headers=headers, json=data)
|
||||
print(f"response: {response.text}")
|
||||
response = response.json()
|
||||
return response
|
||||
|
|
@ -1,129 +1,130 @@
|
|||
#### What this tests ####
|
||||
# This tests calling batch_completions by running 100 messages together
|
||||
# #### What this tests ####
|
||||
# # This tests calling batch_completions by running 100 messages together
|
||||
|
||||
import sys, os, json
|
||||
import traceback
|
||||
import pytest
|
||||
# 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
|
||||
# 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")
|
||||
# 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")
|
||||
# ## 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!"
|
||||
# # 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)}")
|
||||
# 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")
|
||||
# ## 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!"
|
||||
# # 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")
|
||||
# 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)}")
|
||||
# ## 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)}")
|
||||
|
||||
## 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")
|
||||
# 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")
|
||||
# ## 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))
|
||||
# # 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"
|
||||
# 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, 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()
|
||||
# # 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)}")
|
||||
# # 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")
|
||||
# ## 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"
|
||||
# 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)}")
|
||||
# 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()
|
||||
# test_input_text_on_completion()
|
||||
Loading…
Add table
Add a link
Reference in a new issue