diff --git a/dist/litellm-0.1.585-py3-none-any.whl b/dist/litellm-0.1.585-py3-none-any.whl new file mode 100644 index 000000000..a46a1e81f Binary files /dev/null and b/dist/litellm-0.1.585-py3-none-any.whl differ diff --git a/dist/litellm-0.1.585.tar.gz b/dist/litellm-0.1.585.tar.gz new file mode 100644 index 000000000..2b97102a9 Binary files /dev/null and b/dist/litellm-0.1.585.tar.gz differ diff --git a/dist/litellm-0.1.586-py3-none-any.whl b/dist/litellm-0.1.586-py3-none-any.whl new file mode 100644 index 000000000..186bdda5f Binary files /dev/null and b/dist/litellm-0.1.586-py3-none-any.whl differ diff --git a/dist/litellm-0.1.586.tar.gz b/dist/litellm-0.1.586.tar.gz new file mode 100644 index 000000000..e25c44c58 Binary files /dev/null and b/dist/litellm-0.1.586.tar.gz differ diff --git a/dist/litellm-0.1.587-py3-none-any.whl b/dist/litellm-0.1.587-py3-none-any.whl new file mode 100644 index 000000000..7951f470f Binary files /dev/null and b/dist/litellm-0.1.587-py3-none-any.whl differ diff --git a/dist/litellm-0.1.587.tar.gz b/dist/litellm-0.1.587.tar.gz new file mode 100644 index 000000000..e0e7e6e3c Binary files /dev/null and b/dist/litellm-0.1.587.tar.gz differ diff --git a/litellm/__pycache__/__init__.cpython-311.pyc b/litellm/__pycache__/__init__.cpython-311.pyc index 2ec71aeb7..9690c652a 100644 Binary files a/litellm/__pycache__/__init__.cpython-311.pyc and b/litellm/__pycache__/__init__.cpython-311.pyc differ diff --git a/litellm/__pycache__/main.cpython-311.pyc b/litellm/__pycache__/main.cpython-311.pyc index 3bfc32f6e..6841fe91c 100644 Binary files a/litellm/__pycache__/main.cpython-311.pyc and b/litellm/__pycache__/main.cpython-311.pyc differ diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc index 81990b837..70a03ceef 100644 Binary files a/litellm/__pycache__/utils.cpython-311.pyc and b/litellm/__pycache__/utils.cpython-311.pyc differ diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index 98e5b31d5..0ffb4604b 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -14,15 +14,6 @@ class BudgetManager: 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 @@ -50,21 +41,28 @@ class BudgetManager: def update_cost(self, completion_obj: ModelResponse, user: str): cost = litellm.completion_cost(completion_response=completion_obj) model = completion_obj['model'] # if this throws an error try, model = completion_obj['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 {"current_user_cost": self.user_dict[user]["current_cost"], "current_model_cost": self.model_dict[model]} + if "model_cost" in self.user_dict[user]: + self.user_dict[user]["model_cost"][model] = cost + self.user_dict[user]["model_cost"].get(model, 0) + else: + self.user_dict[user]["model_cost"] = {model: cost} + return {"user": self.user_dict[user]} def get_current_cost(self, user): return self.user_dict[user].get("current_cost", 0) + + def get_model_cost(self, user): + return self.user_dict[user].get("model_cost", 0) + + def reset_cost(self, user): + self.user_dict[user]["current_cost"] = 0 + self.user_dict[user]["model_cost"] = {} + return {"user": self.user_dict[user]} def save_data(self): if self.type == "local": import json - # save the model dict - with open("model_cost.json", 'w') as json_file: - json.dump(self.model_dict, json_file, indent=4) # Indent for pretty formatting - # save the user dict with open("user_cost.json", 'w') as json_file: json.dump(self.user_dict, json_file, indent=4) # Indent for pretty formatting diff --git a/litellm/tests/data_map.txt b/litellm/tests/data_map.txt index 96c53e15e..e8077595f 100644 Binary files a/litellm/tests/data_map.txt and b/litellm/tests/data_map.txt differ diff --git a/litellm/tests/test_budget_manager.py b/litellm/tests/test_budget_manager.py index 97b86f555..31cdfa8c5 100644 --- a/litellm/tests/test_budget_manager.py +++ b/litellm/tests/test_budget_manager.py @@ -23,8 +23,6 @@ def test_user_budget_enough(): "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): response = completion(**data) print(budget_manager.update_cost(completion_obj=response, user=user)) @@ -60,4 +58,6 @@ def test_budget_save_to_disk(): ## Scenario 4: Loading budget from disk def test_budget_load_from_disk(): - budget_manager_2 = BudgetManager(type="local") \ No newline at end of file + budget_manager_2 = BudgetManager(type="local") + +## Scenario 5: Test get model cost from user dict \ No newline at end of file diff --git a/litellm/tests/user_cost.json b/litellm/tests/user_cost.json index fb2db88e3..b3411c777 100644 --- a/litellm/tests/user_cost.json +++ b/litellm/tests/user_cost.json @@ -1,7 +1,10 @@ { "1234": { "total_budget": 10, - "current_cost": 7.7e-05 + "current_cost": 7.3e-05, + "model_cost": { + "gpt-3.5-turbo-0613": 7.3e-05 + } }, "12345": { "total_budget": 0