From 42e6e664b22dbfa03d5e05cde816fb45ed951d1f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 15 May 2025 10:12:06 -0700 Subject: [PATCH] [Refactor] Make Pagerduty a free feature (#10857) * refactor: make pagerduty free * refactor: make pagerduty free * fix: pagerduty loc * fix: linting error --- .../pagerduty/pagerduty.py | 28 +++++++++---------- litellm/litellm_core_utils/litellm_logging.py | 5 +++- .../test_pagerduty_alerting.py | 6 +++- .../test_unit_tests_init_callbacks.py | 2 +- 4 files changed, 23 insertions(+), 18 deletions(-) rename {litellm/integrations => enterprise/litellm_enterprise/enterprise_callbacks}/pagerduty/pagerduty.py (93%) diff --git a/litellm/integrations/pagerduty/pagerduty.py b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py similarity index 93% rename from litellm/integrations/pagerduty/pagerduty.py rename to enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py index 6085bc237a..773c34401d 100644 --- a/litellm/integrations/pagerduty/pagerduty.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/pagerduty/pagerduty.py @@ -4,6 +4,10 @@ PagerDuty Alerting Integration Handles two types of alerts: - High LLM API Failure Rate. Configure X fails in Y seconds to trigger an alert. - High Number of Hanging LLM Requests. Configure X hangs in Y seconds to trigger an alert. + +Note: This is a Free feature on the regular litellm docker image. + +However, this is under the enterprise license """ import asyncio @@ -46,8 +50,6 @@ class PagerDutyAlerting(SlackAlerting): def __init__( self, alerting_args: Optional[Union[AlertingConfig, dict]] = None, **kwargs ): - from litellm.proxy.proxy_server import CommonProxyErrors, premium_user - super().__init__() _api_key = os.getenv("PAGERDUTY_API_KEY") if not _api_key: @@ -55,7 +57,7 @@ class PagerDutyAlerting(SlackAlerting): self.api_key: str = _api_key alerting_args = alerting_args or {} - self.alerting_args: AlertingConfig = AlertingConfig( + self.pagerduty_alerting_args: AlertingConfig = AlertingConfig( failure_threshold=alerting_args.get( "failure_threshold", PAGERDUTY_DEFAULT_FAILURE_THRESHOLD ), @@ -76,12 +78,6 @@ class PagerDutyAlerting(SlackAlerting): self._failure_events: List[PagerDutyInternalEvent] = [] self._hanging_events: List[PagerDutyInternalEvent] = [] - # premium user check - if premium_user is not True: - raise ValueError( - f"PagerDutyAlerting is only available for LiteLLM Enterprise users. {CommonProxyErrors.not_premium_user.value}" - ) - # ------------------ MAIN LOGIC ------------------ # async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): @@ -123,8 +119,10 @@ class PagerDutyAlerting(SlackAlerting): ) # Prune + Possibly alert - window_seconds = self.alerting_args.get("failure_threshold_window_seconds", 60) - threshold = self.alerting_args.get("failure_threshold", 1) + window_seconds = self.pagerduty_alerting_args.get( + "failure_threshold_window_seconds", 60 + ) + threshold = self.pagerduty_alerting_args.get("failure_threshold", 1) # If threshold is crossed, send PD alert for failures await self._send_alert_if_thresholds_crossed( @@ -170,10 +168,10 @@ class PagerDutyAlerting(SlackAlerting): If not, we classify it as a hanging request. """ verbose_logger.debug( - f"Inside Hanging Response Handler!..sleeping for {self.alerting_args.get('hanging_threshold_seconds', PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS)} seconds" + f"Inside Hanging Response Handler!..sleeping for {self.pagerduty_alerting_args.get('hanging_threshold_seconds', PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS)} seconds" ) await asyncio.sleep( - self.alerting_args.get( + self.pagerduty_alerting_args.get( "hanging_threshold_seconds", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS ) ) @@ -201,11 +199,11 @@ class PagerDutyAlerting(SlackAlerting): ) # Prune + Possibly alert - window_seconds = self.alerting_args.get( + window_seconds = self.pagerduty_alerting_args.get( "hanging_threshold_window_seconds", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_WINDOW_SECONDS, ) - threshold: int = self.alerting_args.get( + threshold: int = self.pagerduty_alerting_args.get( "hanging_threshold_fails", PAGERDUTY_DEFAULT_HANGING_THRESHOLD_SECONDS ) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index bfcbb3e14d..012b655810 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -53,7 +53,6 @@ from litellm.integrations.arize.arize import ArizeLogger from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.mlflow import MlflowLogger -from litellm.integrations.pagerduty.pagerduty import PagerDutyAlerting from litellm.integrations.vector_stores.bedrock_vector_store import BedrockVectorStore from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import ( @@ -148,6 +147,9 @@ try: from litellm_enterprise.enterprise_callbacks.generic_api_callback import ( GenericAPILogger, ) + from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import ( + PagerDutyAlerting, + ) from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ( ResendEmailLogger, ) @@ -168,6 +170,7 @@ except Exception as e: GenericAPILogger = CustomLogger # type: ignore ResendEmailLogger = CustomLogger # type: ignore SMTPEmailLogger = CustomLogger # type: ignore + PagerDutyAlerting = CustomLogger # type: ignore EnterpriseStandardLoggingPayloadSetupVAR = None _in_memory_loggers: List[Any] = [] diff --git a/tests/logging_callback_tests/test_pagerduty_alerting.py b/tests/logging_callback_tests/test_pagerduty_alerting.py index 00e427d01d..33c24102eb 100644 --- a/tests/logging_callback_tests/test_pagerduty_alerting.py +++ b/tests/logging_callback_tests/test_pagerduty_alerting.py @@ -8,7 +8,11 @@ from typing import Optional sys.path.insert(0, os.path.abspath("../..")) import pytest import litellm -from litellm.integrations.pagerduty.pagerduty import PagerDutyAlerting, AlertingConfig + +from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import ( + PagerDutyAlerting, + AlertingConfig, +) from litellm.proxy._types import UserAPIKeyAuth diff --git a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py index d8137e8aa0..d1b7d46276 100644 --- a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py +++ b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py @@ -20,7 +20,6 @@ from prometheus_client import REGISTRY, CollectorRegistry from litellm.integrations.lago import LagoLogger from litellm.integrations.openmeter import OpenMeterLogger from litellm.integrations.braintrust_logging import BraintrustLogger -from litellm.integrations.pagerduty.pagerduty import PagerDutyAlerting from litellm.integrations.galileo import GalileoObserve from litellm.integrations.langsmith import LangsmithLogger from litellm.integrations.literal_ai import LiteralAILogger @@ -45,6 +44,7 @@ from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHand from litellm_enterprise.enterprise_callbacks.generic_api_callback import GenericAPILogger from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ResendEmailLogger from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import SMTPEmailLogger +from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import PagerDutyAlerting from unittest.mock import patch # clear prometheus collectors / registry