Merge pull request #4209 from BerriAI/litellm_send_email_alerts_budget_exceeded

[Feat] send email alerts when budget exceeded

litellm/integrations/email_alerting.py

@@ -0,0 +1,136 @@
+"""
+Functions for sending Email Alerts
+"""
+
+import os
+from typing import Optional, List
+from litellm.proxy._types import WebhookEvent
+import asyncio
+from litellm._logging import verbose_logger, verbose_proxy_logger
+
+# we use this for the email header, please send a test email if you change this. verify it looks good on email
+LITELLM_LOGO_URL = "https://litellm-listing.s3.amazonaws.com/litellm_logo.png"
+LITELLM_SUPPORT_CONTACT = "support@berri.ai"
+
+
+async def get_all_team_member_emails(team_id: Optional[str] = None) -> list:
+    verbose_logger.debug(
+        "Email Alerting: Getting all team members for team_id=%s", team_id
+    )
+    if team_id is None:
+        return []
+    from litellm.proxy.proxy_server import premium_user, prisma_client
+
+    if prisma_client is None:
+        raise Exception("Not connected to DB!")
+
+    team_row = await prisma_client.db.litellm_teamtable.find_unique(
+        where={
+            "team_id": team_id,
+        }
+    )
+
+    if team_row is None:
+        return []
+
+    _team_members = team_row.members_with_roles
+    verbose_logger.debug(
+        "Email Alerting: Got team members for team_id=%s Team Members: %s",
+        team_id,
+        _team_members,
+    )
+    _team_member_user_ids: List[str] = []
+    for member in _team_members:
+        if member and isinstance(member, dict) and member.get("user_id") is not None:
+            _team_member_user_ids.append(member.get("user_id"))
+
+    sql_query = """
+        SELECT user_email
+        FROM "LiteLLM_UserTable"
+        WHERE user_id = ANY($1::TEXT[]);
+    """
+
+    _result = await prisma_client.db.query_raw(sql_query, _team_member_user_ids)
+
+    verbose_logger.debug("Email Alerting: Got all Emails for team, emails=%s", _result)
+
+    if _result is None:
+        return []
+
+    emails = []
+    for user in _result:
+        if user and isinstance(user, dict) and user.get("user_email", None) is not None:
+            emails.append(user.get("user_email"))
+    return emails
+
+
+async def send_team_budget_alert(webhook_event: WebhookEvent) -> bool:
+    """
+    Send an Email Alert to All Team Members when the Team Budget is crossed
+    Returns -> True if sent, False if not.
+    """
+    from litellm.proxy.utils import send_email
+
+    from litellm.proxy.proxy_server import premium_user, prisma_client
+
+    _team_id = webhook_event.team_id
+    team_alias = webhook_event.team_alias
+    verbose_logger.debug(
+        "Email Alerting: Sending Team Budget Alert for team=%s", team_alias
+    )
+
+    email_logo_url = os.getenv("SMTP_SENDER_LOGO", os.getenv("EMAIL_LOGO_URL", None))
+    email_support_contact = os.getenv("EMAIL_SUPPORT_CONTACT", None)
+
+    # await self._check_if_using_premium_email_feature(
+    #     premium_user, email_logo_url, email_support_contact
+    # )
+
+    if email_logo_url is None:
+        email_logo_url = LITELLM_LOGO_URL
+    if email_support_contact is None:
+        email_support_contact = LITELLM_SUPPORT_CONTACT
+    recipient_emails = await get_all_team_member_emails(_team_id)
+    recipient_emails_str: str = ",".join(recipient_emails)
+    verbose_logger.debug(
+        "Email Alerting: Sending team budget alert to %s", recipient_emails_str
+    )
+
+    event_name = webhook_event.event_message
+    max_budget = webhook_event.max_budget
+    email_html_content = "Alert from LiteLLM Server"
+
+    if recipient_emails_str is None:
+        verbose_proxy_logger.error(
+            "Email Alerting: Trying to send email alert to no recipient, got recipient_emails=%s",
+            recipient_emails_str,
+        )
+
+    email_html_content = f"""
+    <img src="{email_logo_url}" alt="LiteLLM Logo" width="150" height="50" /> <br/><br/><br/>
+
+    Budget Crossed for Team <b> {team_alias} </b> <br/> <br/>
+
+    Your Teams LLM API usage has crossed it's <b> budget of ${max_budget} </b>, current spend is <b>${webhook_event.spend}</b><br /> <br />
+
+    API requests will be rejected until either (a) you increase your budget or (b) your budget gets reset <br /> <br />
+
+    If you have any questions, please send an email to {email_support_contact} <br /> <br />
+
+    Best, <br />
+    The LiteLLM team <br />
+    """
+
+    email_event = {
+        "to": recipient_emails_str,
+        "subject": f"LiteLLM {event_name} for Team {team_alias}",
+        "html": email_html_content,
+    }
+
+    await send_email(
+        receiver_email=email_event["to"],
+        subject=email_event["subject"],
+        html=email_event["html"],
+    )
+
+    return False

