diff --git a/litellm/__init__.py b/litellm/__init__.py index 920833c113..a166dcfcd5 100644 --- a/litellm/__init__.py +++ b/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 \ No newline at end of file diff --git a/litellm/__pycache__/__init__.cpython-311.pyc b/litellm/__pycache__/__init__.cpython-311.pyc index f23105fe1b..8878d126f8 100644 Binary files a/litellm/__pycache__/__init__.cpython-311.pyc and b/litellm/__pycache__/__init__.cpython-311.pyc differ diff --git a/litellm/__pycache__/utils.cpython-311.pyc b/litellm/__pycache__/utils.cpython-311.pyc index 807a839953..e778fcdd33 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 bdda4c1729..729ceb04c3 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -1,8 +1,36 @@ +import os, json import litellm from litellm.utils import ModelResponse + class BudgetManager: - def __init__(self): - self.user_dict = {} + def __init__(self, type: str): + 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'] - self.store_llm_costs(cost, model) + 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 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): - try: - import json - with open("cost.json", 'r') as json_file: - cost_dict = json.load(json_file) - if model in cost_dict: - cost_dict[model] += cost - else: - cost_dict[model] = cost - with open("cost.json", 'w') as json_file: - json.dump(cost_dict, json_file, indent=4) # Indent for pretty formatting - except Exception as e: - return # [Non blocking] \ No newline at end of file + 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/model_cost.json b/litellm/tests/model_cost.json new file mode 100644 index 0000000000..e581433952 --- /dev/null +++ b/litellm/tests/model_cost.json @@ -0,0 +1,3 @@ +{ + "gpt-3.5-turbo-0613": 7.7e-05 +} \ No newline at end of file diff --git a/litellm/tests/test_api_manager.py b/litellm/tests/test_budget_manager.py similarity index 72% rename from litellm/tests/test_api_manager.py rename to litellm/tests/test_budget_manager.py index 18d507a7fa..619b4785f5 100644 --- a/litellm/tests/test_api_manager.py +++ b/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() \ No newline at end of file +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() \ No newline at end of file diff --git a/litellm/tests/user_cost.json b/litellm/tests/user_cost.json new file mode 100644 index 0000000000..fb2db88e3c --- /dev/null +++ b/litellm/tests/user_cost.json @@ -0,0 +1,9 @@ +{ + "1234": { + "total_budget": 10, + "current_cost": 7.7e-05 + }, + "12345": { + "total_budget": 0 + } +} \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 907a4b7777..0907ff1acb 100644 --- a/pyproject.toml +++ b/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"