Merge pull request #3748 from BerriAI/litellm_webhook_support

webhook support for budget alerts
This commit is contained in:
Krish Dholakia 2024-05-20 18:42:04 -07:00 committed by GitHub
commit 1006284129
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 400 additions and 237 deletions

View file

@ -1,4 +1,4 @@
# 🚨 Alerting
# 🚨 Alerting / Webhooks
Get alerts for:
@ -61,8 +61,7 @@ curl -X GET 'http://localhost:4000/health/services?service=slack' \
-H 'Authorization: Bearer sk-1234'
```
## Advanced
### Opting into specific alert types
## Advanced - Opting into specific alert types
Set `alert_types` if you want to Opt into only specific alert types
@ -93,7 +92,7 @@ List[
```
### Using Discord Webhooks
## Advanced - Using Discord Webhooks
Discord provides a slack compatible webhook url that you can use for alerting
@ -125,3 +124,80 @@ environment_variables:
```
That's it ! You're ready to go !
## Advanced - [BETA] Webhooks for Budget Alerts
**Note**: This is a beta feature, so the spec might change.
Set a webhook to get notified for budget alerts.
1. Setup config.yaml
Add url to your environment, for testing you can use a link from [here](https://webhook.site/)
```bash
export WEBHOOK_URL="https://webhook.site/6ab090e8-c55f-4a23-b075-3209f5c57906"
```
Add 'webhook' to config.yaml
```yaml
general_settings:
alerting: ["webhook"] # 👈 KEY CHANGE
```
2. Start proxy
```bash
litellm --config /path/to/config.yaml
# RUNNING on http://0.0.0.0:4000
```
3. Test it!
```bash
curl -X GET --location 'http://0.0.0.0:4000/health/services?service=webhook' \
--header 'Authorization: Bearer sk-1234'
```
**Expected Response**
```bash
{
"spend": 1, # the spend for the 'event_group'
"max_budget": 0, # the 'max_budget' set for the 'event_group'
"token": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",
"user_id": "default_user_id",
"team_id": null,
"user_email": null,
"key_alias": null,
"projected_exceeded_data": null,
"projected_spend": null,
"event": "budget_crossed", # Literal["budget_crossed", "threshold_crossed", "projected_limit_exceeded"]
"event_group": "user",
"event_message": "User Budget: Budget Crossed"
}
```
**API Spec for Webhook Event**
- `spend` *float*: The current spend amount for the 'event_group'.
- `max_budget` *float*: The maximum allowed budget for the 'event_group'.
- `token` *str*: A hashed value of the key, used for authentication or identification purposes.
- `user_id` *str or null*: The ID of the user associated with the event (optional).
- `team_id` *str or null*: The ID of the team associated with the event (optional).
- `user_email` *str or null*: The email of the user associated with the event (optional).
- `key_alias` *str or null*: An alias for the key associated with the event (optional).
- `projected_exceeded_date` *str or null*: The date when the budget is projected to be exceeded, returned when 'soft_budget' is set for key (optional).
- `projected_spend` *float or null*: The projected spend amount, returned when 'soft_budget' is set for key (optional).
- `event` *Literal["budget_crossed", "threshold_crossed", "projected_limit_exceeded"]*: The type of event that triggered the webhook. Possible values are:
* "budget_crossed": Indicates that the spend has exceeded the max budget.
* "threshold_crossed": Indicates that spend has crossed a threshold (currently sent when 85% and 95% of budget is reached).
* "projected_limit_exceeded": For "key" only - Indicates that the projected spend is expected to exceed the soft budget threshold.
- `event_group` *Literal["user", "key", "team", "proxy"]*: The group associated with the event. Possible values are:
* "user": The event is related to a specific user.
* "key": The event is related to a specific key.
* "team": The event is related to a team.
* "proxy": The event is related to a proxy.
- `event_message` *str*: A human-readable description of the event.

View file

@ -1,7 +1,7 @@
#### What this does ####
# Class for sending Slack Alerts #
import dotenv, os
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import UserAPIKeyAuth, CallInfo
from litellm._logging import verbose_logger, verbose_proxy_logger
import litellm, threading
from typing import List, Literal, Any, Union, Optional, Dict
@ -39,6 +39,12 @@ class SlackAlertingArgs(LiteLLMBase):
budget_alert_ttl: int = 24 * 60 * 60 # 24 hours
class WebhookEvent(CallInfo):
event: Literal["budget_crossed", "threshold_crossed", "projected_limit_exceeded"]
event_group: Literal["user", "key", "team", "proxy"]
event_message: str # human-readable description of event
class DeploymentMetrics(LiteLLMBase):
"""
Metrics per deployment, stored in cache
@ -571,20 +577,32 @@ class SlackAlerting(CustomLogger):
alert_type="llm_requests_hanging",
)
async def failed_tracking_alert(self, error_message: str):
"""Raise alert when tracking failed for specific model"""
_cache: DualCache = self.internal_usage_cache
message = "Failed Tracking Cost for" + error_message
_cache_key = "budget_alerts:failed_tracking:{}".format(message)
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
await self.send_alert(
message=message, level="High", alert_type="budget_alerts"
)
await _cache.async_set_cache(
key=_cache_key,
value="SENT",
ttl=self.alerting_args.budget_alert_ttl,
)
async def budget_alerts(
self,
type: Literal[
"token_budget",
"user_budget",
"user_and_proxy_budget",
"failed_budgets",
"failed_tracking",
"team_budget",
"proxy_budget",
"projected_limit_exceeded",
],
user_max_budget: float,
user_current_spend: float,
user_info=None,
error_message="",
user_info: CallInfo,
):
## PREVENTITIVE ALERTING ## - https://github.com/BerriAI/litellm/issues/2727
# - Alert once within 24hr period
@ -598,128 +616,78 @@ class SlackAlerting(CustomLogger):
if "budget_alerts" not in self.alert_types:
return
_id: str = "default_id" # used for caching
if type == "user_and_proxy_budget":
user_info = dict(user_info)
user_id = user_info["user_id"]
_id = user_id
max_budget = user_info["max_budget"]
spend = user_info["spend"]
user_email = user_info["user_email"]
user_info = f"""\nUser ID: {user_id}\nMax Budget: ${max_budget}\nSpend: ${spend}\nUser Email: {user_email}"""
user_info_json = user_info.model_dump(exclude_none=True)
for k, v in user_info_json.items():
user_info_str = "\n{}: {}\n".format(k, v)
event: Optional[
Literal["budget_crossed", "threshold_crossed", "projected_limit_exceeded"]
] = None
event_group: Optional[Literal["user", "team", "key", "proxy"]] = None
event_message: str = ""
webhook_event: Optional[WebhookEvent] = None
if type == "proxy_budget":
event_group = "proxy"
event_message += "Proxy Budget: "
elif type == "user_budget":
event_group = "user"
event_message += "User Budget: "
_id = user_info.user_id or _id
elif type == "team_budget":
event_group = "team"
event_message += "Team Budget: "
_id = user_info.team_id or _id
elif type == "token_budget":
token_info = dict(user_info)
token = token_info["token"]
_id = token
spend = token_info["spend"]
max_budget = token_info["max_budget"]
user_id = token_info["user_id"]
user_info = f"""\nToken: {token}\nSpend: ${spend}\nMax Budget: ${max_budget}\nUser ID: {user_id}"""
elif type == "failed_tracking":
user_id = str(user_info)
_id = user_id
user_info = f"\nUser ID: {user_id}\n Error {error_message}"
message = "Failed Tracking Cost for" + user_info
_cache_key = "budget_alerts:failed_tracking:{}".format(_id)
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
await self.send_alert(
message=message, level="High", alert_type="budget_alerts"
)
await _cache.async_set_cache(
key=_cache_key,
value="SENT",
ttl=self.alerting_args.budget_alert_ttl,
)
return
elif type == "projected_limit_exceeded" and user_info is not None:
"""
Input variables:
user_info = {
"key_alias": key_alias,
"projected_spend": projected_spend,
"projected_exceeded_date": projected_exceeded_date,
}
user_max_budget=soft_limit,
user_current_spend=new_spend
"""
message = f"""\n🚨 `ProjectedLimitExceededError` 💸\n\n`Key Alias:` {user_info["key_alias"]} \n`Expected Day of Error`: {user_info["projected_exceeded_date"]} \n`Current Spend`: {user_current_spend} \n`Projected Spend at end of month`: {user_info["projected_spend"]} \n`Soft Limit`: {user_max_budget}"""
_cache_key = "budget_alerts:projected_limit_exceeded:{}".format(_id)
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
await self.send_alert(
message=message, level="High", alert_type="budget_alerts"
)
await _cache.async_set_cache(
key=_cache_key,
value="SENT",
ttl=self.alerting_args.budget_alert_ttl,
)
return
else:
user_info = str(user_info)
event_group = "key"
event_message += "Key Budget: "
_id = user_info.token
elif type == "projected_limit_exceeded":
event_group = "key"
event_message += "Key Budget: Projected Limit Exceeded"
event = "projected_limit_exceeded"
_id = user_info.token
# percent of max_budget left to spend
if user_max_budget > 0:
percent_left = (user_max_budget - user_current_spend) / user_max_budget
if user_info.max_budget > 0:
percent_left = (
user_info.max_budget - user_info.spend
) / user_info.max_budget
else:
percent_left = 0
verbose_proxy_logger.debug(
f"Budget Alerts: Percent left: {percent_left} for {user_info}"
)
# check if crossed budget
if user_current_spend >= user_max_budget:
verbose_proxy_logger.debug("Budget Crossed for %s", user_info)
message = "Budget Crossed for" + user_info
_cache_key = "budget_alerts:budget_crossed:{}".format(_id)
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
await self.send_alert(
message=message, level="High", alert_type="budget_alerts"
)
await _cache.async_set_cache(
key=_cache_key,
value="SENT",
ttl=self.alerting_args.budget_alert_ttl,
)
return
if user_info.spend >= user_info.max_budget:
event = "budget_crossed"
event_message += "Budget Crossed"
elif percent_left <= 0.05:
event = "threshold_crossed"
event_message += "5% Threshold Crossed"
elif percent_left <= 0.15:
event = "threshold_crossed"
event_message += "15% Threshold Crossed"
# check if 5% of max budget is left
if percent_left <= 0.05:
message = "5% budget left for" + user_info
_cache_key = "budget_alerts:5_perc_budget_crossed:{}".format(_id)
if event is not None and event_group is not None:
_cache_key = "budget_alerts:{}:{}".format(event, _id)
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
webhook_event = WebhookEvent(
event=event,
event_group=event_group,
event_message=event_message,
**user_info_json,
)
await self.send_alert(
message=message,
level="Medium",
alert_type="budget_alerts",
)
await _cache.async_set_cache(
key=_cache_key,
value="SENT",
ttl=self.alerting_args.budget_alert_ttl,
)
return
# check if 15% of max budget is left
if percent_left <= 0.15:
message = "15% budget left for" + user_info
_cache_key = "budget_alerts:15_perc_budget_crossed:{}".format(_id)
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
await self.send_alert(
message=message,
level="Low",
message=event_message + "\n\n" + user_info_str,
level="High",
alert_type="budget_alerts",
user_info=webhook_event,
)
await _cache.async_set_cache(
key=_cache_key,
value="SENT",
ttl=self.alerting_args.budget_alert_ttl,
)
return
return
@ -780,6 +748,34 @@ Model Info:
async def model_removed_alert(self, model_name: str):
pass
async def send_webhook_alert(self, webhook_event: WebhookEvent) -> bool:
"""
Sends structured alert to webhook, if set.
Currently only implemented for budget alerts
Returns -> True if sent, False if not.
"""
webhook_url = os.getenv("WEBHOOK_URL", None)
if webhook_url is None:
raise Exception("Missing webhook_url from environment")
payload = webhook_event.model_dump_json()
headers = {"Content-type": "application/json"}
response = await self.async_http_handler.post(
url=webhook_url,
headers=headers,
data=payload,
)
if response.status_code == 200:
return True
else:
print("Error sending webhook alert. Error=", response.text) # noqa
return False
async def send_alert(
self,
message: str,
@ -795,6 +791,7 @@ Model Info:
"new_model_added",
"cooldown_deployment",
],
user_info: Optional[WebhookEvent] = None,
**kwargs,
):
"""
@ -814,6 +811,16 @@ Model Info:
if self.alerting is None:
return
if (
"webhook" in self.alerting
and alert_type == "budget_alerts"
and user_info is not None
):
await self.send_webhook_alert(webhook_event=user_info)
if "slack" not in self.alerting:
return
if alert_type not in self.alert_types:
return

View file

@ -18,11 +18,3 @@ model_list:
router_settings:
enable_pre_call_checks: true
litellm_settings:
callbacks: ["detect_prompt_injection"]
prompt_injection_params:
heuristics_check: true
similarity_check: true
reject_as_response: true

View file

@ -1030,3 +1030,17 @@ class TokenCountResponse(LiteLLMBase):
request_model: str
model_used: str
tokenizer_type: str
class CallInfo(LiteLLMBase):
"""Used for slack budget alerting"""
spend: float
max_budget: float
token: str = Field(description="Hashed value of that key")
user_id: Optional[str] = None
team_id: Optional[str] = None
user_email: Optional[str] = None
key_alias: Optional[str] = None
projected_exceeded_date: Optional[str] = None
projected_spend: Optional[float] = None

View file

@ -590,17 +590,15 @@ async def user_api_key_auth(
ttl=UserAPIKeyCacheTTLEnum.global_proxy_spend.value,
)
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": "",
}
user_info = CallInfo(
user_id=litellm_proxy_admin_name,
max_budget=litellm.max_budget,
spend=global_proxy_spend,
token=valid_token["token"],
)
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",
type="proxy_budget",
user_info=user_info,
)
)
@ -923,12 +921,18 @@ async def user_api_key_auth(
user_max_budget is not None
and user_current_spend is not None
):
call_info = CallInfo(
token=valid_token.token,
spend=user_current_spend,
max_budget=user_max_budget,
user_id=_user.get("user_id", None),
user_email=_user.get("user_email", None),
key_alias=valid_token.key_alias,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=user_max_budget,
user_current_spend=user_current_spend,
type="user_and_proxy_budget",
user_info=_user,
type="user_budget",
user_info=call_info,
)
)
@ -948,12 +952,20 @@ async def user_api_key_auth(
user_max_budget is not None
and user_current_spend is not None
):
call_info = CallInfo(
token=valid_token.token,
spend=user_current_spend,
max_budget=user_max_budget,
user_id=getattr(user_id_information, "user_id", None),
user_email=getattr(
user_id_information, "user_email", None
),
key_alias=valid_token.key_alias,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=user_max_budget,
user_current_spend=user_current_spend,
type="user_budget",
user_info=user_id_information,
user_info=call_info,
)
)
@ -983,12 +995,17 @@ async def user_api_key_auth(
# Check 4. Token Spend is under budget
if valid_token.spend is not None and valid_token.max_budget is not None:
call_info = CallInfo(
token=valid_token.token,
spend=valid_token.spend,
max_budget=valid_token.max_budget,
user_id=valid_token.user_id,
team_id=valid_token.team_id,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=valid_token.max_budget,
user_current_spend=valid_token.spend,
type="token_budget",
user_info=valid_token,
user_info=call_info,
)
)
@ -1033,39 +1050,24 @@ async def user_api_key_auth(
raise Exception(
f"ExceededModelBudget: Current spend for model: {current_model_spend}; Max Budget for Model: {current_model_budget}"
)
# Check 6. Token spend is under Team budget
if (
valid_token.spend is not None
and hasattr(valid_token, "team_max_budget")
and valid_token.team_max_budget is not None
):
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=valid_token.team_max_budget,
user_current_spend=valid_token.spend,
type="token_budget",
user_info=valid_token,
)
)
if valid_token.spend >= valid_token.team_max_budget:
raise Exception(
f"ExceededTokenBudget: Current spend for token: {valid_token.spend}; Max Budget for Team: {valid_token.team_max_budget}"
)
# Check 7. Team spend is under Team budget
# Check 6. Team spend is under Team budget
if (
hasattr(valid_token, "team_spend")
and valid_token.team_spend is not None
and hasattr(valid_token, "team_max_budget")
and valid_token.team_max_budget is not None
):
call_info = CallInfo(
token=valid_token.token,
spend=valid_token.team_spend,
max_budget=valid_token.team_max_budget,
user_id=valid_token.user_id,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=valid_token.team_max_budget,
user_current_spend=valid_token.team_spend,
type="token_budget",
user_info=valid_token,
type="team_budget",
user_info=call_info,
)
)
@ -1111,18 +1113,17 @@ async def user_api_key_auth(
)
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": "",
}
call_info = CallInfo(
token=valid_token.token,
spend=global_proxy_spend,
max_budget=litellm.max_budget,
user_id=litellm_proxy_admin_name,
team_id=valid_token.team_id,
)
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,
type="proxy_budget",
user_info=call_info,
)
)
_ = common_checks(
@ -1514,13 +1515,8 @@ async def _PROXY_track_cost_callback(
model = kwargs.get("model", "")
metadata = kwargs.get("litellm_params", {}).get("metadata", {})
error_msg += f"\n Args to _PROXY_track_cost_callback\n model: {model}\n metadata: {metadata}\n"
user_id = user_id or "not-found"
asyncio.create_task(
proxy_logging_obj.budget_alerts(
user_max_budget=0,
user_current_spend=0,
type="failed_tracking",
user_info=user_id,
proxy_logging_obj.failed_tracking_alert(
error_message=error_msg,
)
)
@ -1732,14 +1728,14 @@ async def update_cache(
"""
### UPDATE KEY SPEND ###
async def _update_key_cache():
async def _update_key_cache(token: str, response_cost: float):
# Fetch the existing cost for the given token
if isinstance(token, str) and token.startswith("sk-"):
hashed_token = hash_token(token=token)
else:
hashed_token = token
verbose_proxy_logger.debug("_update_key_cache: hashed_token=%s", hashed_token)
existing_spend_obj = await user_api_key_cache.async_get_cache(key=hashed_token)
existing_spend_obj: LiteLLM_VerificationTokenView = await user_api_key_cache.async_get_cache(key=hashed_token) # type: ignore
verbose_proxy_logger.debug(
f"_update_key_cache: existing_spend_obj={existing_spend_obj}"
)
@ -1748,7 +1744,7 @@ async def update_cache(
)
if existing_spend_obj is None:
existing_spend = 0
existing_spend_obj = LiteLLM_VerificationTokenView()
existing_spend_obj = LiteLLM_VerificationTokenView(token=token)
else:
existing_spend = existing_spend_obj.spend
# Calculate the new cost by adding the existing cost and response_cost
@ -1762,29 +1758,36 @@ async def update_cache(
and (
_is_projected_spend_over_limit(
current_spend=new_spend,
soft_budget_limit=existing_spend_obj.litellm_budget_table.soft_budget,
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_budget_limit=existing_spend_obj.litellm_budget_table.get(
"soft_budget", None
),
) # type: ignore
soft_limit = existing_spend_obj.litellm_budget_table.get(
"soft_budget", float("inf")
)
call_info = CallInfo(
token=existing_spend_obj.token or "",
spend=new_spend,
key_alias=existing_spend_obj.key_alias,
max_budget=soft_limit,
user_id=existing_spend_obj.user_id,
projected_spend=projected_spend,
projected_exceeded_date=projected_exceeded_date,
)
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,
user_info=call_info,
)
)
# set cooldown on alert
@ -1794,7 +1797,7 @@ async def update_cache(
existing_spend_obj is not None
and getattr(existing_spend_obj, "team_spend", None) is not None
):
existing_team_spend = existing_spend_obj.team_spend
existing_team_spend = existing_spend_obj.team_spend or 0
# Calculate the new cost by adding the existing cost and response_cost
existing_spend_obj.team_spend = existing_team_spend + response_cost
@ -1911,8 +1914,8 @@ async def update_cache(
f"An error occurred updating end user cache: {str(e)}\n\n{traceback.format_exc()}"
)
if token is not None:
asyncio.create_task(_update_key_cache())
if token is not None and response_cost is not None:
asyncio.create_task(_update_key_cache(token=token, response_cost=response_cost))
asyncio.create_task(_update_user_cache())
@ -10277,7 +10280,7 @@ async def test_endpoint(request: Request):
async def health_services_endpoint(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
service: Literal[
"slack_budget_alerts", "langfuse", "slack", "openmeter"
"slack_budget_alerts", "langfuse", "slack", "openmeter", "webhook"
] = fastapi.Query(description="Specify the service being hit."),
):
"""
@ -10292,7 +10295,13 @@ async def health_services_endpoint(
raise HTTPException(
status_code=400, detail={"error": "Service must be specified."}
)
if service not in ["slack_budget_alerts", "langfuse", "slack", "openmeter"]:
if service not in [
"slack_budget_alerts",
"langfuse",
"slack",
"openmeter",
"webhook",
]:
raise HTTPException(
status_code=400,
detail={
@ -10328,6 +10337,20 @@ async def health_services_endpoint(
"message": "Mock LLM request made - check langfuse.",
}
if service == "webhook":
user_info = CallInfo(
token=user_api_key_dict.token or "",
spend=1,
max_budget=0,
user_id=user_api_key_dict.user_id,
key_alias=user_api_key_dict.key_alias,
team_id=user_api_key_dict.team_id,
)
await proxy_logging_obj.budget_alerts(
type="user_budget",
user_info=user_info,
)
if service == "slack" or service == "slack_budget_alerts":
if "slack" in general_settings.get("alerting", []):
# test_message = f"""\n🚨 `ProjectedLimitExceededError` 💸\n\n`Key Alias:` litellm-ui-test-alert \n`Expected Day of Error`: 28th March \n`Current Spend`: $100.00 \n`Projected Spend at end of month`: $1000.00 \n`Soft Limit`: $700"""
@ -10403,6 +10426,7 @@ async def health_services_endpoint(
},
)
except Exception as e:
traceback.print_exc()
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Authentication Error({str(e)})"),

View file

@ -11,6 +11,7 @@ from litellm.proxy._types import (
LiteLLM_EndUserTable,
LiteLLM_TeamTable,
Member,
CallInfo,
)
from litellm.caching import DualCache, RedisCache
from litellm.router import Deployment, ModelInfo, LiteLLM_Params
@ -312,30 +313,30 @@ class ProxyLogging:
raise e
return data
async def failed_tracking_alert(self, error_message: str):
if self.alerting is None:
return
await self.slack_alerting_instance.failed_tracking_alert(
error_message=error_message
)
async def budget_alerts(
self,
type: Literal[
"token_budget",
"user_budget",
"user_and_proxy_budget",
"failed_budgets",
"failed_tracking",
"team_budget",
"proxy_budget",
"projected_limit_exceeded",
],
user_max_budget: float,
user_current_spend: float,
user_info=None,
error_message="",
user_info: CallInfo,
):
if self.alerting is None:
# do nothing if alerting is not switched on
return
await self.slack_alerting_instance.budget_alerts(
type=type,
user_max_budget=user_max_budget,
user_current_spend=user_current_spend,
user_info=user_info,
error_message=error_message,
)
async def alerting_handler(
@ -391,7 +392,11 @@ class ProxyLogging:
for client in self.alerting:
if client == "slack":
await self.slack_alerting_instance.send_alert(
message=message, level=level, alert_type=alert_type, **extra_kwargs
message=message,
level=level,
alert_type=alert_type,
user_info=None,
**extra_kwargs,
)
elif client == "sentry":
if litellm.utils.sentry_sdk_instance is not None:

View file

@ -1,7 +1,7 @@
# What is this?
## Tests slack alerting on proxy logging object
import sys, json, uuid
import sys, json, uuid, random
import os
import io, asyncio
from datetime import datetime, timedelta
@ -22,6 +22,7 @@ import unittest.mock
from unittest.mock import AsyncMock
import pytest
from litellm.router import AlertingConfig, Router
from litellm.proxy._types import CallInfo
@pytest.mark.parametrize(
@ -123,7 +124,9 @@ from datetime import datetime, timedelta
@pytest.fixture
def slack_alerting():
return SlackAlerting(alerting_threshold=1, internal_usage_cache=DualCache())
return SlackAlerting(
alerting_threshold=1, internal_usage_cache=DualCache(), alerting=["slack"]
)
# Test for hanging LLM responses
@ -161,7 +164,10 @@ async def test_budget_alerts_crossed(slack_alerting):
user_current_spend = 101
with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert:
await slack_alerting.budget_alerts(
"user_budget", user_max_budget, user_current_spend
"user_budget",
user_info=CallInfo(
token="", spend=user_current_spend, max_budget=user_max_budget
),
)
mock_send_alert.assert_awaited_once()
@ -173,12 +179,18 @@ async def test_budget_alerts_crossed_again(slack_alerting):
user_current_spend = 101
with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert:
await slack_alerting.budget_alerts(
"user_budget", user_max_budget, user_current_spend
"user_budget",
user_info=CallInfo(
token="", spend=user_current_spend, max_budget=user_max_budget
),
)
mock_send_alert.assert_awaited_once()
mock_send_alert.reset_mock()
await slack_alerting.budget_alerts(
"user_budget", user_max_budget, user_current_spend
"user_budget",
user_info=CallInfo(
token="", spend=user_current_spend, max_budget=user_max_budget
),
)
mock_send_alert.assert_not_awaited()
@ -417,9 +429,8 @@ async def test_send_daily_reports_all_zero_or_none():
[
"token_budget",
"user_budget",
"user_and_proxy_budget",
"failed_budgets",
"failed_tracking",
"team_budget",
"proxy_budget",
"projected_limit_exceeded",
],
)
@ -428,25 +439,59 @@ async def test_send_token_budget_crossed_alerts(alerting_type):
slack_alerting = SlackAlerting()
with patch.object(slack_alerting, "send_alert", new=AsyncMock()) as mock_send_alert:
if alerting_type == "failed_tracking":
user_info = "ishaan@berri.ai"
else:
user_info = {
"token": "50e55ca5bfbd0759697538e8d23c0cd5031f52d9e19e176d7233b20c7c4d3403",
"spend": uuid.uuid4(),
"max_budget": None,
"user_id": "ishaan@berri.ai",
"user_email": "ishaan@berri.ai",
"key_alias": "my-test-key",
"projected_exceeded_date": "10/20/2024",
"projected_spend": 200,
}
user_info = {
"token": "50e55ca5bfbd0759697538e8d23c0cd5031f52d9e19e176d7233b20c7c4d3403",
"spend": 86,
"max_budget": 100,
"user_id": "ishaan@berri.ai",
"user_email": "ishaan@berri.ai",
"key_alias": "my-test-key",
"projected_exceeded_date": "10/20/2024",
"projected_spend": 200,
}
user_info = CallInfo(**user_info)
for _ in range(50):
await slack_alerting.budget_alerts(
type=alerting_type,
user_info=user_info,
user_current_spend=86,
user_max_budget=100,
)
mock_send_alert.assert_awaited_once()
@pytest.mark.parametrize(
"alerting_type",
[
"token_budget",
"user_budget",
"team_budget",
"proxy_budget",
"projected_limit_exceeded",
],
)
@pytest.mark.asyncio
async def test_webhook_alerting(alerting_type):
slack_alerting = SlackAlerting(alerting=["webhook"])
with patch.object(
slack_alerting, "send_webhook_alert", new=AsyncMock()
) as mock_send_alert:
user_info = {
"token": "50e55ca5bfbd0759697538e8d23c0cd5031f52d9e19e176d7233b20c7c4d3403",
"spend": 1,
"max_budget": 0,
"user_id": "ishaan@berri.ai",
"user_email": "ishaan@berri.ai",
"key_alias": "my-test-key",
"projected_exceeded_date": "10/20/2024",
"projected_spend": 200,
}
user_info = CallInfo(**user_info)
for _ in range(50):
await slack_alerting.budget_alerts(
type=alerting_type,
user_info=user_info,
)
mock_send_alert.assert_awaited_once()