Skip to main content

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.

info

If you want a server to manage user keys, budgets, etc. use our LiteLLM Proxy Server

LiteLLM exposes:

  • litellm.max_budget: a global variable you can use to set the max budget (in USD) across all your litellm calls. If this budget is exceeded, it will raise a BudgetExceededError
  • BudgetManager: A 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.
  • LiteLLM Proxy Server: A server to call 100+ LLMs with an openai-compatible endpoint. Manages user budgets, spend tracking, load balancing etc.

quick start​

import litellm, os 
from litellm import completion

# set env variable 
os.environ["OPENAI_API_KEY"] = "your-api-key"

litellm.max_budget = 0.001 # sets a max budget of $0.001

messages = [{"role": "user", "content": "Hey, how's it going"}]
completion(model="gpt-4", messages=messages)
print(litellm._current_cost)
completion(model="gpt-4", messages=messages)

User-based rate limiting​

Open In Colab
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!"

Implementation Code

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.