mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-08 22:25:25 +00:00
Merge pull request #20553 from BerriAI/litellm_team_soft_budget_email
[Feature] Team Soft Budget Email Alerts
This commit is contained in:
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -30,8 +30,15 @@ from litellm.integrations.email_templates.user_invitation_email import (
|
||||
from litellm.integrations.email_templates.templates import (
|
||||
MAX_BUDGET_ALERT_EMAIL_TEMPLATE,
|
||||
SOFT_BUDGET_ALERT_EMAIL_TEMPLATE,
|
||||
TEAM_SOFT_BUDGET_ALERT_EMAIL_TEMPLATE,
|
||||
)
|
||||
from litellm.proxy._types import (
|
||||
CallInfo,
|
||||
InvitationNew,
|
||||
Litellm_EntityType,
|
||||
UserAPIKeyAuth,
|
||||
WebhookEvent,
|
||||
)
|
||||
from litellm.proxy._types import CallInfo, InvitationNew, UserAPIKeyAuth, WebhookEvent
|
||||
from litellm.secret_managers.main import get_secret_bool
|
||||
from litellm.types.integrations.slack_alerting import LITELLM_LOGO_URL
|
||||
from litellm.constants import (
|
||||
@@ -217,6 +224,78 @@ class BaseEmailLogger(CustomLogger):
|
||||
)
|
||||
pass
|
||||
|
||||
async def send_team_soft_budget_alert_email(self, event: WebhookEvent):
|
||||
"""
|
||||
Send email to team members when team soft budget is crossed
|
||||
Supports multiple recipients via alert_emails field from team metadata
|
||||
"""
|
||||
# Collect all recipient emails
|
||||
recipient_emails: List[str] = []
|
||||
|
||||
# Add additional alert emails from team metadata.soft_budget_alert_emails
|
||||
if hasattr(event, "alert_emails") and event.alert_emails:
|
||||
for email in event.alert_emails:
|
||||
if email and email not in recipient_emails: # Avoid duplicates
|
||||
recipient_emails.append(email)
|
||||
|
||||
# If no recipients found, skip sending
|
||||
if not recipient_emails:
|
||||
verbose_proxy_logger.warning(
|
||||
f"No recipient emails found for team soft budget alert. event={event.model_dump(exclude_none=True)}"
|
||||
)
|
||||
return
|
||||
|
||||
# Validate that we have at least one valid email address
|
||||
first_recipient_email = recipient_emails[0]
|
||||
if not first_recipient_email or not first_recipient_email.strip():
|
||||
verbose_proxy_logger.warning(
|
||||
f"Invalid recipient email found for team soft budget alert. event={event.model_dump(exclude_none=True)}"
|
||||
)
|
||||
return
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"send_team_soft_budget_alert_email_event: {json.dumps(event.model_dump(exclude_none=True), indent=4, default=str)}"
|
||||
)
|
||||
|
||||
# Get email params using the first recipient email (for template formatting)
|
||||
# For team alerts with alert_emails, we don't need user_id lookup since we already have email addresses
|
||||
# Pass user_id=None to prevent _get_email_params from trying to look up email from a potentially None user_id
|
||||
email_params = await self._get_email_params(
|
||||
email_event=EmailEvent.soft_budget_crossed,
|
||||
user_id=None, # Team alerts don't require user_id when alert_emails are provided
|
||||
user_email=first_recipient_email,
|
||||
event_message=event.event_message,
|
||||
)
|
||||
|
||||
# Format budget values
|
||||
soft_budget_str = f"${event.soft_budget}" if event.soft_budget is not None else "N/A"
|
||||
spend_str = f"${event.spend}" if event.spend is not None else "$0.00"
|
||||
max_budget_info = ""
|
||||
if event.max_budget is not None:
|
||||
max_budget_info = f"<b>Maximum Budget:</b> ${event.max_budget} <br />"
|
||||
|
||||
# Use team alias or generic greeting
|
||||
team_alias = event.team_alias or "Team"
|
||||
|
||||
email_html_content = TEAM_SOFT_BUDGET_ALERT_EMAIL_TEMPLATE.format(
|
||||
email_logo_url=email_params.logo_url,
|
||||
team_alias=team_alias,
|
||||
soft_budget=soft_budget_str,
|
||||
spend=spend_str,
|
||||
max_budget_info=max_budget_info,
|
||||
base_url=email_params.base_url,
|
||||
email_support_contact=email_params.support_contact,
|
||||
)
|
||||
|
||||
# Send email to all recipients
|
||||
await self.send_email(
|
||||
from_email=self.DEFAULT_LITELLM_EMAIL,
|
||||
to_email=recipient_emails,
|
||||
subject=email_params.subject,
|
||||
html_body=email_html_content,
|
||||
)
|
||||
pass
|
||||
|
||||
async def send_max_budget_alert_email(self, event: WebhookEvent):
|
||||
"""
|
||||
Send email to user when max budget alert threshold is reached
|
||||
@@ -285,15 +364,36 @@ class BaseEmailLogger(CustomLogger):
|
||||
# - Don't re-alert, if alert already sent
|
||||
_cache: DualCache = self.internal_usage_cache
|
||||
|
||||
# percent of max_budget left to spend
|
||||
if user_info.max_budget is None and user_info.soft_budget is None:
|
||||
return
|
||||
|
||||
# For soft_budget alerts, check if we've already sent an alert
|
||||
if type == "soft_budget":
|
||||
# For team soft budget alerts, we only need team soft_budget to be set
|
||||
# For other entity types, we need either max_budget or soft_budget
|
||||
if user_info.event_group == Litellm_EntityType.TEAM:
|
||||
if user_info.soft_budget is None:
|
||||
return
|
||||
# For team soft budget alerts, require alert_emails to be configured
|
||||
# Team soft budget alerts are sent via metadata.soft_budget_alerting_emails
|
||||
if user_info.alert_emails is None or len(user_info.alert_emails) == 0:
|
||||
verbose_proxy_logger.debug(
|
||||
"Skipping team soft budget email alert: no alert_emails configured",
|
||||
)
|
||||
return
|
||||
else:
|
||||
# For non-team alerts, require either max_budget or soft_budget
|
||||
if user_info.max_budget is None and user_info.soft_budget is None:
|
||||
return
|
||||
if user_info.soft_budget is not None and user_info.spend >= user_info.soft_budget:
|
||||
# Generate cache key based on event type and identifier
|
||||
_id = user_info.token or user_info.user_id or "default_id"
|
||||
# Use appropriate ID based on event_group to ensure unique cache keys per entity type
|
||||
if user_info.event_group == Litellm_EntityType.TEAM:
|
||||
_id = user_info.team_id or "default_id"
|
||||
elif user_info.event_group == Litellm_EntityType.ORGANIZATION:
|
||||
_id = user_info.organization_id or "default_id"
|
||||
elif user_info.event_group == Litellm_EntityType.USER:
|
||||
_id = user_info.user_id or "default_id"
|
||||
else:
|
||||
# For KEY and other types, use token or user_id
|
||||
_id = user_info.token or user_info.user_id or "default_id"
|
||||
_cache_key = f"email_budget_alerts:soft_budget_crossed:{_id}"
|
||||
|
||||
# Check if we've already sent this alert
|
||||
@@ -318,10 +418,15 @@ class BaseEmailLogger(CustomLogger):
|
||||
projected_exceeded_date=user_info.projected_exceeded_date,
|
||||
projected_spend=user_info.projected_spend,
|
||||
event_group=user_info.event_group,
|
||||
alert_emails=user_info.alert_emails,
|
||||
)
|
||||
|
||||
try:
|
||||
await self.send_soft_budget_alert_email(webhook_event)
|
||||
# Use team-specific function for team alerts, otherwise use standard function
|
||||
if user_info.event_group == Litellm_EntityType.TEAM:
|
||||
await self.send_team_soft_budget_alert_email(webhook_event)
|
||||
else:
|
||||
await self.send_soft_budget_alert_email(webhook_event)
|
||||
|
||||
# Cache the alert to prevent duplicate sends
|
||||
await _cache.async_set_cache(
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "litellm-enterprise"
|
||||
version = "0.1.29"
|
||||
version = "0.1.31"
|
||||
description = "Package for LiteLLM Enterprise features"
|
||||
authors = ["BerriAI"]
|
||||
readme = "README.md"
|
||||
@@ -22,7 +22,7 @@ requires = ["poetry-core"]
|
||||
build-backend = "poetry.core.masonry.api"
|
||||
|
||||
[tool.commitizen]
|
||||
version = "0.1.29"
|
||||
version = "0.1.31"
|
||||
version_files = [
|
||||
"pyproject.toml:version",
|
||||
"../requirements.txt:litellm-enterprise==",
|
||||
|
||||
@@ -85,6 +85,30 @@ SOFT_BUDGET_ALERT_EMAIL_TEMPLATE = """
|
||||
The LiteLLM team <br />
|
||||
"""
|
||||
|
||||
TEAM_SOFT_BUDGET_ALERT_EMAIL_TEMPLATE = """
|
||||
<img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" />
|
||||
|
||||
<p> Hi {team_alias} team member, <br/>
|
||||
|
||||
Your LiteLLM team has crossed its <b>soft budget limit of {soft_budget}</b>. <br /> <br />
|
||||
|
||||
<b>Current Spend:</b> {spend} <br />
|
||||
<b>Soft Budget:</b> {soft_budget} <br />
|
||||
{max_budget_info}
|
||||
|
||||
<p style="color: #dc2626; font-weight: 500;">
|
||||
⚠️ Note: Your API requests will continue to work, but you should monitor your usage closely.
|
||||
If you reach your maximum budget, requests will be rejected.
|
||||
</p>
|
||||
|
||||
You can view your usage and manage your budget in the <a href="{base_url}">LiteLLM Dashboard</a>. <br /> <br />
|
||||
|
||||
If you have any questions, please send an email to {email_support_contact} <br /> <br />
|
||||
|
||||
Best, <br />
|
||||
The LiteLLM team <br />
|
||||
"""
|
||||
|
||||
MAX_BUDGET_ALERT_EMAIL_TEMPLATE = """
|
||||
<img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" />
|
||||
|
||||
|
||||
@@ -2190,6 +2190,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken):
|
||||
team_tpm_limit: Optional[int] = None
|
||||
team_rpm_limit: Optional[int] = None
|
||||
team_max_budget: Optional[float] = None
|
||||
team_soft_budget: Optional[float] = None
|
||||
team_models: List = []
|
||||
team_blocked: bool = False
|
||||
soft_budget: Optional[float] = None
|
||||
@@ -2648,6 +2649,10 @@ class CallInfo(LiteLLMPydanticObjectBase):
|
||||
projected_exceeded_date: Optional[str] = None
|
||||
projected_spend: Optional[float] = None
|
||||
event_group: Litellm_EntityType
|
||||
alert_emails: Optional[List[str]] = Field(
|
||||
default=None,
|
||||
description="Additional email addresses to send alerts to (e.g., from team metadata)",
|
||||
)
|
||||
|
||||
|
||||
class WebhookEvent(CallInfo):
|
||||
|
||||
@@ -240,6 +240,13 @@ async def common_checks(
|
||||
valid_token=valid_token,
|
||||
)
|
||||
|
||||
# 3.0.5. If team is over soft budget (alert only, doesn't block)
|
||||
await _team_soft_budget_check(
|
||||
team_object=team_object,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
valid_token=valid_token,
|
||||
)
|
||||
|
||||
# 3.1. If organization is in budget
|
||||
await _organization_max_budget_check(
|
||||
valid_token=valid_token,
|
||||
@@ -2444,6 +2451,75 @@ async def _team_max_budget_check(
|
||||
)
|
||||
|
||||
|
||||
async def _team_soft_budget_check(
|
||||
team_object: Optional[LiteLLM_TeamTable],
|
||||
valid_token: Optional[UserAPIKeyAuth],
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
):
|
||||
"""
|
||||
Triggers a budget alert if the team is over it's soft budget.
|
||||
"""
|
||||
if (
|
||||
team_object is not None
|
||||
and team_object.soft_budget is not None
|
||||
and team_object.spend is not None
|
||||
and team_object.spend >= team_object.soft_budget
|
||||
):
|
||||
verbose_proxy_logger.debug(
|
||||
"Crossed Soft Budget for team %s, spend %s, soft_budget %s",
|
||||
team_object.team_id,
|
||||
team_object.spend,
|
||||
team_object.soft_budget,
|
||||
)
|
||||
if valid_token:
|
||||
# Extract alert emails from team metadata
|
||||
alert_emails: Optional[List[str]] = None
|
||||
if team_object.metadata is not None and isinstance(team_object.metadata, dict):
|
||||
soft_budget_alert_emails = team_object.metadata.get("soft_budget_alerting_emails")
|
||||
if soft_budget_alert_emails is not None:
|
||||
if isinstance(soft_budget_alert_emails, list):
|
||||
alert_emails = [email for email in soft_budget_alert_emails if isinstance(email, str) and email.strip()]
|
||||
elif isinstance(soft_budget_alert_emails, str):
|
||||
# Handle comma-separated string
|
||||
alert_emails = [email.strip() for email in soft_budget_alert_emails.split(",") if email.strip()]
|
||||
# Filter out empty strings
|
||||
if alert_emails:
|
||||
alert_emails = [email for email in alert_emails if email]
|
||||
else:
|
||||
alert_emails = None
|
||||
|
||||
# Only send team soft budget alerts if alert_emails are configured
|
||||
# Team soft budget alerts are sent via metadata.soft_budget_alerting_emails, not global alerting
|
||||
if alert_emails is None or len(alert_emails) == 0:
|
||||
verbose_proxy_logger.debug(
|
||||
"Skipping team soft budget alert for team %s: no alert_emails configured in metadata.soft_budget_alerting_emails",
|
||||
team_object.team_id,
|
||||
)
|
||||
return
|
||||
|
||||
call_info = CallInfo(
|
||||
token=valid_token.token,
|
||||
spend=team_object.spend,
|
||||
max_budget=team_object.max_budget,
|
||||
soft_budget=team_object.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=None, # Team-level alert, no specific user email
|
||||
key_alias=valid_token.key_alias,
|
||||
event_group=Litellm_EntityType.TEAM,
|
||||
alert_emails=alert_emails,
|
||||
)
|
||||
|
||||
asyncio.create_task(
|
||||
proxy_logging_obj.budget_alerts(
|
||||
type="soft_budget",
|
||||
user_info=call_info,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
async def _organization_max_budget_check(
|
||||
valid_token: Optional[UserAPIKeyAuth],
|
||||
team_object: Optional[LiteLLM_TeamTable],
|
||||
|
||||
@@ -1149,6 +1149,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
|
||||
_team_obj: Optional[LiteLLM_TeamTable] = LiteLLM_TeamTable(
|
||||
team_id=valid_token.team_id,
|
||||
max_budget=valid_token.team_max_budget,
|
||||
soft_budget=valid_token.team_soft_budget,
|
||||
spend=valid_token.team_spend,
|
||||
tpm_limit=valid_token.team_tpm_limit,
|
||||
rpm_limit=valid_token.team_rpm_limit,
|
||||
|
||||
+29
-9
@@ -1370,17 +1370,36 @@ class ProxyLogging:
|
||||
],
|
||||
user_info: CallInfo,
|
||||
):
|
||||
if self.alerting is None:
|
||||
# do nothing if alerting is not switched on
|
||||
# For soft_budget alerts with alert_emails set, allow email sending even if alerting is None
|
||||
# This enables team-specific soft budget email alerts via metadata.soft_budget_alerting_emails
|
||||
# Note: user_info is a CallInfo that can represent user/team/org level info. For team budgets,
|
||||
# alert_emails is populated from team_object.metadata.soft_budget_alerting_emails (see auth_checks.py)
|
||||
is_soft_budget_with_alert_emails = (
|
||||
type == "soft_budget"
|
||||
and user_info.alert_emails is not None
|
||||
and len(user_info.alert_emails) > 0
|
||||
)
|
||||
|
||||
if self.alerting is None and not is_soft_budget_with_alert_emails:
|
||||
# do nothing if alerting is not switched on (unless it's a soft_budget alert with team-specific emails)
|
||||
return
|
||||
|
||||
if "slack" in self.alerting:
|
||||
await self.slack_alerting_instance.budget_alerts(
|
||||
type=type,
|
||||
user_info=user_info,
|
||||
)
|
||||
if self.alerting is not None and "slack" in self.alerting:
|
||||
if self.slack_alerting_instance is not None:
|
||||
await self.slack_alerting_instance.budget_alerts(
|
||||
type=type,
|
||||
user_info=user_info,
|
||||
)
|
||||
|
||||
if "email" in self.alerting and self.email_logging_instance is not None:
|
||||
# Call email_logging_instance if:
|
||||
# 1. "email" is in alerting config, OR
|
||||
# 2. It's a soft_budget alert with team-specific alert_emails (bypasses global alerting config)
|
||||
should_send_email = (
|
||||
(self.alerting is not None and "email" in self.alerting)
|
||||
or is_soft_budget_with_alert_emails
|
||||
)
|
||||
|
||||
if should_send_email and self.email_logging_instance is not None:
|
||||
await self.email_logging_instance.budget_alerts(
|
||||
type=type,
|
||||
user_info=user_info,
|
||||
@@ -2607,7 +2626,8 @@ class PrismaClient:
|
||||
SELECT
|
||||
v.*,
|
||||
t.spend AS team_spend,
|
||||
t.max_budget AS team_max_budget,
|
||||
t.max_budget AS team_max_budget,
|
||||
t.soft_budget AS team_soft_budget,
|
||||
t.tpm_limit AS team_tpm_limit,
|
||||
t.rpm_limit AS team_rpm_limit,
|
||||
t.models AS team_models,
|
||||
|
||||
+1
-1
@@ -63,7 +63,7 @@ mcp = {version = ">=1.25.0,<2.0.0", optional = true, python = ">=3.10"}
|
||||
a2a-sdk = {version = "^0.3.22", optional = true, python = ">=3.10"}
|
||||
litellm-proxy-extras = {version = "0.4.31", optional = true}
|
||||
rich = {version = "13.7.1", optional = true}
|
||||
litellm-enterprise = {version = "0.1.27", optional = true}
|
||||
litellm-enterprise = {version = "0.1.31", optional = true}
|
||||
diskcache = {version = "^5.6.1", optional = true}
|
||||
polars = {version = "^1.31.0", optional = true, python = ">=3.10"}
|
||||
semantic-router = {version = ">=0.1.12", optional = true, python = ">=3.9,<3.14"}
|
||||
|
||||
+1
-1
@@ -73,4 +73,4 @@ pypdf>=6.6.2 # for PDF text extraction in RAG ingestion
|
||||
########################
|
||||
# LITELLM ENTERPRISE DEPENDENCIES
|
||||
########################
|
||||
litellm-enterprise==0.1.29
|
||||
litellm-enterprise==0.1.31
|
||||
|
||||
@@ -21,11 +21,13 @@ from litellm.proxy._types import (
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_UserTable,
|
||||
LiteLLM_TeamTable,
|
||||
Litellm_EntityType,
|
||||
)
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
can_team_access_model,
|
||||
_virtual_key_soft_budget_check,
|
||||
_team_soft_budget_check,
|
||||
)
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.proxy.utils import CallInfo
|
||||
@@ -478,6 +480,84 @@ async def test_virtual_key_soft_budget_check(spend, soft_budget, expect_alert):
|
||||
), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"spend, soft_budget, expect_alert, metadata, expected_alert_emails",
|
||||
[
|
||||
(100, 50, False, None, None), # Over soft budget, no metadata - no alert_emails configured, so no alert
|
||||
(50, 50, False, None, None), # At soft budget, no metadata - no alert_emails configured, so no alert
|
||||
(25, 50, False, None, None), # Under soft budget
|
||||
(100, None, False, None, None), # No soft budget set
|
||||
(100, 50, True, {"soft_budget_alerting_emails": ["team1@example.com", "team2@example.com"]}, ["team1@example.com", "team2@example.com"]), # Over soft budget with list of emails
|
||||
(100, 50, True, {"soft_budget_alerting_emails": "team1@example.com,team2@example.com"}, ["team1@example.com", "team2@example.com"]), # Over soft budget with comma-separated emails
|
||||
(100, 50, True, {"soft_budget_alerting_emails": ["team1@example.com", "", " ", "team2@example.com"]}, ["team1@example.com", "team2@example.com"]), # Over soft budget with empty strings filtered
|
||||
],
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_soft_budget_check(spend, soft_budget, expect_alert, metadata, expected_alert_emails):
|
||||
"""
|
||||
Test cases for _team_soft_budget_check:
|
||||
1. Spend over soft budget, no alert_emails configured - should NOT trigger alert (alerts only sent when alert_emails configured)
|
||||
2. Spend at soft budget, no alert_emails configured - should NOT trigger alert (alerts only sent when alert_emails configured)
|
||||
3. Spend under soft budget - should not trigger alert
|
||||
4. No soft budget set - should not trigger alert
|
||||
5. Team with alert emails in metadata (list) - should include alert_emails in CallInfo
|
||||
6. Team with alert emails in metadata (comma-separated string) - should parse and include alert_emails
|
||||
7. Team with alert emails containing empty strings - should filter them out
|
||||
"""
|
||||
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
|
||||
assert type == "soft_budget"
|
||||
assert isinstance(user_info, CallInfo)
|
||||
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="test-token",
|
||||
user_id="test-user",
|
||||
team_id="test-team",
|
||||
team_alias="test-team-alias",
|
||||
key_alias="test-key",
|
||||
)
|
||||
|
||||
team_object = LiteLLM_TeamTable(
|
||||
team_id="test-team",
|
||||
spend=spend,
|
||||
soft_budget=soft_budget,
|
||||
max_budget=100.0,
|
||||
metadata=metadata,
|
||||
)
|
||||
|
||||
proxy_logging_obj = MockProxyLogging()
|
||||
|
||||
await _team_soft_budget_check(
|
||||
team_object=team_object,
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
await asyncio.sleep(0.1) # Allow time for the alert task to complete
|
||||
|
||||
assert (
|
||||
alert_triggered == expect_alert
|
||||
), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
|
||||
|
||||
if expect_alert:
|
||||
assert captured_call_info is not None
|
||||
assert captured_call_info.team_id == "test-team"
|
||||
assert captured_call_info.spend == spend
|
||||
assert captured_call_info.soft_budget == soft_budget
|
||||
assert captured_call_info.event_group == Litellm_EntityType.TEAM
|
||||
# Verify alert_emails if expected
|
||||
if expected_alert_emails is not None:
|
||||
assert captured_call_info.alert_emails == expected_alert_emails
|
||||
else:
|
||||
assert captured_call_info.alert_emails is None or captured_call_info.alert_emails == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_can_user_call_model():
|
||||
from litellm.proxy.auth.auth_checks import can_user_call_model
|
||||
|
||||
@@ -2870,6 +2870,111 @@ class TestProxyLoggingBudgetAlerts:
|
||||
type=alert_type, user_info=user_info
|
||||
)
|
||||
|
||||
async def test_budget_alerts_soft_budget_with_alert_emails_bypasses_alerting_none(self):
|
||||
"""
|
||||
Test that soft_budget alerts with alert_emails bypass the alerting=None check
|
||||
and send emails even when alerting is None.
|
||||
|
||||
This tests the new logic that allows team-specific soft budget email alerts
|
||||
via metadata.soft_budget_alerting_emails to work even when global alerting is disabled.
|
||||
"""
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.proxy._types import CallInfo, Litellm_EntityType
|
||||
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
proxy_logging.alerting = None # Global alerting is disabled
|
||||
proxy_logging.slack_alerting_instance = AsyncMock()
|
||||
proxy_logging.email_logging_instance = AsyncMock()
|
||||
|
||||
# Create CallInfo with alert_emails set (simulating team metadata extraction)
|
||||
user_info = CallInfo(
|
||||
token="test-token",
|
||||
spend=100.0,
|
||||
soft_budget=50.0,
|
||||
user_id="test-user",
|
||||
team_id="test-team",
|
||||
team_alias="test-team-alias",
|
||||
event_group=Litellm_EntityType.TEAM,
|
||||
alert_emails=["team1@example.com", "team2@example.com"],
|
||||
)
|
||||
|
||||
# Should send email even though alerting is None (because of alert_emails)
|
||||
await proxy_logging.budget_alerts(type="soft_budget", user_info=user_info)
|
||||
|
||||
# Verify slack was NOT called (alerting is None)
|
||||
proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called()
|
||||
|
||||
# Verify email WAS called (bypasses alerting=None check)
|
||||
proxy_logging.email_logging_instance.budget_alerts.assert_called_once_with(
|
||||
type="soft_budget", user_info=user_info
|
||||
)
|
||||
|
||||
async def test_budget_alerts_soft_budget_without_alert_emails_respects_alerting_none(self):
|
||||
"""
|
||||
Test that soft_budget alerts WITHOUT alert_emails still respect alerting=None
|
||||
and do not send emails when alerting is None.
|
||||
"""
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.proxy._types import CallInfo, Litellm_EntityType
|
||||
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
proxy_logging.alerting = None
|
||||
proxy_logging.slack_alerting_instance = AsyncMock()
|
||||
proxy_logging.email_logging_instance = AsyncMock()
|
||||
|
||||
# Create CallInfo WITHOUT alert_emails
|
||||
user_info = CallInfo(
|
||||
token="test-token",
|
||||
spend=100.0,
|
||||
soft_budget=50.0,
|
||||
user_id="test-user",
|
||||
team_id="test-team",
|
||||
team_alias="test-team-alias",
|
||||
event_group=Litellm_EntityType.TEAM,
|
||||
alert_emails=None, # No alert emails
|
||||
)
|
||||
|
||||
# Should NOT send email (alerting is None and no alert_emails)
|
||||
await proxy_logging.budget_alerts(type="soft_budget", user_info=user_info)
|
||||
|
||||
# Verify no calls were made
|
||||
proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called()
|
||||
proxy_logging.email_logging_instance.budget_alerts.assert_not_called()
|
||||
|
||||
async def test_budget_alerts_soft_budget_with_empty_alert_emails_respects_alerting_none(self):
|
||||
"""
|
||||
Test that soft_budget alerts with empty alert_emails list still respect alerting=None.
|
||||
"""
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.proxy._types import CallInfo, Litellm_EntityType
|
||||
|
||||
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
|
||||
proxy_logging.alerting = None
|
||||
proxy_logging.slack_alerting_instance = AsyncMock()
|
||||
proxy_logging.email_logging_instance = AsyncMock()
|
||||
|
||||
# Create CallInfo with empty alert_emails list
|
||||
user_info = CallInfo(
|
||||
token="test-token",
|
||||
spend=100.0,
|
||||
soft_budget=50.0,
|
||||
user_id="test-user",
|
||||
team_id="test-team",
|
||||
team_alias="test-team-alias",
|
||||
event_group=Litellm_EntityType.TEAM,
|
||||
alert_emails=[], # Empty list
|
||||
)
|
||||
|
||||
# Should NOT send email (alert_emails is empty)
|
||||
await proxy_logging.budget_alerts(type="soft_budget", user_info=user_info)
|
||||
|
||||
# Verify no calls were made
|
||||
proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called()
|
||||
proxy_logging.email_logging_instance.budget_alerts.assert_not_called()
|
||||
|
||||
|
||||
def test_azure_ai_claude_provider_config():
|
||||
"""Test that Azure AI Claude models return AzureAnthropicConfig for proper tool transformation."""
|
||||
|
||||
Reference in New Issue
Block a user