From 2a555eddfa75ade48ede6b496b499d9e4279992f Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 12 Sep 2023 20:25:59 -0700
Subject: [PATCH] update budget manager docs

---
 docs/my-website/docs/budget_manager.md | 188 +++++++++++++++++++++----
 litellm/budget_manager.py              |   8 +-
 pyproject.toml                         |   2 +-
 3 files changed, 168 insertions(+), 30 deletions(-)

diff --git a/docs/my-website/docs/budget_manager.md b/docs/my-website/docs/budget_manager.md
index f77087461..615bbe99d 100644
--- a/docs/my-website/docs/budget_manager.md
+++ b/docs/my-website/docs/budget_manager.md
@@ -1,6 +1,11 @@
-# 💸 User-based Rate Limiting
+import Tabs from '@theme/Tabs';
+import TabItem from '@theme/TabItem';
 
-LiteLLM exposes the `BudgetManager` class to help with user-based rate limiting. 
+# 💸 User Budget Manager
+
+Don't want to get crazy bills because either while you're calling LLM APIs **or** while your users are calling them? use this. 
+
+LiteLLM exposes the `BudgetManager` class to help set budgets per user. BudgetManager creates a dictionary to manage the user budgets, where the key is user and the object is their current cost + model-specific costs. 
 
 ## quick start
 
