💸 API 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
from litellm import BudgetManager, completion
budget_manager = BudgetManager(project_name="test_project")
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)
# 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!"
use with Text Input / Output
Update cost by just passing in the text input / output and model name.
from litellm import BudgetManager
budget_manager = BudgetManager(project_name="test_project")
user = "12345"
budget_manager.create_budget(total_budget=10, user=user, duration="daily")
input_text = "hello world"
output_text = "it's a sunny day in san francisco"
model = "gpt-3.5-turbo"
budget_manager.update_cost(user=user, model=model, input_text=input_text, output_text=output_text) # 👈
print(budget_manager.get_current_cost(user))
advanced usage
In production, we will need to
- store user budgets in a database
- reset user budgets based on a set duration
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
budget_manager = BudgetManager(project_name="<my-unique-project>", client_type="hosted")
Complete Code
from litellm import BudgetManager, completion
budget_manager = BudgetManager(project_name="<my-unique-project>", client_type="hosted")
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!"
Self-hosted
To use your own db, set the BudgetManager client type to hosted and set the api_base.
Your api is expected to expose /get_budget and /set_budget endpoints. See code for details
Usage
budget_manager = BudgetManager(project_name="<my-unique-project>", client_type="hosted", api_base="your_custom_api")
Complete Code
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!"
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.
Below is a list of public functions exposed by the Budget Manager class and their input/outputs.
init
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.
create_budget
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
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
def get_total_budget(self, user: str)
Returns the total budget of a user.
user(str): user id.
update_cost
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
def get_current_cost(self, user: str)
Returns the current cost of a user.
user(str): The user id.
get_model_cost
def get_model_cost(self, user: str)
Returns the model cost of a user.
user(str): The user id.
is_valid_user
def is_valid_user(self, user: str) -> bool
Checks if a user is valid.
user(str): The user id.
get_users
def get_users(self)
Returns a list of all users.
reset_cost
def reset_cost(self, user: str)
Resets the cost of a user.
user(str): The user id.
reset_on_duration
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
def update_budget_all_users(self)
Updates the budget for all users.
save_data
def save_data(self)
Stores the user dictionary.