Merge pull request #42314 from BerriAI/litellm_alerting_batch_keeps_distinct_alerts

fix(alerting): deliver every distinct alert queued in one flush window

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
         )

litellm/integrations/SlackAlerting/slack_alerting.py

@@ -33,6 +33,7 @@ from litellm.litellm_core_utils.exception_mapping_utils import (
     _add_key_name_and_team_to_alert,
 )
 from litellm.llms.custom_httpx.http_handler import (
+    AsyncHTTPHandler,
     get_async_httpx_client,
     httpxSpecialProvider,
 )
@@ -99,6 +100,7 @@ class SlackAlerting(CustomBatchLogger):
         alerting_args={},
         default_webhook_url: str | None = None,
         alert_type_config: dict[str, dict] | None = None,
+        async_http_handler: AsyncHTTPHandler | None = None,
         **kwargs,
     ):
         if alerting_threshold is None:
@@ -107,7 +109,9 @@ class SlackAlerting(CustomBatchLogger):
         self.alerting = alerting
         self.alert_types = alert_types
         self.internal_usage_cache = internal_usage_cache or DualCache()
-        self.async_http_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback)
+        self.async_http_handler = async_http_handler or get_async_httpx_client(
+            llm_provider=httpxSpecialProvider.LoggingCallback
+        )
         self.alert_to_webhook_url = process_slack_alerting_variables(alert_to_webhook_url=alert_to_webhook_url)
         self.is_running = False
         self.alerting_args = SlackAlertingArgs(**alerting_args)
@@ -1583,12 +1587,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):

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

tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py

@@ -2,18 +2,39 @@ import json
 from typing import Final
 from unittest.mock import AsyncMock, MagicMock, patch
 
+import httpx
 import pytest
+from pydantic import TypeAdapter
 
 from litellm.integrations.SlackAlerting.batching_handler import send_to_webhook
 from litellm.integrations.SlackAlerting.ms_teams import (
     MS_TEAMS_ALERTING_DESTINATION,
     MS_TEAMS_WEBHOOK_URL_ENV,
+    MSTeamsMessage,
     build_ms_teams_payload,
     get_ms_teams_webhook_url,
 )
 from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.proxy._types import AlertType
 
+_MS_TEAMS_MESSAGE: Final = TypeAdapter(MSTeamsMessage)
+
+
+def _webhook_accepting_posts() -> AsyncMock:
+    response: Final = MagicMock(spec=httpx.Response)
+    response.status_code = 200
+    http_handler: Final = AsyncMock(spec=AsyncHTTPHandler)
+    http_handler.post.return_value = response
+    return http_handler
+
+
+def _posted_card_texts(http_handler: AsyncMock) -> tuple[str, ...]:
+    return tuple(
+        _MS_TEAMS_MESSAGE.validate_json(call.kwargs["data"])["attachments"][0]["content"]["body"][0]["text"]
+        for call in http_handler.post.call_args_list
+    )
+
 
 def test_build_ms_teams_payload_wraps_text_in_adaptive_card():
     payload: Final = build_ms_teams_payload("hello alert")
@@ -80,11 +101,8 @@ async def test_send_alert_slack_and_ms_teams_enqueue_both(monkeypatch):
 
 @pytest.mark.asyncio
 async def test_send_to_webhook_posts_adaptive_card_for_ms_teams_items():
-    slack_alerting: Final = SlackAlerting(alerting=["ms_teams"])
-    mock_response: Final = MagicMock()
-    mock_response.status_code = 200
-    slack_alerting.async_http_handler = MagicMock()
-    slack_alerting.async_http_handler.post = AsyncMock(return_value=mock_response)
+    http_handler: Final = _webhook_accepting_posts()
+    slack_alerting: Final = SlackAlerting(alerting=["ms_teams"], async_http_handler=http_handler)
 
     item: Final = {
         "url": "https://teams.example/webhook",
@@ -95,7 +113,7 @@ async def test_send_to_webhook_posts_adaptive_card_for_ms_teams_items():
     }
     await send_to_webhook(slackAlertingInstance=slack_alerting, item=item, count=1)
 
-    call_kwargs: Final = slack_alerting.async_http_handler.post.call_args.kwargs
+    call_kwargs: Final = http_handler.post.call_args.kwargs
     assert call_kwargs["url"] == "https://teams.example/webhook"
     sent_body: Final = json.loads(call_kwargs["data"])
     assert sent_body["type"] == "message"
@@ -104,11 +122,8 @@ async def test_send_to_webhook_posts_adaptive_card_for_ms_teams_items():
 
 @pytest.mark.asyncio
 async def test_send_to_webhook_keeps_slack_payload_shape():
