update budget manager with hosted client

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

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 
+## 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")

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"