litellm/integrations/slack_alerting.py

@@ -1333,7 +1333,9 @@ Model Info:
             verbose_proxy_logger.error("Error sending email alert %s", str(e))
             return False
 
-    async def send_email_alert_using_smtp(self, webhook_event: WebhookEvent) -> bool:
+    async def send_email_alert_using_smtp(
+        self, webhook_event: WebhookEvent, alert_type: str
+    ) -> bool:
         """
         Sends structured Email alert to an SMTP server
 
@@ -1342,7 +1344,6 @@ Model Info:
         Returns -> True if sent, False if not.
         """
         from litellm.proxy.utils import send_email
-
         from litellm.proxy.proxy_server import premium_user, prisma_client
 
         email_logo_url = os.getenv(
@@ -1396,6 +1397,10 @@ Model Info:
             subject=email_event["subject"],
             html=email_event["html"],
         )
+        if webhook_event.event_group == "team":
+            from litellm.integrations.email_alerting import send_team_budget_alert
+
+            await send_team_budget_alert(webhook_event=webhook_event)
 
         return False
 
@@ -1438,7 +1443,9 @@ Model Info:
             and user_info is not None
         ):
             # only send budget alerts over Email
-            await self.send_email_alert_using_smtp(webhook_event=user_info)
+            await self.send_email_alert_using_smtp(
+                webhook_event=user_info, alert_type=alert_type
+            )
 
         if "slack" not in self.alerting:
             return
@@ -1480,7 +1487,7 @@ Model Info:
             slack_webhook_url = os.getenv("SLACK_WEBHOOK_URL", None)
 
         if slack_webhook_url is None:
-            raise Exception("Missing SLACK_WEBHOOK_URL from environment")
+            raise ValueError("Missing SLACK_WEBHOOK_URL from environment")
         payload = {"text": formatted_message}
         headers = {"Content-type": "application/json"}
 

litellm/proxy/_types.py

@@ -1362,6 +1362,7 @@ class CallInfo(LiteLLMBase):
     customer_id: Optional[str] = None
     user_id: Optional[str] = None
     team_id: Optional[str] = None
+    team_alias: Optional[str] = None
     user_email: Optional[str] = None
     key_alias: Optional[str] = None
     projected_exceeded_date: Optional[str] = None

litellm/proxy/common_utils/management_endpoint_utils.py

@@ -2,6 +2,7 @@ from datetime import datetime
 from functools import wraps
 from litellm.proxy._types import UserAPIKeyAuth, ManagementEndpointLoggingPayload
 from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
+from litellm._logging import verbose_logger
 from fastapi import Request
 
 
@@ -20,53 +21,58 @@ def management_endpoint_wrapper(func):
         try:
             result = await func(*args, **kwargs)
             end_time = datetime.now()
+            try:
+                if kwargs is None:
+                    kwargs = {}
+                user_api_key_dict: UserAPIKeyAuth = (
+                    kwargs.get("user_api_key_dict") or UserAPIKeyAuth()
+                )
+                _http_request: Request = kwargs.get("http_request")
+                parent_otel_span = user_api_key_dict.parent_otel_span
+                if parent_otel_span is not None:
+                    from litellm.proxy.proxy_server import open_telemetry_logger
 
