updates to budget manager naming schema

docs/my-website/docs/budget_manager.md

@@ -11,7 +11,8 @@ budget_manager = BudgetManager(project_name="test_project")
 
 user = "1234"
 
-# create a budget for a user
+# create a budget if new user user
+if not budget_manager.is_valid_user(user):
     budget_manager.create_budget(total_budget=10, user=user)
 
 # check if a given call can be made
@@ -40,7 +41,7 @@ budget_manager.get_model_cost(user=user) # {"gpt-3.5-turbo-0613": 7.3e-05}
 
 ### save budget to disk
 
-When you call `save_data()` it will check for the self.type (by default this is set to client), and save the dictionary to a local `user_cost.json` file. 
+When you call `save_data()` it will check for the self.client_type (by default this is set to local), and save the dictionary to a local `user_cost.json` file. 
 
 ```python
 # ...
@@ -49,20 +50,20 @@ budget_manager.save_data() # 👈 save to user_cost.json()
 
 [**Implementation Code**](https://github.com/BerriAI/litellm/blob/817798c692207569a17c26186d10541aa83f04e7/litellm/budget_manager.py#L83)
 
-### save budget to db (LiteLLM)
+### save budget to hosted db (LiteLLM)
 
-Set the BudgetManager type to `client`.
+Set the BudgetManager type to `client_type`.
 ```python
-budget_manager = BudgetManager(project_name="test_project", type="client")
+budget_manager = BudgetManager(project_name="test_project", client_type="hosted")
 # ...
 budget_manager.save_data() # 👈 saved to hosted db 
 ```
 
 [**Implementation Code**](https://github.com/BerriAI/litellm/blob/817798c692207569a17c26186d10541aa83f04e7/litellm/budget_manager.py#L11)
 
-### save budget to db (self-hosted)
+### save budget to hosted db (self-hosted)
 
-Set the BudgetManager type to `client`. Overwrite the api_base
+Set the BudgetManager type to `client_type`. Overwrite the api_base
 ```python
 budget_manager = BudgetManager(project_name="test_project", type="client", api_base="your_custom_api")
 # ...

litellm/budget_manager.py

@@ -5,8 +5,8 @@ import requests
 from typing import Optional
 
 class BudgetManager:
-    def __init__(self, project_name: str, type: str = "local", api_base: Optional[str] = None):
-        self.type = type
+    def __init__(self, project_name: str, client_type: str = "local", api_base: Optional[str] = None):
+        self.client_type = client_type
         self.project_name = project_name
         self.api_base = api_base or "https://api.litellm.ai"
         ## load the data or init the initial dictionaries
@@ -17,7 +17,7 @@ class BudgetManager:
             print(print_statement)
     
     def load_data(self):
-        if self.type == "local":
+        if self.client_type == "local":
             # Check if user dict file exists
             if os.path.isfile("user_cost.json"):
                 # Load the user dict
@@ -26,7 +26,8 @@ class BudgetManager:
             else:
                 self.print_verbose("User Dictionary not found!")
                 self.user_dict = {} 
-        elif self.type == "client":
+            self.print_verbose(f"user dict from local: {self.user_dict}")
+        elif self.client_type == "hosted":
             # Load the user_dict from hosted db
             url = self.api_base + "/get_budget"
             headers = {'Content-Type': 'application/json'}
@@ -71,6 +72,9 @@ class BudgetManager:
     def get_model_cost(self, user):
         return self.user_dict[user].get("model_cost", 0)
     
+    def is_valid_user(self, user: str) -> bool:
+        return user in self.user_dict
+    
     def get_users(self):
         return list(self.user_dict.keys())
     
@@ -80,14 +84,14 @@ class BudgetManager:
         return {"user": self.user_dict[user]}
 
     def save_data(self):
-        if self.type == "local":
+        if self.client_type == "local":
             import json 
             
             # 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":
+        elif self.client_type == "hosted":
             url = self.api_base + "/set_budget"
             headers = {'Content-Type': 'application/json'}
             data = {

litellm/tests/test_budget_manager.py

@@ -9,9 +9,10 @@ sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path
 import litellm 
+litellm.set_verbose = True
 from litellm import BudgetManager, completion 
 
-budget_manager = BudgetManager(project_name="test_project")
+budget_manager = BudgetManager(project_name="test_project", client_type="hosted")
 
 ## Scenario 1: User budget enough to make call
 def test_user_budget_enough():
@@ -32,8 +33,10 @@ def test_user_budget_enough():
             response = "Sorry - no budget!"
 
         print(f"response: {response}")
-    except:
-        pytest.fail(f"An error occurred")
+    except Exception as e:
+        pytest.fail(f"An error occurred - {str(e)}")
+
+test_user_budget_enough()
 
 ## Scenario 2: User budget not enough to make call
 def test_user_budget_not_enough():

pyproject.toml

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