diff --git a/litellm/integrations/slack_alerting.py b/litellm/integrations/slack_alerting.py index a27e16b3c..5546f7c33 100644 --- a/litellm/integrations/slack_alerting.py +++ b/litellm/integrations/slack_alerting.py @@ -797,6 +797,48 @@ Model Info: ) ) + async def _run_scheduler_helper(self, llm_router: litellm.Router) -> bool: + """ + Returns: + - True -> report sent + - False -> report not sent + """ + report_sent_bool = False + + report_sent = await self.internal_usage_cache.async_get_cache( + key=SlackAlertingCacheKeys.report_sent_key.value + ) # None | datetime + + current_time = litellm.utils.get_utc_datetime() + + if report_sent is None: + _current_time = current_time.isoformat() + await self.internal_usage_cache.async_set_cache( + key=SlackAlertingCacheKeys.report_sent_key.value, + value=_current_time, + ) + else: + # check if current time - interval >= time last sent + delta = current_time - timedelta( + seconds=self.alerting_args.daily_report_frequency + ) + + if isinstance(report_sent, str): + report_sent = dt.fromisoformat(report_sent) + + if delta >= report_sent: + # Sneak in the reporting logic here + await self.send_daily_reports(router=llm_router) + # Also, don't forget to update the report_sent time after sending the report! + _current_time = current_time.isoformat() + await self.internal_usage_cache.async_set_cache( + key=SlackAlertingCacheKeys.report_sent_key.value, + value=_current_time, + ) + report_sent_bool = True + + return report_sent_bool + async def _run_scheduled_daily_report(self, llm_router: Optional[litellm.Router]): """ If 'daily_reports' enabled @@ -810,29 +852,7 @@ Model Info: if "daily_reports" in self.alert_types: while True: - report_sent = await self.internal_usage_cache.async_get_cache( - key=SlackAlertingCacheKeys.report_sent_key.value - ) # None | datetime - - if report_sent is None: - await self.internal_usage_cache.async_set_cache( - key=SlackAlertingCacheKeys.report_sent_key.value, - value=litellm.utils.get_utc_datetime(), - ) - else: - # check if current time - interval >= time last sent - current_time = litellm.utils.get_utc_datetime() - delta = current_time - timedelta( - seconds=self.alerting_args.daily_report_frequency - ) - if delta >= report_sent: - # Sneak in the reporting logic here - await self.send_daily_reports(router=llm_router) - # Also, don't forget to update the report_sent time after sending the report! - await self.internal_usage_cache.async_set_cache( - key=SlackAlertingCacheKeys.report_sent_key.value, - value=litellm.utils.get_utc_datetime(), - ) + await self._run_scheduler_helper(llm_router=llm_router) interval = random.randint( self.alerting_args.report_check_interval - 3, self.alerting_args.report_check_interval + 3, diff --git a/litellm/tests/test_alerting.py b/litellm/tests/test_alerting.py index ba1022472..3734c29d2 100644 --- a/litellm/tests/test_alerting.py +++ b/litellm/tests/test_alerting.py @@ -1,7 +1,7 @@ # What is this? ## Tests slack alerting on proxy logging object -import sys +import sys, json import os import io, asyncio from datetime import datetime, timedelta @@ -10,7 +10,7 @@ from datetime import datetime, timedelta # logging.basicConfig(level=logging.DEBUG) sys.path.insert(0, os.path.abspath("../..")) from litellm.proxy.utils import ProxyLogging -from litellm.caching import DualCache +from litellm.caching import DualCache, RedisCache import litellm import pytest import asyncio @@ -273,3 +273,43 @@ async def test_daily_reports_completion(slack_alerting): assert response_val == True mock_send_alert.assert_awaited() + + +@pytest.mark.asyncio +async def test_daily_reports_redis_cache_scheduler(): + redis_cache = RedisCache() + slack_alerting = SlackAlerting( + internal_usage_cache=DualCache(redis_cache=redis_cache) + ) + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-5", + "litellm_params": { + "model": "gpt-3.5-turbo", + }, + } + ] + ) + + with patch.object( + slack_alerting, "send_alert", new=AsyncMock() + ) as mock_send_alert, patch.object( + redis_cache, "async_set_cache", new=AsyncMock() + ) as mock_redis_set_cache: + # initial call - expect empty + await slack_alerting._run_scheduler_helper(llm_router=router) + + try: + json.dumps(mock_redis_set_cache.call_args[0][1]) + except Exception as e: + pytest.fail( + "Cache value can't be json dumped - {}".format( + mock_redis_set_cache.call_args[0][1] + ) + ) + + mock_redis_set_cache.assert_awaited_once() + + # second call - expect empty + await slack_alerting._run_scheduler_helper(llm_router=router)