-    slack_alerting: Final = SlackAlerting(alerting=["slack"])
-    mock_response: Final = MagicMock()
-    mock_response.status_code = 200
-    slack_alerting.async_http_handler = MagicMock()
-    slack_alerting.async_http_handler.post = AsyncMock(return_value=mock_response)
+    http_handler: Final = _webhook_accepting_posts()
+    slack_alerting: Final = SlackAlerting(alerting=["slack"], async_http_handler=http_handler)
 
     item: Final = {
         "url": "https://hooks.slack.com/services/test",
@@ -118,5 +133,27 @@ async def test_send_to_webhook_keeps_slack_payload_shape():
     }
     await send_to_webhook(slackAlertingInstance=slack_alerting, item=item, count=1)
 
-    call_kwargs: Final = slack_alerting.async_http_handler.post.call_args.kwargs
+    call_kwargs: Final = http_handler.post.call_args.kwargs
     assert json.loads(call_kwargs["data"]) == {"text": "alert body"}
+
+
+@pytest.mark.asyncio
+async def test_async_send_batch_delivers_every_distinct_ms_teams_alert(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv(MS_TEAMS_WEBHOOK_URL_ENV, "https://teams.example/webhook")
+    http_handler: Final = _webhook_accepting_posts()
+    slack_alerting: Final = SlackAlerting(alerting=["ms_teams"], async_http_handler=http_handler)
+    slack_alerting.periodic_started = True
+    for message in ("User Budget: 15% or less of budget remaining", "User Budget: Budget Crossed"):
+        await slack_alerting.send_alert(
+            message=message,
+            level="High",
+            alert_type=AlertType.budget_alerts,
+            alerting_metadata={},
+        )
+
+    await slack_alerting.async_send_batch()
+
+    card_texts: Final = _posted_card_texts(http_handler)
+    assert len(card_texts) == 2
+    assert "User Budget: 15% or less of budget remaining" in card_texts[0]
+    assert "User Budget: Budget Crossed" in card_texts[1]

tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py

@@ -6,14 +6,18 @@ import unittest
 from typing import Final, List, Optional, Tuple
 from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch
 
+import httpx
 import pytest
+from pydantic import TypeAdapter
+from typing_extensions import ReadOnly, TypedDict
 
 import litellm
 from litellm.caching.caching import DualCache
 from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type
 from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
 from litellm.proxy._types import CallInfo, Litellm_EntityType
-from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingCacheKeys
+from litellm.types.integrations.slack_alerting import AlertQueueItem, AlertType, SlackAlertingCacheKeys
 
 
 class TestSlackAlerting(unittest.TestCase):
@@ -434,3 +438,91 @@ async def test_send_alert_raises_when_no_webhook_url_configured(monkeypatch):
             alert_type=AlertType.budget_alerts,
             alerting_metadata={},
         )
+
+
+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`"
+
+
+class _SlackWebhookBody(TypedDict):
+    text: ReadOnly[str]
+
+
+_SLACK_WEBHOOK_BODY: Final = TypeAdapter(_SlackWebhookBody)
+
+
+def _webhook_accepting_posts() -> AsyncMock:
+    response: Final = MagicMock(spec=httpx.Response)
+    response.status_code = 200
+    http_handler: Final = AsyncMock(spec=AsyncHTTPHandler)
+    http_handler.post.return_value = response
+    return http_handler
+
+
+def _slack_alerting_flushing_to(http_handler: AsyncHTTPHandler) -> SlackAlerting:
+    slack_alerting: Final = SlackAlerting(alerting=["slack"], async_http_handler=http_handler)
+    slack_alerting.periodic_started = True
+    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_slack_bodies(http_handler: AsyncMock) -> tuple[_SlackWebhookBody, ...]:
+    return tuple(_SLACK_WEBHOOK_BODY.validate_json(call.kwargs["data"]) for call in 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: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("SLACK_WEBHOOK_URL", SLACK_WEBHOOK_URL)
+    http_handler: Final = _webhook_accepting_posts()
+    slack_alerting: Final = _slack_alerting_flushing_to(http_handler)
+    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_slack_bodies(http_handler))
+    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() -> None:
+    http_handler: Final = _webhook_accepting_posts()
+    slack_alerting: Final = _slack_alerting_flushing_to(http_handler)
+    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_slack_bodies(http_handler) == (
+        {"text": f"[Num Alerts: 2]\n\n{THRESHOLD_ALERT}"},
+        {"text": CROSSED_ALERT},
+    )