scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 114

Merge pull request #17995 from BerriAI/litellm_email_budget_alerts_2

Browse source 
[Feature] Email Budget Alerts
This commit is contained in:
yuneng-jiang 2025-12-17 19:16:17 -08:00 • committed by GitHub
commit 866d1eded5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 1164 additions and 49 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

BIN
enterprise/dist/litellm_enterprise-0.1.26-py3-none-any.whl vendored Normal file
View file

Binary file not shown.

BIN
enterprise/dist/litellm_enterprise-0.1.26.tar.gz vendored Normal file
View file

Binary file not shown.

BIN
enterprise/dist/litellm_enterprise-0.1.27-py3-none-any.whl vendored Normal file
View file

Binary file not shown.

BIN
enterprise/dist/litellm_enterprise-0.1.27.tar.gz vendored Normal file
View file

Binary file not shown.

240
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/base_email.py
View file

@ -5,7 +5,7 @@ Base class for sending emails to user after creating keys or invite links
import json
import os
from typing import List, Optional
from typing import List, Literal, Optional
from litellm_enterprise.types.enterprise_callbacks.send_emails import (
EmailEvent,
@ -15,6 +15,7 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import (
)
from litellm._logging import verbose_proxy_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER
from litellm.integrations.email_templates.key_created_email import (
@ -26,9 +27,17 @@ from litellm.integrations.email_templates.key_rotated_email import (
from litellm.integrations.email_templates.user_invitation_email import (
USER_INVITATION_EMAIL_TEMPLATE,
)
from litellm.proxy._types import InvitationNew, UserAPIKeyAuth, WebhookEvent
from litellm.integrations.email_templates.templates import (
MAX_BUDGET_ALERT_EMAIL_TEMPLATE,
SOFT_BUDGET_ALERT_EMAIL_TEMPLATE,
)
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 (
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
EMAIL_BUDGET_ALERT_TTL,
)
class BaseEmailLogger(CustomLogger):
@ -40,6 +49,21 @@ class BaseEmailLogger(CustomLogger):
EmailEvent.virtual_key_rotated: "LiteLLM: {event_message}",
}
def __init__(
self,
internal_usage_cache: Optional[DualCache] = None,
**kwargs,
):
"""
Initialize BaseEmailLogger
Args:
internal_usage_cache: DualCache instance for preventing duplicate alerts
**kwargs: Additional arguments passed to CustomLogger
"""
super().__init__(**kwargs)
self.internal_usage_cache = internal_usage_cache or DualCache()
async def send_user_invitation_email(self, event: WebhookEvent):
"""
Send email to user after inviting them to the team
@ -154,6 +178,218 @@ class BaseEmailLogger(CustomLogger):
)
pass
async def send_soft_budget_alert_email(self, event: WebhookEvent):
"""
Send email to user when soft budget is crossed
"""
email_params = await self._get_email_params(
email_event=EmailEvent.soft_budget_crossed, # Reuse existing event type for subject template
user_id=event.user_id,
user_email=event.user_email,
event_message=event.event_message,
)
verbose_proxy_logger.debug(
f"send_soft_budget_alert_email_event: {json.dumps(event.model_dump(exclude_none=True), indent=4, default=str)}"
)
# 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 />"
email_html_content = SOFT_BUDGET_ALERT_EMAIL_TEMPLATE.format(
email_logo_url=email_params.logo_url,
recipient_email=email_params.recipient_email,
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,
)
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
async def send_max_budget_alert_email(self, event: WebhookEvent):
"""
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,
)
verbose_proxy_logger.debug(
f"send_max_budget_alert_email_event: {json.dumps(event.model_dump(exclude_none=True), indent=4, default=str)}"
)
# Format budget values
spend_str = f"${event.spend}" if event.spend is not None else "$0.00"
max_budget_str = f"${event.max_budget}" if event.max_budget is not None else "N/A"
# Calculate percentage and alert threshold
percentage = int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100)
alert_threshold_str = f"${event.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE:.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
async def budget_alerts(
self,
type: Literal[
"token_budget",
"soft_budget",
"max_budget_alert",
"user_budget",
"team_budget",
"organization_budget",
"proxy_budget",
"projected_limit_exceeded",
],
user_info: CallInfo,
):
"""
Send a budget alert via email
Args:
type: The type of budget alert to send
user_info: The user info to send the alert for
"""
## PREVENTITIVE ALERTING ##
# - Alert once within 24hr period
# - Cache this information
# - 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":
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"
_cache_key = f"email_budget_alerts:soft_budget_crossed:{_id}"
# Check if we've already sent this alert
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
# Create WebhookEvent for soft budget alert
event_message = f"Soft Budget Crossed - Total Soft Budget: ${user_info.soft_budget}"
webhook_event = WebhookEvent(
event="soft_budget_crossed",
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_soft_budget_alert_email(webhook_event)
# Cache the alert to prevent duplicate sends
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 soft budget alert email: {e}",
exc_info=True,
)
return
# 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:
alert_threshold = user_info.max_budget * EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
# Only alert if we've crossed the threshold but haven't exceeded max_budget yet
if user_info.spend >= alert_threshold and user_info.spend < user_info.max_budget:
# Generate cache key based on event type and identifier
_id = user_info.token or user_info.user_id or "default_id"
_cache_key = f"email_budget_alerts:max_budget_alert:{_id}"
# Check if we've already sent this alert
result = await _cache.async_get_cache(key=_cache_key)
if result is None:
# Calculate percentage
percentage = int(EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE * 100)
# Create WebhookEvent for max budget alert
event_message = f"Max Budget Alert - {percentage}% 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)
# Cache the alert to prevent duplicate sends
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 max budget alert email: {e}",
exc_info=True,
)
return
async def _get_email_params(
self,
email_event: EmailEvent,

3
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/resend_email.py
View file

@ -19,7 +19,8 @@ RESEND_API_ENDPOINT = "https://api.resend.com/emails"
class ResendEmailLogger(BaseEmailLogger):
def __init__(self):
def __init__(self, internal_usage_cache=None, **kwargs):
super().__init__(internal_usage_cache=internal_usage_cache, **kwargs)
self.async_httpx_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)

