tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-11 19:48:29 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #25989 from BerriAI/litellm_feat-multi_threshold_budget_alerts

Browse Source 
feat: configurable multi-threshold budget alerts for virtual keys
This commit is contained in:
ryan-crabbe-berri
2026-04-18 16:21:29 -07:00
committed by GitHub
parent ecff06df65 c8b7c1bafa
commit 67bf18dfe0
7 changed files with 674 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py
+168 -28
View File
@@ -3,6 +3,7 @@ Base class for sending emails to user after creating keys or invite links
"""
import html
import json
import os
from typing import List, Literal, Optional
@@ -47,6 +48,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 +322,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 +349,67 @@ 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:
# 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 = html.escape(
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 +521,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 +586,87 @@ 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.
"""
if not user_info.max_budget_alert_emails or user_info.max_budget is None:
return
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)
if not emails:
verbose_proxy_logger.warning(
"No recipients for %d%% threshold on key %s, skipping alert",
threshold_pct,
_id,
)
continue
recipient_emails = list(set(emails))
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,
litellm/__init__.py
+1
View File
@@ -384,6 +384,7 @@ datadog_params: Optional[Union[DatadogInitParams, Dict]] = None
aws_sqs_callback_params: Optional[Dict] = None
generic_logger_headers: Optional[Dict] = None
default_key_generate_params: Optional[Dict] = None
default_key_max_budget_alert_emails: Optional[Dict[str, list]] = None
upperbound_key_generate_params: Optional[LiteLLM_UpperboundKeyGenerateParams] = None
key_generation_settings: Optional["StandardKeyGenerationConfig"] = None
default_internal_user_params: Optional[Dict] = None
litellm/proxy/_types.py
+4
View File
@@ -3135,6 +3135,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):
litellm/proxy/auth/auth_checks.py
+102 -15
View File
@@ -3059,6 +3059,52 @@ async def _virtual_key_soft_budget_check(
)
def _parse_email_list(raw: Any) -> 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 []
def _normalize_alert_emails(
cfg: Optional[Dict[str, Any]],
) -> Dict[str, List[str]]:
"""Coerce user-supplied threshold→recipients mapping to Dict[str, List[str]].
Values may legitimately arrive as list, comma-separated string, or None
from YAML/metadata; _parse_email_list tolerates all three.
"""
if not cfg:
return {}
return {k: _parse_email_list(v) for k, v in cfg.items()}
def _merge_budget_alert_email_configs(
global_cfg: Optional[Dict[str, Any]],
per_key_cfg: Optional[Dict[str, Any]],
) -> Optional[Dict[str, List[str]]]:
"""
Per-threshold additive merge: each threshold's recipient list is the union
of global + per-key entries (deduped, global-first ordering). Missing
thresholds on one side are inherited from the other.
"""
global_cfg_normalized = _normalize_alert_emails(global_cfg)
per_key_cfg_normalized = _normalize_alert_emails(per_key_cfg)
if not global_cfg_normalized and not per_key_cfg_normalized:
return None
thresholds = set(global_cfg_normalized) | set(per_key_cfg_normalized)
return {
t: list(
dict.fromkeys(
global_cfg_normalized.get(t, []) + per_key_cfg_normalized.get(t, [])
)
)
for t in thresholds
}
async def _virtual_key_max_budget_alert_check(
valid_token: UserAPIKeyAuth,
proxy_logging_obj: ProxyLogging,
@@ -3076,22 +3122,25 @@ 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: Optional[Dict[str, List[str]]] = (
_merge_budget_alert_email_configs(
global_cfg=litellm.default_key_max_budget_alert_emails,
per_key_cfg=(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: only create task if spend has crossed the lowest threshold
min_pct = min(
(int(k) for k in alert_email_config if k.isdigit()),
default=None,
)
if min_pct is None or valid_token.spend < valid_token.max_budget * (min_pct / 100.0):
return
call_info = CallInfo(
token=valid_token.token,
spend=valid_token.spend,
@@ -3101,17 +3150,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(
litellm/proxy/auth/user_api_key_auth.py
+10 -8
View File
@@ -1368,7 +1368,16 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
if not skip_budget_checks:
with tracer.trace("litellm.proxy.auth.budget_checks"):
# Check 4. Token Spend is under budget
# Check 4. Max Budget Alert Check (runs before budget enforcement
# so multi-threshold 100% alerts fire on the request that crosses
# max_budget, before BudgetExceededError is raised below)
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 5. Token Spend is under budget
if RouteChecks.is_llm_api_route(route=route):
await _virtual_key_max_budget_check(
valid_token=valid_token,
@@ -1376,13 +1385,6 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
user_obj=user_obj,
)
# Check 5. Max Budget Alert Check
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 6. Soft Budget Check
await _virtual_key_soft_budget_check(
valid_token=valid_token,
tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py
+197 -1
View File
@@ -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"
tests/test_litellm/proxy/auth/test_auth_checks.py
+192
View File
@@ -1558,6 +1558,198 @@ 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=60.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_virtual_key_max_budget_alert_check_global_fallback():
"""Test that litellm.default_key_max_budget_alert_emails is used when key metadata has no map"""
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
global_config = {
"50": ["global-finance@co.com"],
"75": ["global-finance@co.com", "global-lead@co.com"],
}
valid_token = UserAPIKeyAuth(
token="test-token",
spend=60.0,
max_budget=100.0,
user_id="test-user",
key_alias="test-key",
metadata={}, # no per-key config
)
import litellm
original = litellm.default_key_max_budget_alert_emails
try:
litellm.default_key_max_budget_alert_emails = global_config
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.max_budget_alert_emails == global_config
finally:
litellm.default_key_max_budget_alert_emails = original
@pytest.mark.asyncio
async def test_virtual_key_max_budget_alert_check_per_key_merges_with_global():
"""Test that per-key and global configs are additively merged"""
captured_call_info = None
class MockProxyLogging:
async def budget_alerts(self, type, user_info):
nonlocal captured_call_info
captured_call_info = user_info
per_key_config = {"50": ["per-key@co.com"]}
global_config = {"75": ["global@co.com"]}
valid_token = UserAPIKeyAuth(
token="test-token",
spend=60.0,
max_budget=100.0,
user_id="test-user",
key_alias="test-key",
metadata={"max_budget_alert_emails": per_key_config},
)
import litellm
original = litellm.default_key_max_budget_alert_emails
try:
litellm.default_key_max_budget_alert_emails = global_config
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=MockProxyLogging(),
user_obj=None,
)
await asyncio.sleep(0.1)
# Additive merge: both thresholds present, recipients merged per threshold
assert captured_call_info.max_budget_alert_emails == {
"50": ["per-key@co.com"],
"75": ["global@co.com"],
}
finally:
litellm.default_key_max_budget_alert_emails = original
@pytest.mark.asyncio
async def test_get_fuzzy_user_object_case_insensitive_email():
"""Test that _get_fuzzy_user_object uses case-insensitive email lookup"""