rename apimanager to budget manager

litellm/__init__.py

@@ -36,9 +36,11 @@ caching_with_models = False  # if you want the caching key to be model + prompt
 cache: Optional[Cache] = None # cache object
 model_alias_map: Dict[str, str] = {}
 
-####### APIManager ###################
+####### BudgetManager ###################
-from .apimanager import APIManager
+from .budget_manager import BudgetManager
-apiManager = APIManager()
+budget_manager = BudgetManager()
+
+#############################################
 
 def get_model_cost_map():
     url = "https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"

litellm/budget_manager.py

@@ -1,6 +1,6 @@
 import litellm 
 from litellm.utils import ModelResponse
-class APIManager:
+class BudgetManager:
     def __init__(self):
         self.user_dict = {}
 
@@ -22,4 +22,7 @@ class APIManager:
     def update_cost(self, completion_obj: ModelResponse, user: str):
         cost = litellm.completion_cost(completion_response=completion_obj)
         self.user_dict[user]["current_cost"] = cost + self.user_dict[user].get("current_cost", 0)
-        return self.user_dict[user]["current_cost"]
+        return self.user_dict[user]["current_cost"]
+    
+    def get_current_cost(self, user):
+        return self.user_dict[user].get("current_cost", 0)

litellm/tests/test_api_manager.py

@@ -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 
+from litellm import budget_manager, completion 
-
-litellm.success_callback = ["api_manager"]
-
 
 ## 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)
+    if budget_manager.get_current_cost(user=user) < budget_manager.get_total_budget(user=user):
-    print(f"projectedCost: {projectedCost}")
-    totalBudget = apiManager.get_total_budget(user)
-    print(f"totalBudget: {totalBudget}")
-    if projectedCost <= totalBudget:
         response = completion(**data)
+        budget_manager.update_cost(completion_obj=response, user=user)
     else:
         response = "Sorry - no budget!"