From 0bdf31c4f3c20e5db5adc29cebf36ad31ad90b11 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 7 May 2024 18:03:04 -0700
Subject: [PATCH] router- initialize alerting

---
 litellm/router.py       | 24 +++++++++++++++++++++++-
 litellm/types/router.py |  8 ++++----
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/litellm/router.py b/litellm/router.py
index 3f2bef4768..ddb006d527 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -44,6 +44,7 @@ from litellm.types.router import (
     updateDeployment,
     updateLiteLLMParams,
     RetryPolicy,
+    AlertingConfig,
 )
 from litellm.integrations.custom_logger import CustomLogger
 
@@ -103,6 +104,7 @@ class Router:
         ] = "simple-shuffle",
         routing_strategy_args: dict = {},  # just for latency-based routing
         semaphore: Optional[asyncio.Semaphore] = None,
+        alerting_config: Optional[AlertingConfig] = None,
     ) -> None:
         """
         Initialize the Router class with the given parameters for caching, reliability, and routing strategy.
@@ -131,7 +133,7 @@ class Router:
             cooldown_time (float): Time to cooldown a deployment after failure in seconds. Defaults to 1.
             routing_strategy (Literal["simple-shuffle", "least-busy", "usage-based-routing", "latency-based-routing", "cost-based-routing"]): Routing strategy. Defaults to "simple-shuffle".
             routing_strategy_args (dict): Additional args for latency-based routing. Defaults to {}.
-
+            alerting_config (AlertingConfig): Slack alerting configuration. Defaults to None.
         Returns:
             Router: An instance of the litellm.Router class.
 
@@ -316,6 +318,9 @@ class Router:
         self.model_group_retry_policy: Optional[Dict[str, RetryPolicy]] = (
             model_group_retry_policy
         )
+        self.alerting_config: Optional[AlertingConfig] = alerting_config
+        if self.alerting_config is not None:
+            self._initialize_alerting()
 
     def routing_strategy_init(self, routing_strategy: str, routing_strategy_args: dict):
         if routing_strategy == "least-busy":
@@ -3320,6 +3325,23 @@ class Router:
         ):
             return retry_policy.ContentPolicyViolationErrorRetries
 
+    def _initialize_alerting(self):
+        from litellm.integrations.slack_alerting import SlackAlerting
+
+        router_alerting_config: AlertingConfig = self.alerting_config
+
+        _slack_alerting_logger = SlackAlerting(
+            alerting_threshold=router_alerting_config.alerting_threshold,
+            alerting=["slack"],
+            default_webhook_url=router_alerting_config.webhook_url,
+        )
+
+        litellm.callbacks.append(_slack_alerting_logger)
+        litellm.success_callback.append(
+            _slack_alerting_logger.response_taking_too_long_callback
+        )
+        print("\033[94m\nInitialized Alerting for litellm.Router\033[0m\n")  # noqa
+
     def flush_cache(self):
         litellm.cache = None
         self.cache.flush_cache()
diff --git a/litellm/types/router.py b/litellm/types/router.py
index d79fb2e2e1..f3fa893246 100644
--- a/litellm/types/router.py
+++ b/litellm/types/router.py
@@ -348,7 +348,7 @@ class RetryPolicy(BaseModel):
     InternalServerErrorRetries: Optional[int] = None
 
 
-class RouterAlerting(BaseModel):
+class AlertingConfig(BaseModel):
     """
     Use this configure alerting for the router. Receive alerts on the following events
     - LLM API Exceptions
@@ -356,9 +356,9 @@ class RouterAlerting(BaseModel):
     - LLM Requests Hanging
 
     Args:
-        webhook_url: Optional[str] = None            - webhook url for alerting
+        webhook_url: str            - webhook url for alerting, slack provides a webhook url to send alerts to
         alerting_threshold: Optional[float] = None - threhshold for slow / hanging llm responses (in seconds)
     """
 
-    webhook_url: Optional[str] = None
-    alerting_threshold: Optional[float] = None
+    webhook_url: str
+    alerting_threshold: Optional[float] = 300