From e5a9b5c113b9d8107d4d60bef61cba1354aa77b2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 13:51:03 -0700 Subject: [PATCH] fix(alerting): deliver every distinct alert queued in one flush window --- .../SlackAlerting/batching_handler.py | 40 ++++---- .../SlackAlerting/slack_alerting.py | 12 +-- litellm/types/integrations/slack_alerting.py | 15 ++- .../SlackAlerting/test_batching_handler.py | 98 +++++++++++++++++++ 4 files changed, 136 insertions(+), 29 deletions(-) create mode 100644 tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py diff --git a/litellm/integrations/SlackAlerting/batching_handler.py b/litellm/integrations/SlackAlerting/batching_handler.py index 1c35a15d5a1..d152985a2c5 100644 --- a/litellm/integrations/SlackAlerting/batching_handler.py +++ b/litellm/integrations/SlackAlerting/batching_handler.py @@ -1,14 +1,18 @@ """ Handles Batching + sending Httpx Post requests to slack -Slack alerts are sent every 10s or when events are greater than X events +Slack alerts are sent every DEFAULT_FLUSH_INTERVAL_SECONDS or when events are greater than X events see custom_batch_logger.py for more details / defaults """ +from collections import Counter +from collections.abc import Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger +from litellm.types.integrations.slack_alerting import AlertQueueItem, AlertType from .ms_teams import MS_TEAMS_ALERTING_DESTINATION, build_ms_teams_payload @@ -20,26 +24,20 @@ else: SlackAlertingType = Any -def squash_payloads(queue): - squashed: Final = {} - if len(queue) == 0: - return squashed - if len(queue) == 1: - return {"key": {"item": queue[0], "count": 1}} +@dataclass(frozen=True, slots=True) +class SquashedAlert: + item: AlertQueueItem + count: int - for item in queue: - url = item["url"] - alert_type = item["alert_type"] - _key = (url, alert_type) - if _key in squashed: - squashed[_key]["count"] += 1 - # Merge the payloads +def _squash_key(item: AlertQueueItem) -> tuple[str, AlertType | str, str]: + return (item["url"], item["alert_type"], item["payload"]["text"]) - else: - squashed[_key] = {"item": item, "count": 1} - return squashed +def squash_payloads(queue: Sequence[AlertQueueItem]) -> tuple[SquashedAlert, ...]: + counts: Final = Counter(_squash_key(item) for item in queue) + first_item_by_key: Final = {_squash_key(item): item for item in reversed(queue)} + return tuple(SquashedAlert(item=first_item_by_key[key], count=count) for key, count in counts.items()) def _print_alerting_payload_warning(payload: dict, slackAlertingInstance: SlackAlertingType): @@ -53,17 +51,15 @@ def _print_alerting_payload_warning(payload: dict, slackAlertingInstance: SlackA verbose_proxy_logger.warning(payload) -async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item, count): +async def send_to_webhook(slackAlertingInstance: SlackAlertingType, item: AlertQueueItem, count: int) -> None: """ Send a single slack alert to the webhook """ import json - payload: Final = item.get("payload", {}) + text: Final = item["payload"]["text"] + payload: Final = {"text": text if count == 1 else f"[Num Alerts: {count}]\n\n{text}"} try: - if count > 1: - payload["text"] = f"[Num Alerts: {count}]\n\n{payload['text']}" - request_body: Final = ( build_ms_teams_payload(payload["text"]) if item.get("format") == MS_TEAMS_ALERTING_DESTINATION else payload ) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 66e2754d5ad..7b579546403 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -1583,12 +1583,12 @@ Model Info: if not self.log_queue: return - squashed_queue: Final = squash_payloads(self.log_queue) - tasks: Final = [ - send_to_webhook(slackAlertingInstance=self, item=item["item"], count=item["count"]) - for item in squashed_queue.values() - ] - await asyncio.gather(*tasks) + await asyncio.gather( + *( + send_to_webhook(slackAlertingInstance=self, item=squashed.item, count=squashed.count) + for squashed in squash_payloads(self.log_queue) + ) + ) self.log_queue.clear() async def _flush_digest_buckets(self): diff --git a/litellm/types/integrations/slack_alerting.py b/litellm/types/integrations/slack_alerting.py index 64c0c530e9b..33bb446364e 100644 --- a/litellm/types/integrations/slack_alerting.py +++ b/litellm/types/integrations/slack_alerting.py @@ -1,11 +1,12 @@ import os import time +from collections.abc import Mapping from datetime import datetime as dt from enum import Enum from typing import Any, Final, Literal, Optional, Union from pydantic import BaseModel, Field -from typing_extensions import TypedDict +from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm.types.utils import LiteLLMPydanticObjectBase @@ -235,6 +236,18 @@ DEFAULT_ALERT_TYPES: Final[list[AlertType]] = [ ] +class AlertText(TypedDict): + text: ReadOnly[str] + + +class AlertQueueItem(TypedDict): + url: ReadOnly[str] + headers: ReadOnly[Mapping[str, str]] + payload: ReadOnly[AlertText] + alert_type: ReadOnly[AlertType | str] + format: NotRequired[ReadOnly[str]] + + class HangingRequestData(BaseModel): request_id: str model: str diff --git a/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py b/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py new file mode 100644 index 00000000000..9052cb8bb5d --- /dev/null +++ b/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py @@ -0,0 +1,98 @@ +import json +from typing import Final +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from litellm.integrations.SlackAlerting.ms_teams import MS_TEAMS_WEBHOOK_URL_ENV +from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting +from litellm.types.integrations.slack_alerting import AlertQueueItem, AlertType + +SLACK_WEBHOOK_URL: Final = "https://hooks.slack.com/services/test" +THRESHOLD_ALERT: Final = "User Budget: 15% or less of budget remaining\n\n*user_id:* `user-a`" +CROSSED_ALERT: Final = "User Budget: Budget Crossed\n\n*user_id:* `user-b`" + + +def _slack_alerting_recording_posts(alerting: list[str]) -> SlackAlerting: + slack_alerting: Final = SlackAlerting(alerting=alerting) + slack_alerting.periodic_started = True + response: Final = MagicMock() + response.status_code = 200 + slack_alerting.async_http_handler = MagicMock() + slack_alerting.async_http_handler.post = AsyncMock(return_value=response) + return slack_alerting + + +def _queued_slack_alert(text: str) -> AlertQueueItem: + return { + "url": SLACK_WEBHOOK_URL, + "headers": {"Content-type": "application/json"}, + "payload": {"text": text}, + "alert_type": AlertType.budget_alerts, + } + + +def _posted_bodies(slack_alerting: SlackAlerting) -> tuple[dict, ...]: + return tuple(json.loads(call.kwargs["data"]) for call in slack_alerting.async_http_handler.post.call_args_list) + + +async def _send_budget_alert(slack_alerting: SlackAlerting, message: str) -> None: + await slack_alerting.send_alert( + message=message, + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) + + +@pytest.mark.asyncio +async def test_async_send_batch_delivers_every_distinct_alert_queued_in_one_flush(monkeypatch): + monkeypatch.setenv("SLACK_WEBHOOK_URL", SLACK_WEBHOOK_URL) + slack_alerting: Final = _slack_alerting_recording_posts(["slack"]) + await _send_budget_alert(slack_alerting, THRESHOLD_ALERT) + await _send_budget_alert(slack_alerting, CROSSED_ALERT) + + await slack_alerting.async_send_batch() + + posted_texts: Final = tuple(body["text"] for body in _posted_bodies(slack_alerting)) + assert len(posted_texts) == 2 + assert THRESHOLD_ALERT in posted_texts[0] + assert CROSSED_ALERT in posted_texts[1] + assert not any(text.startswith("[Num Alerts") for text in posted_texts) + assert slack_alerting.log_queue == [] + + +@pytest.mark.asyncio +async def test_async_send_batch_collapses_only_identical_alerts(): + slack_alerting: Final = _slack_alerting_recording_posts(["slack"]) + slack_alerting.log_queue.extend( + ( + _queued_slack_alert(THRESHOLD_ALERT), + _queued_slack_alert(CROSSED_ALERT), + _queued_slack_alert(THRESHOLD_ALERT), + ) + ) + + await slack_alerting.async_send_batch() + + assert _posted_bodies(slack_alerting) == ( + {"text": f"[Num Alerts: 2]\n\n{THRESHOLD_ALERT}"}, + {"text": CROSSED_ALERT}, + ) + + +@pytest.mark.asyncio +async def test_async_send_batch_delivers_every_distinct_ms_teams_alert(monkeypatch): + monkeypatch.setenv(MS_TEAMS_WEBHOOK_URL_ENV, "https://teams.example/webhook") + slack_alerting: Final = _slack_alerting_recording_posts(["ms_teams"]) + await _send_budget_alert(slack_alerting, THRESHOLD_ALERT) + await _send_budget_alert(slack_alerting, CROSSED_ALERT) + + await slack_alerting.async_send_batch() + + card_texts: Final = tuple( + body["attachments"][0]["content"]["body"][0]["text"] for body in _posted_bodies(slack_alerting) + ) + assert len(card_texts) == 2 + assert THRESHOLD_ALERT in card_texts[0] + assert CROSSED_ALERT in card_texts[1]