-            if kwargs is None:
-                kwargs = {}
-            user_api_key_dict: UserAPIKeyAuth = (
-                kwargs.get("user_api_key_dict") or UserAPIKeyAuth()
-            )
-            _http_request: Request = kwargs.get("http_request")
-            parent_otel_span = user_api_key_dict.parent_otel_span
-            if parent_otel_span is not None:
-                from litellm.proxy.proxy_server import open_telemetry_logger
+                    if open_telemetry_logger is not None:
+                        if _http_request:
+                            _route = _http_request.url.path
+                            _request_body: dict = await _read_request_body(
+                                request=_http_request
+                            )
+                            _response = dict(result) if result is not None else None
+
+                            logging_payload = ManagementEndpointLoggingPayload(
+                                route=_route,
+                                request_data=_request_body,
+                                response=_response,
+                                start_time=start_time,
+                                end_time=end_time,
+                            )
+
+                            await open_telemetry_logger.async_management_endpoint_success_hook(
+                                logging_payload=logging_payload,
+                                parent_otel_span=parent_otel_span,
+                            )
 
-                if open_telemetry_logger is not None:
                     if _http_request:
                         _route = _http_request.url.path
-                        _request_body: dict = await _read_request_body(
-                            request=_http_request
-                        )
-                        _response = dict(result) if result is not None else None
+                        # Flush user_api_key cache if this was an update/delete call to /key, /team, or /user
+                        if _route in [
+                            "/key/update",
+                            "/key/delete",
+                            "/team/update",
+                            "/team/delete",
+                            "/user/update",
+                            "/user/delete",
+                            "/customer/update",
+                            "/customer/delete",
+                        ]:
+                            from litellm.proxy.proxy_server import user_api_key_cache
 
-                        logging_payload = ManagementEndpointLoggingPayload(
-                            route=_route,
-                            request_data=_request_body,
-                            response=_response,
-                            start_time=start_time,
-                            end_time=end_time,
-                        )
-
-                        await open_telemetry_logger.async_management_endpoint_success_hook(
-                            logging_payload=logging_payload,
-                            parent_otel_span=parent_otel_span,
-                        )
-            if _http_request:
-                _route = _http_request.url.path
-                # Flush user_api_key cache if this was an update/delete call to /key, /team, or /user
-                if _route in [
-                    "/key/update",
-                    "/key/delete",
-                    "/team/update",
-                    "/team/delete",
-                    "/user/update",
-                    "/user/delete",
-                    "/customer/update",
-                    "/customer/delete",
-                ]:
-                    from litellm.proxy.proxy_server import user_api_key_cache
-
-                    user_api_key_cache.flush_cache()
+                            user_api_key_cache.flush_cache()
+            except Exception as e:
+                # Non-Blocking Exception
+                verbose_logger.debug("Error in management endpoint wrapper: %s", str(e))
+                pass
 
             return result
         except Exception as e:

litellm/proxy/proxy_config.yaml

@@ -20,6 +20,7 @@ model_list:
 
 general_settings:
   master_key: sk-1234
+  alerting: ["slack", "email"]
 
 litellm_settings:
   success_callback: ["prometheus"]

litellm/proxy/proxy_server.py

@@ -1269,6 +1269,8 @@ async def user_api_key_auth(
                     spend=valid_token.team_spend,
                     max_budget=valid_token.team_max_budget,
                     user_id=valid_token.user_id,
+                    team_id=valid_token.team_id,
+                    team_alias=valid_token.team_alias,
                 )
                 asyncio.create_task(
                     proxy_logging_obj.budget_alerts(
@@ -8913,7 +8915,7 @@ async def new_user(data: NewUserRequest):
                 ),
             ),
             http_request=Request(
-                scope={"type": "http"},
+                scope={"type": "http", "path": "/user/new"},
             ),
         )
 

litellm/proxy/utils.py

@@ -1967,6 +1967,9 @@ async def send_email(receiver_email, subject, html):
     email_message["From"] = sender_email
     email_message["To"] = receiver_email
     email_message["Subject"] = subject
+    verbose_proxy_logger.debug(
+        "sending email from %s to %s", sender_email, receiver_email
+    )
 
     # Attach the body to the email
     email_message.attach(MIMEText(html, "html"))