@@ -26,48 +31,177 @@ else:
 [**Implementation Code**](https://github.com/BerriAI/litellm/blob/main/litellm/budget_manager.py)
 
 ## advanced usage
+In production, we will need to 
+* store user budgets in a database
+* reset user budgets based on a set duration 
 
-BudgetManager creates a dictionary to manage the user budgets, where the key is user and the object is their current cost + model-specific costs. By default this is saved to a local, but you can change it to be stored to a hosted client (either self-hosted or LiteLLM one).
 
-### get model-breakdown per user 
 
+### LiteLLM API
+
+The LiteLLM API provides both. It stores the user object in a hosted db, and runs a cron job daily to reset user-budgets based on the set duration (e.g. reset budget daily/weekly/monthly/etc.). 
+
+**Usage**
 ```python
+budget_manager = BudgetManager(project_name="<my-unique-project>", client_type="hosted")
+```
+
+**Complete Code**
+```python
+from litellm import BudgetManager, completion 
+
+budget_manager = BudgetManager(project_name="<my-unique-project>", client_type="hosted")
+
 user = "1234"
-# ...
-budget_manager.get_model_cost(user=user) # {"gpt-3.5-turbo-0613": 7.3e-05}
+
+# create a budget if new user user
+if not budget_manager.is_valid_user(user):
+    budget_manager.create_budget(total_budget=10, user=user, duration="monthly") # 👈 duration = 'daily'/'weekly'/'monthly'/'yearly'
+
+# check if a given call can be made
+if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user):
+    response = completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hey, how's it going?"}])
+    budget_manager.update_cost(completion_obj=response, user=user)
+else:
+    response = "Sorry - no budget!"
 ```
 
-[**Implementation Code**](https://github.com/BerriAI/litellm/blob/817798c692207569a17c26186d10541aa83f04e7/litellm/budget_manager.py#L71)
+### Self-hosted
 
-### save budget to disk
+To use your own db, set the BudgetManager client type to `hosted` **and** set the api_base. 
 
-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. 
+Your api is expected to expose `/get_budget` and `/set_budget` endpoints. [See code for details](https://github.com/BerriAI/litellm/blob/27f1051792176a7eb1fe3b72b72bccd6378d24e9/litellm/budget_manager.py#L7)
 
+**Usage**
 ```python
-# ...
-budget_manager.save_data() # 👈 save to user_cost.json()
+budget_manager = BudgetManager(project_name="<my-unique-project>", client_type="hosted", api_base="your_custom_api")
 ```
-
-[**Implementation Code**](https://github.com/BerriAI/litellm/blob/817798c692207569a17c26186d10541aa83f04e7/litellm/budget_manager.py#L83)
-
-### save budget to hosted db (LiteLLM)
-
-Set the BudgetManager client type to `hosted`.
+**Complete Code**
 ```python
-budget_manager = BudgetManager(project_name="test_project", client_type="hosted")
-# ...
-budget_manager.save_data() # 👈 saved to hosted db 
+from litellm import BudgetManager, completion 
+
+budget_manager = BudgetManager(project_name="<my-unique-project>", client_type="hosted", api_base="your_custom_api")
+
+user = "1234"
+
+# create a budget if new user user
+if not budget_manager.is_valid_user(user):
+    budget_manager.create_budget(total_budget=10, user=user, duration="monthly") # 👈 duration = 'daily'/'weekly'/'monthly'/'yearly'
+
+# check if a given call can be made
+if budget_manager.get_current_cost(user=user) <= budget_manager.get_total_budget(user):
+    response = completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hey, how's it going?"}])
+    budget_manager.update_cost(completion_obj=response, user=user)
+else:
+    response = "Sorry - no budget!"
 ```
 
-[**Implementation Code**](https://github.com/BerriAI/litellm/blob/817798c692207569a17c26186d10541aa83f04e7/litellm/budget_manager.py#L11)
+## Budget Manager Class
+The `BudgetManager` class is used to manage budgets for different users. It provides various functions to create, update, and retrieve budget information. 
 
-### save budget to hosted db (self-hosted)
+Below is a list of public functions exposed by the Budget Manager class and their input/outputs. 
 
-Set the BudgetManager client type to `hosted` **And** Overwrite the api_base
+### __init__
 ```python
-budget_manager = BudgetManager(project_name="test_project", type="client", api_base="your_custom_api")
-# ...
-budget_manager.save_data() # 👈 saved to self-hosted db 
+def __init__(self, project_name: str, client_type: str = "local", api_base: Optional[str] = None)
 ```
+- `project_name` (str): The name of the project.
+- `client_type` (str): The client type ("local" or "hosted"). Defaults to "local".
+- `api_base` (Optional[str]): The base URL of the API. Defaults to None.
 
-[**Implementation Code**](https://github.com/BerriAI/litellm/blob/817798c692207569a17c26186d10541aa83f04e7/litellm/budget_manager.py#L11)
\ No newline at end of file
+
+### create_budget
+```python
+def create_budget(self, total_budget: float, user: str, duration: Literal["daily", "weekly", "monthly", "yearly"], created_at: float = time.time())
+```
+Creates a budget for a user.
+
+- `total_budget` (float): The total budget of the user.
+- `user` (str): The user id.
+- `duration` (Literal["daily", "weekly", "monthly", "yearly"]): The budget duration.
+- `created_at` (float): The creation time. Default is the current time.
+
+### projected_cost
+```python
+def projected_cost(self, model: str, messages: list, user: str)
+```
+Computes the projected cost for a session.
+
+- `model` (str): The name of the model.
+- `messages` (list): The list of messages.
+- `user` (str): The user id.
+
+### get_total_budget
+```python
+def get_total_budget(self, user: str)
+```
+Returns the total budget of a user.
+
+- `user` (str): user id.
+
+### update_cost
+```python
+def update_cost(self, completion_obj: ModelResponse, user: str)
+```
+Updates the user's cost.
+
+- `completion_obj` (ModelResponse): The completion object received from the model.
+- `user` (str): The user id.
+
+### get_current_cost
+```python
+def get_current_cost(self, user: str)
+```
+Returns the current cost of a user.
+
+- `user` (str): The user id.
+
+### get_model_cost 
+```python
+def get_model_cost(self, user: str)
+```
+Returns the model cost of a user.
+
+- `user` (str): The user id.
+
+### is_valid_user 
+```python
+def is_valid_user(self, user: str) -> bool
+```
+Checks if a user is valid.
+
+- `user` (str): The user id.
+
+### get_users 
+```python
+def get_users(self)
+```
+Returns a list of all users.
+
+### reset_cost 
+```python
+def reset_cost(self, user: str)
+```
+Resets the cost of a user.
+
+- `user` (str): The user id.
+
+### reset_on_duration 
+```python
+def reset_on_duration(self, user: str)
+```
+Resets the cost of a user based on the duration.
+
+- `user` (str): The user id.
+
+### update_budget_all_users 
+```python
+def update_budget_all_users(self)
+```
+Updates the budget for all users.
+
+### save_data 
+```python
+def save_data(self)
+```
+Stores the user dictionary.
\ No newline at end of file
diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py
index ff7de9b7b..bb7b5dff9 100644
--- a/litellm/budget_manager.py
+++ b/litellm/budget_manager.py
@@ -41,7 +41,11 @@ class BudgetManager:
             else:
                 self.user_dict = response["data"]
 
-    def create_budget(self, total_budget: float, user: str, duration: Literal["daily", "weekly", "monthly", "yearly"], created_at: float = time.time()): 
+    def create_budget(self, total_budget: float, user: str, duration: Literal["daily", "weekly", "monthly", "yearly"] = None, created_at: float = time.time()): 
+        self.user_dict[user] = {"total_budget": total_budget}
+        if duration is None:
+            return self.user_dict[user]
+        
         if duration == 'daily':
             duration_in_days = 1
         elif duration == 'weekly':
@@ -51,7 +55,7 @@ class BudgetManager:
         elif duration == 'yearly':
             duration_in_days = 365
         else:
-            raise ValueError('Invalid duration')
+            raise ValueError("""duration needs to be one of ["daily", "weekly", "monthly", "yearly"]""")
         self.user_dict[user] = {"total_budget": total_budget, "duration": duration_in_days, "created_at": created_at, "last_updated_at": created_at}
         return self.user_dict[user]
     
diff --git a/pyproject.toml b/pyproject.toml
index 117814b8c..0c26a5e66 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "0.1.607"
+version = "0.1.608"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"