3
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/sendgrid_email.py
View file

@ -27,7 +27,8 @@ class SendGridEmailLogger(BaseEmailLogger):
- SENDGRID_API_KEY
"""
def __init__(self):
def __init__(self, internal_usage_cache=None, **kwargs):
super().__init__(internal_usage_cache=internal_usage_cache, **kwargs)
self.async_httpx_client = get_async_httpx_client(
llm_provider=httpxSpecialProvider.LoggingCallback
)

3
enterprise/litellm_enterprise/enterprise_callbacks/send_emails/smtp_email.py
View file

@ -21,7 +21,8 @@ class SMTPEmailLogger(BaseEmailLogger):
- SMTP_SENDER_EMAIL
"""
def __init__(self):
def __init__(self, internal_usage_cache=None, **kwargs):
super().__init__(internal_usage_cache=internal_usage_cache, **kwargs)
verbose_logger.debug("SMTP Email Logger initialized....")
async def send_email(

4
enterprise/litellm_enterprise/types/enterprise_callbacks/send_emails.py
View file

@ -36,6 +36,8 @@ class EmailEvent(str, enum.Enum):
virtual_key_created = "Virtual Key Created"
new_user_invitation = "New User Invitation"
virtual_key_rotated = "Virtual Key Rotated"
soft_budget_crossed = "Soft Budget Crossed"
max_budget_alert = "Max Budget Alert"
class EmailEventSettings(BaseModel):
event: EmailEvent
@ -51,6 +53,8 @@ class DefaultEmailSettings(BaseModel):
EmailEvent.virtual_key_created: True, # On by default
EmailEvent.new_user_invitation: True, # On by default
EmailEvent.virtual_key_rotated: True, # On by default
EmailEvent.soft_budget_crossed: True, # On by default
EmailEvent.max_budget_alert: True, # On by default
}
)
def to_dict(self) -> Dict[str, bool]:

4
enterprise/pyproject.toml
View file

