feat: add configurable multi-threshold budget alerts for virtual keys

Users can set metadata.max_budget_alert_emails as a JSON map of threshold
percentages to email recipients on virtual keys. When configured, the email
handler loops over each threshold, checks per-threshold dedup cache, and
sends to the configured recipients (auto-including the key owner's email).

When no map is set, the existing single 80% threshold behavior is preserved
unchanged. Teams support is out of scope for this v0.
This commit is contained in:
Ryan Crabbe
2026-04-17 17:12:35 -07:00
parent c0fc4c4234
commit 779a9fab8e
5 changed files with 509 additions and 45 deletions
@@ -47,6 +47,15 @@ from litellm.secret_managers.main import get_secret_bool
from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL
def _parse_email_list(raw) -> List[str]:
"""Parse emails from a list or comma-separated string."""
if isinstance(raw, list):
return [e.strip() for e in raw if isinstance(e, str) and e.strip()]
elif isinstance(raw, str):
return [e.strip() for e in raw.split(",") if e.strip()]
return []
class BaseEmailLogger(CustomLogger):
DEFAULT_LITELLM_EMAIL = "notifications@alerts.litellm.ai"
DEFAULT_SUPPORT_EMAIL = "support@berri.ai"
@@ -312,17 +321,22 @@ class BaseEmailLogger(CustomLogger):
)
pass
async def send_max_budget_alert_email(self, event: WebhookEvent):
async def send_max_budget_alert_email(
self,
event: WebhookEvent,
threshold_pct: Optional[int] = None,
recipient_emails: Optional[List[str]] = None,
):
"""
Send email to user when max budget alert threshold is reached
"""
email_params = await self._get_email_params(
email_event=EmailEvent.max_budget_alert,
user_id=event.user_id,
user_email=event.user_email,
event_message=event.event_message,
)
Send email to user when max budget alert threshold is reached.
Args:
event: The webhook event with spend/budget info
threshold_pct: Override percentage for multi-threshold alerts (e.g. 50, 75, 100).
When None, uses EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE (old behavior).
recipient_emails: Override recipient list for multi-threshold alerts.
When None, resolves single owner email via _get_email_params (old behavior).
"""
verbose_proxy_logger.debug(
f"send_max_budget_alert_email_event: {json.dumps(event.model_dump(exclude_none=True), indent=4, default=str)}"
)
@@ -334,30 +348,65 @@ class BaseEmailLogger(CustomLogger):
)
# Calculate percentage and alert threshold
percentage = int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100)
percentage = threshold_pct if threshold_pct is not None else int(
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100
)
threshold_fraction = percentage / 100.0
alert_threshold_str = (
f"${event.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE:.2f}"
f"${event.max_budget * threshold_fraction:.2f}"
if event.max_budget is not None
else "N/A"
)
email_html_content = MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
recipient_email=email_params.recipient_email,
percentage=percentage,
spend=spend_str,
max_budget=max_budget_str,
alert_threshold=alert_threshold_str,
base_url=email_params.base_url,
email_support_contact=email_params.support_contact,
)
await self.send_email(
from_email=self.DEFAULT_LITELLM_EMAIL,
to_email=[email_params.recipient_email],
subject=email_params.subject,
html_body=email_html_content,
)
pass
if recipient_emails is not None:
# Multi-threshold path: batch send with generic key-based greeting
email_params = await self._get_email_params(
email_event=EmailEvent.max_budget_alert,
user_id=event.user_id,
user_email=event.user_email or recipient_emails[0],
event_message=event.event_message,
)
greeting = event.user_email or event.key_alias or event.token or ""
email_html_content = MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
recipient_email=greeting,
percentage=percentage,
spend=spend_str,
max_budget=max_budget_str,
alert_threshold=alert_threshold_str,
base_url=email_params.base_url,
email_support_contact=email_params.support_contact,
)
await self.send_email(
from_email=self.DEFAULT_LITELLM_EMAIL,
to_email=recipient_emails,
subject=email_params.subject,
html_body=email_html_content,
)
else:
# Old path: single recipient resolved from user_id/user_email
email_params = await self._get_email_params(
email_event=EmailEvent.max_budget_alert,
user_id=event.user_id,
user_email=event.user_email,
event_message=event.event_message,
)
email_html_content = MAX_BUDGET_ALERT_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
recipient_email=email_params.recipient_email,
percentage=percentage,
spend=spend_str,
max_budget=max_budget_str,
alert_threshold=alert_threshold_str,
base_url=email_params.base_url,
email_support_contact=email_params.support_contact,
)
await self.send_email(
from_email=self.DEFAULT_LITELLM_EMAIL,
to_email=[email_params.recipient_email],
subject=email_params.subject,
html_body=email_html_content,
)
async def budget_alerts(
self,
@@ -469,6 +518,13 @@ class BaseEmailLogger(CustomLogger):
# For max_budget_alert, check if we've already sent an alert
if type == "max_budget_alert":
if user_info.max_budget is not None and user_info.spend is not None:
if user_info.max_budget_alert_emails:
# New path: multi-threshold alerts
await self._handle_multi_threshold_max_budget_alert(
user_info=user_info, _cache=_cache
)
return
alert_threshold = (
user_info.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
)
@@ -527,6 +583,77 @@ class BaseEmailLogger(CustomLogger):
)
return
async def _handle_multi_threshold_max_budget_alert(
self,
user_info: CallInfo,
_cache: DualCache,
):
"""
Loop over configured thresholds in max_budget_alert_emails,
check cache per threshold, and send to configured recipients.
"""
for threshold_str, raw_emails in user_info.max_budget_alert_emails.items():
try:
threshold_pct = int(threshold_str)
except (ValueError, TypeError):
continue
threshold_amount = user_info.max_budget * (threshold_pct / 100.0)
if user_info.spend < threshold_amount:
continue
_id = user_info.token or user_info.user_id or "default_id"
_cache_key = (
f"email_budget_alerts:max_budget_alert:{threshold_pct}:{_id}"
)
result = await _cache.async_get_cache(key=_cache_key)
if result is not None:
continue
# Parse emails + auto-include owner
emails = _parse_email_list(raw_emails)
if user_info.user_email:
emails.append(user_info.user_email)
recipient_emails = list(set(emails)) if emails else None
event_message = f"Max Budget Alert - {threshold_pct}% of Maximum Budget Reached"
webhook_event = WebhookEvent(
event="max_budget_alert",
event_message=event_message,
spend=user_info.spend,
max_budget=user_info.max_budget,
soft_budget=user_info.soft_budget,
token=user_info.token,
customer_id=user_info.customer_id,
user_id=user_info.user_id,
team_id=user_info.team_id,
team_alias=user_info.team_alias,
organization_id=user_info.organization_id,
user_email=user_info.user_email,
key_alias=user_info.key_alias,
projected_exceeded_date=user_info.projected_exceeded_date,
projected_spend=user_info.projected_spend,
event_group=user_info.event_group,
)
try:
await self.send_max_budget_alert_email(
webhook_event,
threshold_pct=threshold_pct,
recipient_emails=recipient_emails,
)
await _cache.async_set_cache(
key=_cache_key,
value="SENT",
ttl=EMAIL_BUDGET_ALERT_TTL,
)
except Exception as e:
verbose_proxy_logger.error(
f"Error sending multi-threshold max budget alert email for {threshold_pct}%: {e}",
exc_info=True,
)
async def _get_email_params(
self,
email_event: EmailEvent,
+4
View File
@@ -3095,6 +3095,10 @@ class CallInfo(LiteLLMPydanticObjectBase):
default=None,
description="Additional email addresses to send alerts to (e.g., from team metadata)",
)
max_budget_alert_emails: Optional[Dict[str, List[str]]] = Field(
default=None,
description="Map of threshold percentage to email recipients (e.g., {'50': ['a@co.com'], '75': ['a@co.com', 'b@co.com']})",
)
class WebhookEvent(CallInfo):
+45 -16
View File
@@ -2963,22 +2963,13 @@ async def _virtual_key_max_budget_alert_check(
and valid_token.spend is not None
and valid_token.spend > 0
):
alert_threshold = (
valid_token.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
owner_email = user_obj.user_email if user_obj else None
alert_email_config = (valid_token.metadata or {}).get(
"max_budget_alert_emails"
)
# Only alert if we've crossed the threshold but haven't exceeded max_budget yet
if (
valid_token.spend >= alert_threshold
and valid_token.spend < valid_token.max_budget
):
verbose_proxy_logger.debug(
"Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s",
valid_token.token,
valid_token.spend,
valid_token.max_budget,
alert_threshold,
)
if isinstance(alert_email_config, dict) and alert_email_config:
# New path: pass the map through, let the email handler decide what to fire
call_info = CallInfo(
token=valid_token.token,
spend=valid_token.spend,
@@ -2988,17 +2979,55 @@ async def _virtual_key_max_budget_alert_check(
team_id=valid_token.team_id,
team_alias=valid_token.team_alias,
organization_id=valid_token.org_id,
user_email=user_obj.user_email if user_obj else None,
user_email=owner_email,
key_alias=valid_token.key_alias,
event_group=Litellm_EntityType.KEY,
max_budget_alert_emails=alert_email_config,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="max_budget_alert",
user_info=call_info,
)
)
else:
# Old path: existing single 80% threshold — completely unchanged
alert_threshold = (
valid_token.max_budget
* EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
)
if (
valid_token.spend >= alert_threshold
and valid_token.spend < valid_token.max_budget
):
verbose_proxy_logger.debug(
"Reached Max Budget Alert Threshold for token %s, spend %s, max_budget %s, alert_threshold %s",
valid_token.token,
valid_token.spend,
valid_token.max_budget,
alert_threshold,
)
call_info = CallInfo(
token=valid_token.token,
spend=valid_token.spend,
max_budget=valid_token.max_budget,
soft_budget=valid_token.soft_budget,
user_id=valid_token.user_id,
team_id=valid_token.team_id,
team_alias=valid_token.team_alias,
organization_id=valid_token.org_id,
user_email=owner_email,
key_alias=valid_token.key_alias,
event_group=Litellm_EntityType.KEY,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="max_budget_alert",
user_info=call_info,
)
)
async def _check_team_member_budget(
@@ -878,4 +878,200 @@ async def test_budget_alerts_max_budget_alert_crossed(
cache_call_args = mock_cache.async_set_cache.call_args[1]
assert cache_call_args["key"] == "email_budget_alerts:max_budget_alert:test_user"
assert cache_call_args["value"] == "SENT"
assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL
assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL
@pytest.mark.asyncio
async def test_multi_threshold_sends_crossed_thresholds(
base_email_logger, mock_send_email
):
"""Test that multi-threshold path sends emails for all crossed thresholds"""
user_info = CallInfo(
token="hashed_key_1",
user_id="test_user",
user_email="owner@co.com",
spend=80.0,
max_budget=100.0,
event_group=Litellm_EntityType.KEY,
max_budget_alert_emails={
"50": ["finance@co.com"],
"75": ["finance@co.com", "bu_lead@co.com"],
"100": ["cto@co.com"],
},
)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
# spend=80 crosses 50% ($50) and 75% ($75), but not 100% ($100)
assert mock_send_email.call_count == 2
# Check cache keys include threshold percentage
cache_keys = [
c[1]["key"] for c in mock_cache.async_set_cache.call_args_list
]
assert "email_budget_alerts:max_budget_alert:50:hashed_key_1" in cache_keys
assert "email_budget_alerts:max_budget_alert:75:hashed_key_1" in cache_keys
@pytest.mark.asyncio
async def test_multi_threshold_dedup_cache_prevents_resend(
base_email_logger, mock_send_email
):
"""Test that cached thresholds are not re-sent"""
user_info = CallInfo(
token="hashed_key_1",
user_id="test_user",
user_email="owner@co.com",
spend=80.0,
max_budget=100.0,
event_group=Litellm_EntityType.KEY,
max_budget_alert_emails={
"50": ["finance@co.com"],
"75": ["finance@co.com"],
},
)
# Simulate 50% already sent (cached), 75% not yet sent
async def cache_get(key):
if "50:" in key:
return "SENT"
return None
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(side_effect=cache_get)
mock_cache.async_set_cache = mock.AsyncMock()
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
# Only 75% should fire
assert mock_send_email.call_count == 1
cache_key = mock_cache.async_set_cache.call_args[1]["key"]
assert "75:" in cache_key
@pytest.mark.asyncio
async def test_multi_threshold_owner_email_auto_included(
base_email_logger, mock_send_email
):
"""Test that the owner email is auto-appended and deduplicated"""
user_info = CallInfo(
token="hashed_key_1",
user_id="test_user",
user_email="owner@co.com",
spend=60.0,
max_budget=100.0,
event_group=Litellm_EntityType.KEY,
max_budget_alert_emails={
"50": ["finance@co.com", "owner@co.com"], # owner already in list
},
)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
mock_send_email.assert_called_once()
to_emails = mock_send_email.call_args[1]["to_email"]
# owner@co.com should appear exactly once (deduplicated)
assert sorted(to_emails) == ["finance@co.com", "owner@co.com"]
@pytest.mark.asyncio
async def test_multi_threshold_malformed_keys_skipped(
base_email_logger, mock_send_email
):
"""Test that non-numeric threshold keys are skipped"""
user_info = CallInfo(
token="hashed_key_1",
user_id="test_user",
user_email="owner@co.com",
spend=60.0,
max_budget=100.0,
event_group=Litellm_EntityType.KEY,
max_budget_alert_emails={
"fifty": ["finance@co.com"], # invalid
"50": ["finance@co.com"], # valid, crossed
},
)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
# Only the valid "50" threshold should fire
assert mock_send_email.call_count == 1
@pytest.mark.asyncio
async def test_multi_threshold_empty_emails_only_owner(
base_email_logger, mock_send_email
):
"""Test that empty email list for a threshold sends only to owner"""
user_info = CallInfo(
token="hashed_key_1",
user_id="test_user",
user_email="owner@co.com",
spend=60.0,
max_budget=100.0,
event_group=Litellm_EntityType.KEY,
max_budget_alert_emails={
"50": [], # empty list
},
)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
mock_send_email.assert_called_once()
to_emails = mock_send_email.call_args[1]["to_email"]
assert to_emails == ["owner@co.com"]
@pytest.mark.asyncio
async def test_no_map_preserves_old_single_threshold(
base_email_logger, mock_send_email
):
"""Test that without max_budget_alert_emails, the old 80% single-threshold path works"""
user_info = CallInfo(
user_id="test_user",
user_email="test@example.com",
spend=165.0,
max_budget=200.0,
event_group=Litellm_EntityType.USER,
)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value=None)
mock_cache.async_set_cache = mock.AsyncMock()
base_email_logger.internal_usage_cache = mock_cache
with mock.patch.dict(os.environ, {"PROXY_BASE_URL": "http://test.com"}):
await base_email_logger.budget_alerts(type="max_budget_alert", user_info=user_info)
mock_send_email.assert_called_once()
call_args = mock_send_email.call_args[1]
assert call_args["to_email"] == ["test@example.com"]
# Old path cache key has no threshold percentage
cache_key = mock_cache.async_set_cache.call_args[1]["key"]
assert cache_key == "email_budget_alerts:max_budget_alert:test_user"
@@ -1544,6 +1544,114 @@ async def test_virtual_key_max_budget_alert_check_scenarios(
), f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}"
@pytest.mark.asyncio
async def test_virtual_key_max_budget_alert_check_with_multi_threshold_map():
"""Test that max_budget_alert_emails map from metadata is attached to CallInfo on the new path"""
alert_triggered = False
captured_call_info = None
class MockProxyLogging:
async def budget_alerts(self, type, user_info):
nonlocal alert_triggered, captured_call_info
alert_triggered = True
captured_call_info = user_info
alert_config = {
"50": ["finance@co.com"],
"75": ["finance@co.com", "bu_lead@co.com"],
}
valid_token = UserAPIKeyAuth(
token="test-token",
spend=30.0,
max_budget=100.0,
user_id="test-user",
key_alias="test-key",
metadata={"max_budget_alert_emails": alert_config},
)
user_obj = LiteLLM_UserTable(
user_id="test-user",
user_email="owner@co.com",
max_budget=None,
)
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=MockProxyLogging(),
user_obj=user_obj,
)
await asyncio.sleep(0.1)
assert alert_triggered is True
assert captured_call_info is not None
assert captured_call_info.max_budget_alert_emails == alert_config
assert captured_call_info.user_email == "owner@co.com"
assert captured_call_info.event_group == Litellm_EntityType.KEY
@pytest.mark.asyncio
async def test_virtual_key_max_budget_alert_check_old_path_no_map():
"""Test that old single-threshold path is used when no max_budget_alert_emails in metadata"""
alert_triggered = False
captured_call_info = None
class MockProxyLogging:
async def budget_alerts(self, type, user_info):
nonlocal alert_triggered, captured_call_info
alert_triggered = True
captured_call_info = user_info
# spend=90 is above 80% of 100 → old path should fire
valid_token = UserAPIKeyAuth(
token="test-token",
spend=90.0,
max_budget=100.0,
user_id="test-user",
key_alias="test-key",
metadata={},
)
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=MockProxyLogging(),
user_obj=None,
)
await asyncio.sleep(0.1)
assert alert_triggered is True
assert captured_call_info is not None
assert captured_call_info.max_budget_alert_emails is None
@pytest.mark.asyncio
async def test_virtual_key_max_budget_alert_check_old_path_below_threshold_no_alert():
"""Test that old path does NOT fire when spend is below 80% and no map is set"""
alert_triggered = False
class MockProxyLogging:
async def budget_alerts(self, type, user_info):
nonlocal alert_triggered
alert_triggered = True
# spend=50 is below 80% of 100 → should NOT fire
valid_token = UserAPIKeyAuth(
token="test-token",
spend=50.0,
max_budget=100.0,
user_id="test-user",
key_alias="test-key",
metadata={},
)
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=MockProxyLogging(),
user_obj=None,
)
await asyncio.sleep(0.1)
assert alert_triggered is False
@pytest.mark.asyncio
async def test_get_fuzzy_user_object_case_insensitive_email():
"""Test that _get_fuzzy_user_object uses case-insensitive email lookup"""