phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(proxy_server.py): actual implementation of slack soft budget alerting

Browse source
This commit is contained in:
Krrish Dholakia 2024-03-02 18:34:18 -08:00
parent b30cbd0d55
commit ac085a4643
5 changed files with 142 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

50
litellm/proxy/proxy_server.py
View file

@ -94,6 +94,8 @@ from litellm.proxy.utils import (
_read_request_body,
_is_valid_team_configs,
_is_user_proxy_admin,
_is_projected_spend_over_limit,
_get_projected_spend_over_limit,
)
from litellm.proxy.secret_managers.google_kms import load_google_kms
import pydantic
@ -360,7 +362,6 @@ async def user_api_key_auth(
valid_token = await prisma_client.get_data(
token=api_key, table_name="combined_view"
)
elif custom_db_client is not None:
try:
valid_token = await custom_db_client.get_data(
@ -932,6 +933,7 @@ async def _PROXY_track_cost_callback(
f"user_api_key {user_api_key}, prisma_client: {prisma_client}, custom_db_client: {custom_db_client}"
)
if user_api_key is not None:
## UPDATE DATABASE
await update_database(
token=user_api_key,
response_cost=response_cost,
@ -1085,6 +1087,39 @@ async def update_database(
# Calculate the new cost by adding the existing cost and response_cost
new_spend = existing_spend + response_cost
## CHECK IF USER PROJECTED SPEND > SOFT LIMIT
soft_budget_cooldown = existing_spend_obj.soft_budget_cooldown
if existing_spend_obj.soft_budget_cooldown == False and (
_is_projected_spend_over_limit(
current_spend=new_spend,
soft_budget_limit=existing_spend_obj.litellm_budget_table.soft_budget,
)
== True
):
key_alias = existing_spend_obj.key_alias
projected_spend, projected_exceeded_date = (
_get_projected_spend_over_limit(
current_spend=new_spend,
soft_budget_limit=existing_spend_obj.litellm_budget_table.soft_budget,
)
)
soft_limit = existing_spend_obj.litellm_budget_table.soft_budget
user_info = {
"key_alias": key_alias,
"projected_spend": projected_spend,
"projected_exceeded_date": projected_exceeded_date,
}
# alert user
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="projected_limit_exceeded",
user_info=user_info,
user_max_budget=soft_limit,
user_current_spend=new_spend,
)
)
# set cooldown on alert
soft_budget_cooldown = True
# track cost per model, for the given key
spend_per_model = existing_spend_obj.model_spend or {}
current_model = kwargs.get("model")
@ -1100,7 +1135,11 @@ async def update_database(
# Update the cost column for the given token
await prisma_client.update_data(
token=token,
data={"spend": new_spend, "model_spend": spend_per_model},
data={
"spend": new_spend,
"model_spend": spend_per_model,
"soft_budget_cooldown": soft_budget_cooldown,
},
)
valid_token = user_api_key_cache.get_cache(key=token)
@ -1874,17 +1913,18 @@ async def generate_key_helper_fn(
allowed_cache_controls = allowed_cache_controls
# TODO: @ishaan-jaff: Migrate all budget tracking to use LiteLLM_BudgetTable
if prisma_client is not None:
_budget_id = None
if prisma_client is not None and key_soft_budget is not None:
# create the Budget Row for the LiteLLM Verification Token
budget_row = LiteLLM_BudgetTable(
soft_budget=key_soft_budget or litellm.default_soft_budget,
soft_budget=key_soft_budget,
model_max_budget=model_max_budget or {},
created_by=user_id,
updated_by=user_id,
)
new_budget = prisma_client.jsonify_object(budget_row.json(exclude_none=True))
_budget = await prisma_client.db.litellm_budgettable.create(data={**new_budget}) # type: ignore
_budget_id = getattr(_budget, "id", None)
_budget_id = getattr(_budget, "budget_id", None)
try:
# Create a new verification token (you may want to enhance this logic based on your needs)