@ -1,6 +1,6 @@
[tool.poetry]
name = "litellm-enterprise"
version = "0.1.25"
version = "0.1.27"
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.25"
version = "0.1.27"
version_files = [
"pyproject.toml:version",
"../requirements.txt:litellm-enterprise==",

2
litellm/constants.py
View file

@ -313,6 +313,8 @@ DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE = os.getenv(
"DD_TRACER_STREAMING_CHUNK_YIELD_RESOURCE", "streaming.chunk.yield"
)
EMAIL_BUDGET_ALERT_TTL = int(os.getenv("EMAIL_BUDGET_ALERT_TTL", 24 * 60 * 60)) # 24 hours in seconds
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE = float(os.getenv("EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE", 0.8)) # 80% of max budget
############### LLM Provider Constants ###############
### ANTHROPIC CONSTANTS ###
ANTHROPIC_SKILLS_API_BETA_VERSION = "skills-2025-10-02"

48
litellm/integrations/email_templates/templates.py
View file

@ -60,3 +60,51 @@ USER_INVITED_EMAIL_TEMPLATE = """
Best, <br />
The LiteLLM team <br />
"""
SOFT_BUDGET_ALERT_EMAIL_TEMPLATE = """
<img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" />
<p> Hi {recipient_email}, <br/>
Your LiteLLM API key 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" />
<p> Hi {recipient_email}, <br/>
Your LiteLLM API key has reached <b>{percentage}% of its maximum budget</b>. <br /> <br />
<b>Current Spend:</b> {spend} <br />
<b>Maximum Budget:</b> {max_budget} <br />
<b>Alert Threshold:</b> {alert_threshold} ({percentage}%) <br />
<p style="color: #dc2626; font-weight: 500;">
⚠️ Warning: You are approaching your maximum budget limit.
Once you reach your maximum budget of {max_budget}, all API 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 />
"""

1
litellm/proxy/_types.py
View file

@ -2484,6 +2484,7 @@ class CallInfo(LiteLLMPydanticObjectBase):
class WebhookEvent(CallInfo):
event: Literal[
"budget_crossed",
"max_budget_alert",
"soft_budget_crossed",
"threshold_crossed",
"projected_limit_exceeded",

56
litellm/proxy/auth/auth_checks.py
View file

@ -24,6 +24,7 @@ from litellm.constants import (
DEFAULT_IN_MEMORY_TTL,
DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL,
DEFAULT_MAX_RECURSE_DEPTH,
EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE,
)
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.proxy._types import (
@ -1911,6 +1912,7 @@ async def _virtual_key_max_budget_check(
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,
organization_id=valid_token.org_id,
@ -1939,6 +1941,7 @@ async def _virtual_key_max_budget_check(
async def _virtual_key_soft_budget_check(
valid_token: UserAPIKeyAuth,
proxy_logging_obj: ProxyLogging,
user_obj: Optional[LiteLLM_UserTable] = None,
):
"""
Triggers a budget alert if the token is over it's soft budget.
@ -1961,10 +1964,11 @@ async def _virtual_key_soft_budget_check(
team_id=valid_token.team_id,
team_alias=valid_token.team_alias,
organization_id=valid_token.org_id,
user_email=None,
user_email=user_obj.user_email if user_obj else None,
key_alias=valid_token.key_alias,
event_group=Litellm_EntityType.KEY,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="soft_budget",
@ -1973,6 +1977,56 @@ async def _virtual_key_soft_budget_check(
)
async def _virtual_key_max_budget_alert_check(
valid_token: UserAPIKeyAuth,
proxy_logging_obj: ProxyLogging,
user_obj: Optional[LiteLLM_UserTable] = None,
):
"""
Triggers a budget alert if the token has reached EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE
(default 80%) of its max budget.
This is a warning alert before the token actually exceeds the max budget.
"""
if (
valid_token.max_budget is not None
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
# 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,
)
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=user_obj.user_email if user_obj else None,
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 _team_max_budget_check(
team_object: Optional[LiteLLM_TeamTable],
valid_token: Optional[UserAPIKeyAuth],

11
litellm/proxy/auth/user_api_key_auth.py
View file

@ -29,6 +29,7 @@ from litellm.proxy.auth.auth_checks import (
_get_user_role,
_is_user_proxy_admin,
_virtual_key_max_budget_check,
_virtual_key_max_budget_alert_check,
_virtual_key_soft_budget_check,
can_key_call_model,
common_checks,
@ -1062,10 +1063,18 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
user_obj=user_obj,
)
# Check 5. Soft Budget Check
# 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,
proxy_logging_obj=proxy_logging_obj,
user_obj=user_obj,
)
# Check 5. Token Model Spend is under Model budget

7
litellm/proxy/proxy_config.yaml
View file

@ -3,7 +3,8 @@ model_list:
litellm_params:
model: gemini/*
litellm_settings:
callbacks: ["langfuse"]
callbacks: ["dynamic_rate_limiter_v3"]
priority_reservation:
"prod": 0.9 # 90% reserved for production
"dev": 0.1 # 10% reserved for development

70
litellm/proxy/proxy_server.py
View file

@ -2795,19 +2795,25 @@ class ProxyConfig:
verbose_proxy_logger.debug(f"_alerting_callbacks: {general_settings}")
if _alerting_callbacks is None:
return
# Ensure proxy_logging_obj.alerting is set for all alerting types
_alerting_value = general_settings.get("alerting", None)
verbose_proxy_logger.debug(f"_load_alerting_settings: Calling update_values with alerting={_alerting_value}")
proxy_logging_obj.update_values(
alerting=_alerting_value,
alerting_threshold=general_settings.get("alerting_threshold", 600),
alert_types=general_settings.get("alert_types", None),
alert_to_webhook_url=general_settings.get(
"alert_to_webhook_url", None
),
alerting_args=general_settings.get("alerting_args", None),
redis_cache=redis_usage_cache,
)
for _alert in _alerting_callbacks:
if _alert == "slack":
# [OLD] v0 implementation
proxy_logging_obj.update_values(
alerting=general_settings.get("alerting", None),
alerting_threshold=general_settings.get("alerting_threshold", 600),
alert_types=general_settings.get("alert_types", None),
alert_to_webhook_url=general_settings.get(
"alert_to_webhook_url", None
),
alerting_args=general_settings.get("alerting_args", None),
redis_cache=redis_usage_cache,
)
# [OLD] v0 implementation - already handled by update_values above
pass
else:
# [NEW] v1 implementation - init as a custom logger
if _alert in litellm._known_custom_logger_compatible_callbacks:
@ -3274,6 +3280,7 @@ class ProxyConfig:
proxy_logging_obj: ProxyLogging
"""
_general_settings = config_data.get("general_settings", {})
if _general_settings is not None and "alerting" in _general_settings:
if (
general_settings is not None
@ -3282,29 +3289,36 @@ class ProxyConfig:
and _general_settings.get("alerting", None) is not None
and isinstance(_general_settings["alerting"], list)
):
verbose_proxy_logger.debug(
"Overriding Default 'alerting' values with db 'alerting' values."
)
general_settings["alerting"] = _general_settings[
"alerting"
] # override yaml values with db
proxy_logging_obj.alerting = general_settings["alerting"]
proxy_logging_obj.slack_alerting_instance.alerting = general_settings[
"alerting"
# Merge DB and YAML/config alerting values instead of overriding
_yaml_alerting = set(general_settings["alerting"])
_db_alerting = set(_general_settings["alerting"])
_merged_alerting = list(_yaml_alerting.union(_db_alerting))
# Preserve order: YAML values first, then DB values
_merged_alerting = list(general_settings["alerting"]) + [
item for item in _general_settings["alerting"]
if item not in general_settings["alerting"]
]
verbose_proxy_logger.debug(
f"Merging alerting values: YAML={general_settings['alerting']}, DB={_general_settings['alerting']}, Merged={_merged_alerting}"
)
general_settings["alerting"] = _merged_alerting
# Use update_values to properly set alerting for both slack and email
proxy_logging_obj.update_values(
alerting=general_settings["alerting"],
)
elif general_settings is None:
general_settings = {}
general_settings["alerting"] = _general_settings["alerting"]
proxy_logging_obj.alerting = general_settings["alerting"]
proxy_logging_obj.slack_alerting_instance.alerting = general_settings[
"alerting"
]
# Use update_values to properly set alerting for both slack and email
proxy_logging_obj.update_values(
alerting=general_settings["alerting"],
)
elif isinstance(general_settings, dict):
general_settings["alerting"] = _general_settings["alerting"]
proxy_logging_obj.alerting = general_settings["alerting"]
proxy_logging_obj.slack_alerting_instance.alerting = general_settings[
"alerting"
]
# Use update_values to properly set alerting for both slack and email
proxy_logging_obj.update_values(
alerting=general_settings["alerting"],
)
if _general_settings is not None and "alert_types" in _general_settings:
general_settings["alert_types"] = _general_settings["alert_types"]

63
litellm/proxy/utils.py
View file

@ -35,6 +35,17 @@ from litellm.proxy._types import (
from litellm.types.guardrails import GuardrailEventHooks
from litellm.types.utils import CallTypes, CallTypesLiteral
try:
from litellm_enterprise.enterprise_callbacks.send_emails.base_email import BaseEmailLogger
from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import SendGridEmailLogger
from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import SMTPEmailLogger
from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ResendEmailLogger
except ImportError:
BaseEmailLogger = None # type: ignore
SendGridEmailLogger = None # type: ignore
SMTPEmailLogger = None # type: ignore
ResendEmailLogger = None # type: ignore
try:
import backoff
except ImportError:
@ -128,6 +139,33 @@ def print_verbose(print_statement):
print(f"LiteLLM Proxy: {print_statement}") # noqa
def _get_email_logger_class():
"""
Determine which email logger class to use based on environment variables.
Priority: SendGrid > Resend > SMTP > BaseEmailLogger (fallback)
Returns:
The email logger class to use, or None if BaseEmailLogger is not available
"""
if BaseEmailLogger is None:
return None
# Check for SendGrid API key
if SendGridEmailLogger is not None and os.getenv("SENDGRID_API_KEY"):
return SendGridEmailLogger
# Check for Resend API key
if ResendEmailLogger is not None and os.getenv("RESEND_API_KEY"):
return ResendEmailLogger
# Check for SMTP configuration
if SMTPEmailLogger is not None and os.getenv("SMTP_HOST"):
return SMTPEmailLogger
# Fallback to BaseEmailLogger (though it won't actually send emails)
return BaseEmailLogger
class InternalUsageCache:
def __init__(self, dual_cache: DualCache):
self.dual_cache: DualCache = dual_cache
@ -266,6 +304,14 @@ class ProxyLogging:
alerting=self.alerting,
internal_usage_cache=self.internal_usage_cache.dual_cache,
)
self.email_logging_instance: Optional[Any] = None
if BaseEmailLogger is not None:
email_logger_class = _get_email_logger_class()
if email_logger_class is not None:
# All email logger classes now accept internal_usage_cache
self.email_logging_instance = email_logger_class(
internal_usage_cache=self.internal_usage_cache.dual_cache,
)
self.premium_user = premium_user
self.service_logging_obj = ServiceLogging()
self.db_spend_update_writer = DBSpendUpdateWriter()
@ -1149,6 +1195,7 @@ class ProxyLogging:
"token_budget",
"user_budget",
"soft_budget",
"max_budget_alert",
"team_budget",
"organization_budget",
"proxy_budget",
@ -1159,10 +1206,18 @@ class ProxyLogging:
if self.alerting is None:
# do nothing if alerting is not switched on
return
await self.slack_alerting_instance.budget_alerts(
type=type,
user_info=user_info,
)
if "slack" in self.alerting:
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:
await self.email_logging_instance.budget_alerts(
type=type,
user_info=user_info,
)
async def alerting_handler(
self,

2
pyproject.toml
View file

@ -61,7 +61,7 @@ redisvl = {version = "^0.4.1", optional = true, markers = "python_version >= '3.
mcp = {version = "^1.21.2", optional = true, python = ">=3.10"}
litellm-proxy-extras = {version = "0.4.14", optional = true}
rich = {version = "13.7.1", optional = true}
litellm-enterprise = {version = "0.1.25", optional = true}
litellm-enterprise = {version = "0.1.27", 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"}

2
requirements.txt
View file

@ -67,4 +67,4 @@ soundfile==0.12.1 # for audio file processing
########################
# LITELLM ENTERPRISE DEPENDENCIES
########################
litellm-enterprise==0.1.25
litellm-enterprise==0.1.27

277
tests/test_litellm/enterprise/enterprise_callbacks/send_emails/test_base_email.py
View file

@ -19,7 +19,8 @@ from litellm_enterprise.types.enterprise_callbacks.send_emails import (
)
from litellm.integrations.email_templates.email_footer import EMAIL_FOOTER
from litellm.proxy._types import Litellm_EntityType, WebhookEvent
from litellm.proxy._types import CallInfo, Litellm_EntityType, WebhookEvent
from litellm.constants import EMAIL_BUDGET_ALERT_TTL
@pytest.fixture(autouse=True)
@ -605,4 +606,276 @@ async def test_get_email_params_default_templates(monkeypatch):
)
assert key_params.subject == "LiteLLM: API Key Created"
assert key_params.signature == EMAIL_FOOTER
assert key_params.signature == EMAIL_FOOTER
@pytest.mark.asyncio
async def test_send_soft_budget_alert_email(
base_email_logger, mock_send_email, mock_lookup_user_email
):
"""Test that send_soft_budget_alert_email sends an email with the correct parameters and content"""
event = WebhookEvent(
user_id="test_user",
user_email="test@example.com",
event_group=Litellm_EntityType.USER,
event="soft_budget_crossed",
event_message="Soft Budget Crossed - Total Soft Budget: $100.0",
spend=105.0,
max_budget=200.0,
soft_budget=100.0,
)
with mock.patch.dict(
os.environ,
{
"EMAIL_LOGO_URL": "https://litellm-listing.s3.amazonaws.com/litellm_logo.png",
"EMAIL_SUPPORT_CONTACT": "support@berri.ai",
"PROXY_BASE_URL": "http://test.com",
},
):
await base_email_logger.send_soft_budget_alert_email(event)
mock_send_email.assert_called_once()
call_args = mock_send_email.call_args[1]
assert call_args["from_email"] == BaseEmailLogger.DEFAULT_LITELLM_EMAIL
assert call_args["to_email"] == ["test@example.com"]
assert call_args["subject"] == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0"
assert "$100.0" in call_args["html_body"] # soft_budget
assert "$105.0" in call_args["html_body"] # spend
assert "$200.0" in call_args["html_body"] # max_budget
@pytest.mark.asyncio
async def test_send_soft_budget_alert_email_no_max_budget(
base_email_logger, mock_send_email, mock_lookup_user_email
):
"""Test that send_soft_budget_alert_email handles missing max_budget correctly"""
event = WebhookEvent(
user_id="test_user",
user_email="test@example.com",
event_group=Litellm_EntityType.USER,
event="soft_budget_crossed",
event_message="Soft Budget Crossed - Total Soft Budget: $100.0",
spend=105.0,
max_budget=None,
soft_budget=100.0,
)
with mock.patch.dict(
os.environ,
{
"PROXY_BASE_URL": "http://test.com",
},
):
await base_email_logger.send_soft_budget_alert_email(event)
mock_send_email.assert_called_once()
call_args = mock_send_email.call_args[1]
assert "$100.0" in call_args["html_body"] # soft_budget
assert "$105.0" in call_args["html_body"] # spend
assert "Maximum Budget" not in call_args["html_body"] # max_budget should not be shown
@pytest.mark.asyncio
async def test_budget_alerts_soft_budget_crossed(
base_email_logger, mock_send_email
):
"""Test that budget_alerts sends email when soft budget is crossed"""
user_info = CallInfo(
user_id="test_user",
user_email="test@example.com",
spend=105.0,
max_budget=200.0,
soft_budget=100.0,
event_group=Litellm_EntityType.USER,
)
# Mock the cache to return None (no previous alert sent)
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="soft_budget", user_info=user_info)
# Verify email was sent
mock_send_email.assert_called_once()
call_args = mock_send_email.call_args[1]
assert call_args["to_email"] == ["test@example.com"]
# Verify cache was set to prevent duplicate alerts
mock_cache.async_set_cache.assert_called_once()
cache_call_args = mock_cache.async_set_cache.call_args[1]
assert cache_call_args["key"] == "email_budget_alerts:soft_budget_crossed:test_user"
assert cache_call_args["value"] == "SENT"
assert cache_call_args["ttl"] == EMAIL_BUDGET_ALERT_TTL
@pytest.mark.asyncio
async def test_budget_alerts_soft_budget_not_crossed(
base_email_logger, mock_send_email
):
"""Test that budget_alerts does not send email when soft budget is not crossed"""
user_info = CallInfo(
user_id="test_user",
user_email="test@example.com",
spend=50.0,
max_budget=200.0,
soft_budget=100.0,
event_group=Litellm_EntityType.USER,
)
mock_cache = mock.AsyncMock()
base_email_logger.internal_usage_cache = mock_cache
await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info)
# Verify email was NOT sent
mock_send_email.assert_not_called()
@pytest.mark.asyncio
async def test_budget_alerts_soft_budget_duplicate_prevention(
base_email_logger, mock_send_email
):
"""Test that budget_alerts does not send duplicate alerts within TTL period"""
user_info = CallInfo(
user_id="test_user",
user_email="test@example.com",
spend=105.0,
max_budget=200.0,
soft_budget=100.0,
event_group=Litellm_EntityType.USER,
)
# Mock the cache to return "SENT" (previous alert already sent)
mock_cache = mock.AsyncMock()
mock_cache.async_get_cache = mock.AsyncMock(return_value="SENT")
base_email_logger.internal_usage_cache = mock_cache
await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info)
# Verify email was NOT sent (duplicate prevention)
mock_send_email.assert_not_called()
@pytest.mark.asyncio
async def test_budget_alerts_no_budgets(
base_email_logger, mock_send_email
):
"""Test that budget_alerts returns early when no budgets are set"""
user_info = CallInfo(
user_id="test_user",
user_email="test@example.com",
spend=50.0,
max_budget=None,
soft_budget=None,
event_group=Litellm_EntityType.USER,
)
await base_email_logger.budget_alerts(type="soft_budget", user_info=user_info)
# Verify email was NOT sent
mock_send_email.assert_not_called()
@pytest.mark.asyncio
async def test_budget_alerts_uses_token_for_cache_key(
base_email_logger, mock_send_email
):
"""Test that budget_alerts uses token for cache key when available"""
user_info = CallInfo(
user_id="test_user",
user_email="test@example.com",
token="hashed_token_123",
spend=105.0,
max_budget=200.0,
soft_budget=100.0,
event_group=Litellm_EntityType.KEY,
)
# Mock the cache to return None (no previous alert sent)
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="soft_budget", user_info=user_info)
# Verify cache key uses token instead of user_id
mock_cache.async_set_cache.assert_called_once()
cache_call_args = mock_cache.async_set_cache.call_args[1]
assert cache_call_args["key"] == "email_budget_alerts:soft_budget_crossed:hashed_token_123"
@pytest.mark.asyncio
async def test_get_email_params_soft_budget_crossed(
base_email_logger, mock_lookup_user_email
):
"""Test that _get_email_params handles soft_budget_crossed event correctly"""
with mock.patch.dict(
os.environ,
{
"PROXY_BASE_URL": "http://test.com",
},
):
result = await base_email_logger._get_email_params(
email_event=EmailEvent.soft_budget_crossed,
user_email="test@example.com",
event_message="Soft Budget Crossed - Total Soft Budget: $100.0",
)
# Should use default subject template for soft_budget_crossed
assert result.subject == "LiteLLM: Soft Budget Crossed - Total Soft Budget: $100.0"
assert result.recipient_email == "test@example.com"
assert result.base_url == "http://test.com"
@pytest.mark.asyncio
async def test_budget_alerts_max_budget_alert_crossed(
base_email_logger, mock_send_email
):
"""Test that budget_alerts sends email when max budget alert threshold is crossed"""
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"]
assert "Max Budget Alert" in call_args["subject"]
mock_cache.async_set_cache.assert_called_once()
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

289
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -14,9 +14,11 @@ import pytest
import litellm
from litellm.proxy._types import (
CallInfo,
LiteLLM_ObjectPermissionTable,
LiteLLM_TeamTable,
LiteLLM_UserTable,
Litellm_EntityType,
LitellmUserRoles,
ProxyErrorTypes,
ProxyException,
@ -27,6 +29,8 @@ from litellm.proxy.auth.auth_checks import (
ExperimentalUIJWTToken,
_can_object_call_vector_stores,
_get_team_db_check,
_virtual_key_max_budget_alert_check,
_virtual_key_soft_budget_check,
get_user_object,
vector_store_access_check,
)
@ -988,3 +992,288 @@ async def test_reject_clientside_metadata_tags_non_llm_route():
)
assert result is True
@pytest.mark.asyncio
async def test_virtual_key_soft_budget_check_with_user_obj():
"""Test _virtual_key_soft_budget_check includes user_email when user_obj is provided"""
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",
spend=100.0,
soft_budget=50.0,
user_id="test-user",
team_id="test-team",
team_alias="test-team-alias",
org_id="test-org",
key_alias="test-key",
max_budget=200.0,
)
user_obj = LiteLLM_UserTable(
user_id="test-user",
user_email="test@example.com",
max_budget=None,
)
proxy_logging_obj = MockProxyLogging()
await _virtual_key_soft_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
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.user_email == "test@example.com"
assert captured_call_info.token == "test-token"
assert captured_call_info.spend == 100.0
assert captured_call_info.soft_budget == 50.0
assert captured_call_info.max_budget == 200.0
assert captured_call_info.user_id == "test-user"
assert captured_call_info.team_id == "test-team"
assert captured_call_info.team_alias == "test-team-alias"
assert captured_call_info.organization_id == "test-org"
assert captured_call_info.key_alias == "test-key"
assert captured_call_info.event_group == Litellm_EntityType.KEY
@pytest.mark.asyncio
async def test_virtual_key_soft_budget_check_without_user_obj():
"""Test _virtual_key_soft_budget_check sets user_email to None when user_obj is not provided"""
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",
spend=100.0,
soft_budget=50.0,
user_id="test-user",
team_id="test-team",
key_alias="test-key",
)
proxy_logging_obj = MockProxyLogging()
await _virtual_key_soft_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=None,
)
await asyncio.sleep(0.1)
assert alert_triggered is True
assert captured_call_info is not None
assert captured_call_info.user_email is None
@pytest.mark.parametrize(
"spend, soft_budget, expect_alert",
[
(100.0, 50.0, True), # Over soft budget
(50.0, 50.0, True), # At soft budget
(25.0, 50.0, False), # Under soft budget
(100.0, None, False), # No soft budget set
],
)
@pytest.mark.asyncio
async def test_virtual_key_soft_budget_check_scenarios(
spend, soft_budget, expect_alert
):
"""Test _virtual_key_soft_budget_check with various spend and soft_budget scenarios"""
alert_triggered = False
class MockProxyLogging:
async def budget_alerts(self, type, user_info):
nonlocal alert_triggered
alert_triggered = True
assert type == "soft_budget"
assert isinstance(user_info, CallInfo)
valid_token = UserAPIKeyAuth(
token="test-token",
spend=spend,
soft_budget=soft_budget,
user_id="test-user",
key_alias="test-key",
)
proxy_logging_obj = MockProxyLogging()
await _virtual_key_soft_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=None,
)
await asyncio.sleep(0.1)
assert (
alert_triggered == expect_alert
), f"Expected alert_triggered to be {expect_alert} for spend={spend}, soft_budget={soft_budget}"
@pytest.mark.asyncio
async def test_virtual_key_max_budget_alert_check_with_user_obj():
"""Test _virtual_key_max_budget_alert_check includes user_email when user_obj is provided"""
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 == "max_budget_alert"
assert isinstance(user_info, CallInfo)
valid_token = UserAPIKeyAuth(
token="test-token",
spend=90.0,
max_budget=100.0,
user_id="test-user",
team_id="test-team",
team_alias="test-team-alias",
org_id="test-org",
key_alias="test-key",
soft_budget=50.0,
)
user_obj = LiteLLM_UserTable(
user_id="test-user",
user_email="test@example.com",
max_budget=None,
)
proxy_logging_obj = MockProxyLogging()
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
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.user_email == "test@example.com"
assert captured_call_info.token == "test-token"
assert captured_call_info.spend == 90.0
assert captured_call_info.max_budget == 100.0
assert captured_call_info.soft_budget == 50.0
assert captured_call_info.user_id == "test-user"
assert captured_call_info.team_id == "test-team"
assert captured_call_info.team_alias == "test-team-alias"
assert captured_call_info.organization_id == "test-org"
assert captured_call_info.key_alias == "test-key"
assert captured_call_info.event_group == Litellm_EntityType.KEY
@pytest.mark.asyncio
async def test_virtual_key_max_budget_alert_check_without_user_obj():
"""Test _virtual_key_max_budget_alert_check sets user_email to None when user_obj is not provided"""
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 == "max_budget_alert"
assert isinstance(user_info, CallInfo)
valid_token = UserAPIKeyAuth(
token="test-token",
spend=90.0,
max_budget=100.0,
user_id="test-user",
team_id="test-team",
key_alias="test-key",
)
proxy_logging_obj = MockProxyLogging()
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=None,
)
await asyncio.sleep(0.1)
assert alert_triggered is True
assert captured_call_info is not None
assert captured_call_info.user_email is None
@pytest.mark.parametrize(
"spend, max_budget, expect_alert",
[
(80.0, 100.0, True), # At 80% threshold (alert threshold)
(90.0, 100.0, True), # Above threshold, below max_budget
(79.0, 100.0, False), # Below threshold
(100.0, 100.0, False), # At max_budget (not below, so no alert)
(110.0, 100.0, False), # Above max_budget (already exceeded)
(100.0, None, False), # No max_budget set
(0.0, 100.0, False), # Spend is 0
],
)
@pytest.mark.asyncio
async def test_virtual_key_max_budget_alert_check_scenarios(
spend, max_budget, expect_alert
):
"""Test _virtual_key_max_budget_alert_check with various spend and max_budget scenarios"""
alert_triggered = False
class MockProxyLogging:
async def budget_alerts(self, type, user_info):
nonlocal alert_triggered
alert_triggered = True
assert type == "max_budget_alert"
assert isinstance(user_info, CallInfo)
valid_token = UserAPIKeyAuth(
token="test-token",
spend=spend,
max_budget=max_budget,
user_id="test-user",
key_alias="test-key",
)
proxy_logging_obj = MockProxyLogging()
await _virtual_key_max_budget_alert_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
user_obj=None,
)
await asyncio.sleep(0.1)
assert (
alert_triggered == expect_alert
), f"Expected alert_triggered to be {expect_alert} for spend={spend}, max_budget={max_budget}"

