phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-28 04:04:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy_server.py): fix max budget check to also fire slack alert

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-29 16:24:40 -07:00
parent fbe4d6d332
commit f7004a94df
3 changed files with 43 additions and 16 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_new_secret_config.yaml
View file

@ -5,6 +5,10 @@ model_list:
api_key: my-fake-key api_key: my-fake-key
api_base: https://exampleopenaiendpoint-production.up.railway.app/ api_base: https://exampleopenaiendpoint-production.up.railway.app/
litellm_settings:
max_budget: 600020
budget_duration: 30d
general_settings: general_settings:
master_key: sk-1234 master_key: sk-1234
proxy_batch_write_at: 60 # 👈 Frequency of batch writing logs to server (in seconds) proxy_batch_write_at: 60 # 👈 Frequency of batch writing logs to server (in seconds)

13
litellm/proxy/auth/auth_checks.py
View file

@ -81,14 +81,11 @@ def common_checks(
f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}" f"'user' param not passed in. 'enforce_user_param'={general_settings['enforce_user_param']}"
) )
# 6. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget # 6. [OPTIONAL] If 'litellm.max_budget' is set (>0), is proxy under budget
if ( if litellm.max_budget > 0 and global_proxy_spend is not None:
litellm.max_budget > 0 if global_proxy_spend > litellm.max_budget:
and global_proxy_spend is not None raise Exception(
and global_proxy_spend > litellm.max_budget f"'ExceededBudget: LiteLLM Proxy has exceeded its budget. Current spend: {global_proxy_spend}; Max Budget: {litellm.max_budget}"
): )
raise Exception(
f"'ExceededBudget: LiteLLM Proxy has exceeded its budget. Current spend: {global_proxy_spend}; Max Budget: {litellm.max_budget}"
)
return True return True

42
litellm/proxy/proxy_server.py
View file

@ -446,7 +446,7 @@ async def user_api_key_auth(
) )
if global_proxy_spend is None and prisma_client is not None: if global_proxy_spend is None and prisma_client is not None:
# get from db # get from db
sql_query = """SELECT SUM(spend) as total_spend FROM MONTHLYGLOBALSPEND;""" sql_query = """SELECT SUM(spend) as total_spend FROM "MonthlyGlobalSpend";"""
response = await prisma_client.db.query_raw(query=sql_query) response = await prisma_client.db.query_raw(query=sql_query)
@ -457,7 +457,21 @@ async def user_api_key_auth(
value=global_proxy_spend, value=global_proxy_spend,
ttl=60, ttl=60,
) )
if global_proxy_spend is not None:
user_info = {
"user_id": litellm_proxy_admin_name,
"max_budget": litellm.max_budget,
"spend": global_proxy_spend,
"user_email": "",
}
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=litellm.max_budget,
user_current_spend=global_proxy_spend,
type="user_and_proxy_budget",
user_info=user_info,
)
)
# run through common checks # run through common checks
_ = common_checks( _ = common_checks(
request_body=request_data, request_body=request_data,
@ -903,20 +917,33 @@ async def user_api_key_auth(
) )
if global_proxy_spend is None: if global_proxy_spend is None:
# get from db # get from db
sql_query = ( sql_query = """SELECT SUM(spend) as total_spend FROM "MonthlyGlobalSpend";"""
"""SELECT SUM(spend) as total_spend FROM MONTHLYGLOBALSPEND;"""
)
response = await prisma_client.db.query_raw(query=sql_query) response = await prisma_client.db.query_raw(query=sql_query)
global_proxy_spend = response[0].total_spend global_proxy_spend = response[0]["total_spend"]
await user_api_key_cache.async_set_cache( await user_api_key_cache.async_set_cache(
key="{}:spend".format(litellm_proxy_admin_name), key="{}:spend".format(litellm_proxy_admin_name),
value=global_proxy_spend, value=global_proxy_spend,
ttl=60, ttl=60,
) )
if global_proxy_spend is not None:
user_info = {
"user_id": litellm_proxy_admin_name,
"max_budget": litellm.max_budget,
"spend": global_proxy_spend,
"user_email": "",
}
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=litellm.max_budget,
user_current_spend=global_proxy_spend,
type="user_and_proxy_budget",
user_info=user_info,
)
)
_ = common_checks( _ = common_checks(
request_body=request_data, request_body=request_data,
team_object=_team_obj, team_object=_team_obj,
@ -4086,7 +4113,6 @@ async def generate_key_fn(
) )
_budget_id = getattr(_budget, "budget_id", None) _budget_id = getattr(_budget, "budget_id", None)
data_json = data.json() # type: ignore data_json = data.json() # type: ignore
# if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users # if we get max_budget passed to /key/generate, then use it as key_max_budget. Since generate_key_helper_fn is used to make new users
if "max_budget" in data_json: if "max_budget" in data_json:
data_json["key_max_budget"] = data_json.pop("max_budget", None) data_json["key_max_budget"] = data_json.pop("max_budget", None)