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

rename apimanager to budget manager

Browse source
This commit is contained in:
Krrish Dholakia 2023-09-09 16:10:41 -07:00
parent a39756bfda
commit 8ed85b0523
6 changed files with 16 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

18
litellm/tests/test_api_manager.py
View file

@ -8,16 +8,13 @@ sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import litellm
from litellm import apiManager, completion
litellm.success_callback = ["api_manager"]
from litellm import budget_manager, completion
## Scenario 1: User budget enough to make call
def test_user_budget_enough():
user = "1234"
# create a budget for a user
apiManager.create_budget(total_budget=10, user=user)
budget_manager.create_budget(total_budget=10, user=user)
# check if a given call can be made
data = {
@ -26,7 +23,7 @@ def test_user_budget_enough():
}
model = data["model"]
messages = data["messages"]
if apiManager.projected_cost(**data, user=user) <= apiManager.get_total_budget(user):
if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user):
response = completion(**data)
else:
response = "Sorry - no budget!"
@ -37,7 +34,7 @@ def test_user_budget_enough():
def test_user_budget_not_enough():
user = "12345"
# create a budget for a user
apiManager.create_budget(total_budget=0, user=user)
budget_manager.create_budget(total_budget=0, user=user)
# check if a given call can be made
data = {
@ -46,12 +43,9 @@ def test_user_budget_not_enough():
}
model = data["model"]
messages = data["messages"]
projectedCost = apiManager.projected_cost(**data, user=user)
print(f"projectedCost: {projectedCost}")
totalBudget = apiManager.get_total_budget(user)
print(f"totalBudget: {totalBudget}")
if projectedCost <= totalBudget:
if budget_manager.get_current_cost(user=user) < budget_manager.get_total_budget(user=user):
response = completion(**data)
budget_manager.update_cost(completion_obj=response, user=user)
else:
response = "Sorry - no budget!"