update budget manager

litellm/__init__.py

@@ -36,10 +36,6 @@ 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] = {}
 
-####### BudgetManager ###################
-from .budget_manager import BudgetManager
-budget_manager = BudgetManager()
-
 #############################################
 
 def get_model_cost_map():
@@ -240,7 +236,7 @@ from .utils import (
     get_max_tokens,
     register_prompt_template,
     validate_environment,
-    check_valid_key
+    check_valid_key,
 )
 from .main import *  # type: ignore
 from .integrations import *
@@ -253,3 +249,4 @@ from .exceptions import (
     ContextWindowExceededError,
 
 )
+from .budget_manager import BudgetManager

litellm/budget_manager.py

@@ -1,8 +1,36 @@
+import os, json
 import litellm 
 from litellm.utils import ModelResponse
+
 class BudgetManager:
-    def __init__(self):
+    def __init__(self, type: str):
-        self.user_dict = {}
+        self.type = type
+        ## load the data or init the initial dictionaries
+        self.load_data() 
+    
+    def print_verbose(self, print_statement):
+        if litellm.print_verbose:
+            print(print_statement)
+    
+    def load_data(self):
+        if self.type == "local":
+            # Check if model dict file exists
+            if os.path.isfile("model_cost.json"):
+                # Load the model dict
+                with open("model_cost.json", 'r') as json_file:
+                    self.model_dict = json.load(json_file)
+            else:
+                self.print_verbose("Model Dictionary not found!")
+                self.model_dict = {}
+                    
+            # Check if user dict file exists
+            if os.path.isfile("user_cost.json"):
+                # Load the user dict
+                with open("user_cost.json", 'r') as json_file:
+                    self.user_dict = json.load(json_file)
+            else:
+                self.print_verbose("User Dictionary not found!")
+                self.user_dict = {} 
 
     def create_budget(self, total_budget: float, user: str):
         self.user_dict[user] = {"total_budget": total_budget}
@@ -21,24 +49,22 @@ class BudgetManager:
 
     def update_cost(self, completion_obj: ModelResponse, user: str):
         cost = litellm.completion_cost(completion_response=completion_obj)
-        model = completion_obj.get("model", "") # if this throws an error try, model = completion_obj['model']
+        model = completion_obj['model'] # if this throws an error try, model = completion_obj['model']
-        self.store_llm_costs(cost, model)
+        self.model_dict[model] = cost + self.model_dict.get(model, 0)
         self.user_dict[user]["current_cost"] = cost + self.user_dict[user].get("current_cost", 0)
-        return self.user_dict[user]["current_cost"]
+        return {"current_user_cost": self.user_dict[user]["current_cost"], "current_model_cost": self.model_dict[model]}
     
     def get_current_cost(self, user):
         return self.user_dict[user].get("current_cost", 0)
 
-    def store_llm_costs(self, cost, model):
+    def save_data(self):
-        try:
+        if self.type == "local":
-            import json
+            import json 
-            with open("cost.json", 'r') as json_file:
+            
-                cost_dict = json.load(json_file)
+            # save the model dict
-                if model in cost_dict:
+            with open("model_cost.json", 'w') as json_file:
-                    cost_dict[model] += cost
+                json.dump(self.model_dict, json_file, indent=4)  # Indent for pretty formatting
-                else:
+            
-                    cost_dict[model] = cost
+            # save the user dict 
-                with open("cost.json", 'w') as json_file:
+            with open("user_cost.json", 'w') as json_file:
-                    json.dump(cost_dict, json_file, indent=4)  # Indent for pretty formatting
+                json.dump(self.user_dict, json_file, indent=4)  # Indent for pretty formatting
-        except Exception as e:
-            return # [Non blocking]

litellm/tests/model_cost.json

@@ -0,0 +1,3 @@
+{
+    "gpt-3.5-turbo-0613": 7.7e-05
+}

litellm/tests/test_budget_manager.py

@@ -8,7 +8,9 @@ sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
 import litellm 
-from litellm import budget_manager, completion 
+from litellm import BudgetManager, completion 
+
+budget_manager = BudgetManager(type="local")
 
 ## Scenario 1: User budget enough to make call
 def test_user_budget_enough():
@@ -25,11 +27,14 @@ def test_user_budget_enough():
     messages = data["messages"]
     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}")
 
+test_user_budget_enough()
+
 ## Scenario 2: User budget not enough to make call
 def test_user_budget_not_enough():
     user = "12345"
@@ -45,10 +50,22 @@ def test_user_budget_not_enough():
     messages = data["messages"]
     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)
+        print(budget_manager.update_cost(completion_obj=response, user=user))
     else:
         response = "Sorry - no budget!"
 
     print(f"response: {response}")
 
-test_user_budget_not_enough()
+test_user_budget_not_enough()
+
+## Scenario 3: Saving budget to disk 
+def test_budget_save_to_disk():
+    budget_manager.save_data()
+
+test_budget_save_to_disk()
+
+## Scenario 4: Loading budget from disk
+def test_budget_load_from_disk():
+    budget_manager_2 = BudgetManager(type="local") 
+
+test_budget_load_from_disk()

litellm/tests/user_cost.json

@@ -0,0 +1,9 @@
+{
+    "1234": {
+        "total_budget": 10,
+        "current_cost": 7.7e-05
+    },
+    "12345": {
+        "total_budget": 0
+    }
+}

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "0.1.583"
+version = "0.1.582"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"