128
tests/test_litellm/test_utils.py
View file

@ -1,7 +1,7 @@
import json
import os
import sys
from unittest.mock import MagicMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from jsonschema import validate
@ -2604,6 +2604,132 @@ class TestIsCachedMessage:
assert is_cached_message(message) is False
@pytest.mark.asyncio
class TestProxyLoggingBudgetAlerts:
"""Test budget_alerts method in ProxyLogging class."""
async def test_budget_alerts_when_alerting_is_none(self):
"""Test that budget_alerts returns early when alerting is None."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging.alerting = None
proxy_logging.slack_alerting_instance = AsyncMock()
proxy_logging.email_logging_instance = AsyncMock()
user_info = MagicMock()
# Should return without calling any alerting instances
await proxy_logging.budget_alerts(type="user_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_with_slack_only(self):
"""Test that budget_alerts calls slack_alerting_instance when slack is in alerting."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging.alerting = ["slack"]
proxy_logging.slack_alerting_instance = AsyncMock()
user_info = MagicMock()
await proxy_logging.budget_alerts(type="token_budget", user_info=user_info)
proxy_logging.slack_alerting_instance.budget_alerts.assert_called_once_with(
type="token_budget", user_info=user_info
)
async def test_budget_alerts_with_email_only(self):
"""Test that budget_alerts calls email_logging_instance when email is in alerting."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging.alerting = ["email"]
proxy_logging.email_logging_instance = AsyncMock()
user_info = MagicMock()
await proxy_logging.budget_alerts(type="team_budget", user_info=user_info)
proxy_logging.email_logging_instance.budget_alerts.assert_called_once_with(
type="team_budget", user_info=user_info
)
async def test_budget_alerts_with_email_when_instance_is_none(self):
"""Test that budget_alerts does not call email_logging_instance when it is None."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging.alerting = ["email"]
proxy_logging.email_logging_instance = None
user_info = MagicMock()
# Should not raise an error
await proxy_logging.budget_alerts(type="organization_budget", user_info=user_info)
async def test_budget_alerts_with_both_slack_and_email(self):
"""Test that budget_alerts calls both slack and email instances when both are in alerting."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging.alerting = ["slack", "email"]
proxy_logging.slack_alerting_instance = AsyncMock()
proxy_logging.email_logging_instance = AsyncMock()
user_info = MagicMock()
await proxy_logging.budget_alerts(type="proxy_budget", user_info=user_info)
proxy_logging.slack_alerting_instance.budget_alerts.assert_called_once_with(
type="proxy_budget", user_info=user_info
)
proxy_logging.email_logging_instance.budget_alerts.assert_called_once_with(
type="proxy_budget", user_info=user_info
)
@pytest.mark.parametrize(
"alert_type",
[
"token_budget",
"user_budget",
"soft_budget",
"team_budget",
"organization_budget",
"proxy_budget",
"projected_limit_exceeded",
],
)
async def test_budget_alerts_with_all_alert_types(self, alert_type):
"""Test that budget_alerts works with all supported alert types."""
from litellm.caching.caching import DualCache
from litellm.proxy.utils import ProxyLogging
proxy_logging = ProxyLogging(user_api_key_cache=DualCache())
proxy_logging.alerting = ["slack", "email"]
proxy_logging.slack_alerting_instance = AsyncMock()
proxy_logging.email_logging_instance = AsyncMock()
user_info = MagicMock()
await proxy_logging.budget_alerts(type=alert_type, user_info=user_info)
proxy_logging.slack_alerting_instance.budget_alerts.assert_called_once_with(
type=alert_type, user_info=user_info
)
proxy_logging.email_logging_instance.budget_alerts.assert_called_once_with(
type=alert_type, user_info=user_info
)
def test_azure_ai_claude_provider_config():
"""Test that Azure AI Claude models return AzureAnthropicConfig for proper tool transformation."""
from litellm import AzureAnthropicConfig, AzureAIStudioConfig