diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 7b579546403..8d0d044ff93 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/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) diff --git a/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py b/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py deleted file mode 100644 index 9052cb8bb5d..00000000000 --- a/tests/test_litellm/integrations/SlackAlerting/test_batching_handler.py +++ /dev/null @@ -1,98 +0,0 @@ -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] diff --git a/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py b/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py index 41b7f3b969b..44426c00628 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_ms_teams.py +++ b/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] diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 4bc6c08bd63..2d5eb78950c 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/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}, + )