diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index 0ffb4604b9e..0d5b6e7fea6 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -1,10 +1,13 @@ import os, json import litellm from litellm.utils import ModelResponse +import requests class BudgetManager: - def __init__(self, type: str): + def __init__(self, project_name: str, type: str = "client", api_base: str = None): self.type = type + self.project_name = project_name + self.api_base = api_base or "https://api.litellm.ai" ## load the data or init the initial dictionaries self.load_data() @@ -22,6 +25,19 @@ class BudgetManager: else: self.print_verbose("User Dictionary not found!") self.user_dict = {} + elif self.type == "client": + # Load the user_dict from hosted db + url = self.api_base + "/get_budget" + headers = {'Content-Type': 'application/json'} + data = { + 'project_name' : self.project_name + } + response = requests.post(url, headers=headers, json=data) + response = response.json() + if response["status"] == "error": + self.user_dict = {} # assume this means the user dict hasn't been stored yet + else: + self.user_dict = response["data"] def create_budget(self, total_budget: float, user: str): self.user_dict[user] = {"total_budget": total_budget} @@ -66,3 +82,14 @@ class BudgetManager: # 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 + return {"status": "success"} + elif self.type == "client": + url = self.api_base + "/set_budget" + headers = {'Content-Type': 'application/json'} + data = { + 'project_name' : self.project_name, + "user_dict": self.user_dict + } + response = requests.post(url, headers=headers, json=data) + response = response.json() + return response \ No newline at end of file diff --git a/litellm/tests/test_budget_manager.py b/litellm/tests/test_budget_manager.py index 31cdfa8c5d6..e8506548458 100644 --- a/litellm/tests/test_budget_manager.py +++ b/litellm/tests/test_budget_manager.py @@ -1,8 +1,9 @@ #### What this tests #### # This tests calling batch_completions by running 100 messages together -import sys, os +import sys, os, json import traceback +import pytest sys.path.insert( 0, os.path.abspath("../..") @@ -10,54 +11,60 @@ sys.path.insert( import litellm from litellm import BudgetManager, completion -budget_manager = BudgetManager(type="local") +budget_manager = BudgetManager(project_name="test_project") ## Scenario 1: User budget enough to make call def test_user_budget_enough(): - user = "1234" - # create a budget for a user - budget_manager.create_budget(total_budget=10, user=user) + try: + user = "1234" + # create a budget for a user + budget_manager.create_budget(total_budget=10, user=user) - # check if a given call can be made - data = { - "model": "gpt-3.5-turbo", - "messages": [{"role": "user", "content": "Hey, how's it going?"}] - } - 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!" + # check if a given call can be made + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hey, how's it going?"}] + } + 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}") + print(f"response: {response}") + except: + pytest.fail(f"An error occurred") ## Scenario 2: User budget not enough to make call def test_user_budget_not_enough(): - user = "12345" - # create a budget for a user - budget_manager.create_budget(total_budget=0, user=user) + try: + user = "12345" + # create a budget for a user + budget_manager.create_budget(total_budget=0, user=user) - # check if a given call can be made - data = { - "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=user): - response = completion(**data) - print(budget_manager.update_cost(completion_obj=response, user=user)) - else: - response = "Sorry - no budget!" + # check if a given call can be made + data = { + "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=user): + response = completion(**data) + print(budget_manager.update_cost(completion_obj=response, user=user)) + else: + response = "Sorry - no budget!" - print(f"response: {response}") + print(f"response: {response}") + except: + pytest.fail(f"An error occurred") -## Scenario 3: Saving budget to disk -def test_budget_save_to_disk(): - budget_manager.save_data() - -## Scenario 4: Loading budget from disk -def test_budget_load_from_disk(): - budget_manager_2 = BudgetManager(type="local") - -## Scenario 5: Test get model cost from user dict \ No newline at end of file +## Scenario 3: Saving budget to client +def test_save_user_budget(): + try: + response = budget_manager.save_data() + if response["status"] == "error": + raise Exception(f"An error occurred - {json.dumps(response)}") + print(response) + except: + pytest.fail(f"An error occurred") \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 0dfbbbefa5d..0c00739e8bd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "0.1.591" +version = "0.1.592" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT License"