diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index 1a3be203fec..f32e893dbbd 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -99,11 +99,46 @@ class ServiceLogging(CustomLogger): self, service: ServiceTypes, duration: float, error: Exception, call_type: str ): """ - [TODO] Not implemented for sync calls yet. V0 is focused on async monitoring (used by proxy). + Handles both sync and async monitoring by checking for existing event loop. """ if self.mock_testing: self.mock_testing_sync_failure_hook += 1 + try: + # Try to get the current event loop + loop = asyncio.get_event_loop() + # Check if the loop is running + if loop.is_running(): + # If we're in a running loop, create a task + loop.create_task( + self.async_service_failure_hook( + service=service, + duration=duration, + error=error, + call_type=call_type, + ) + ) + else: + # Loop exists but not running, we can use run_until_complete + loop.run_until_complete( + self.async_service_failure_hook( + service=service, + duration=duration, + error=error, + call_type=call_type, + ) + ) + except RuntimeError: + # No event loop exists, create a new one and run + asyncio.run( + self.async_service_failure_hook( + service=service, + duration=duration, + error=error, + call_type=call_type, + ) + ) + async def async_service_success_hook( self, service: ServiceTypes, diff --git a/tests/test_litellm/test_service_logger.py b/tests/test_litellm/test_service_logger.py index ed44fe9b9f2..0a2bc895632 100644 --- a/tests/test_litellm/test_service_logger.py +++ b/tests/test_litellm/test_service_logger.py @@ -5,11 +5,14 @@ Regression test for KeyError: 'call_type' when async_log_success_event is called without call_type in kwargs (e.g. from batch polling callbacks). """ -import pytest -from datetime import datetime, timedelta +import asyncio +from datetime import datetime from unittest.mock import AsyncMock, patch +import pytest + from litellm._service_logger import ServiceLogging +from litellm.types.services import ServiceTypes @pytest.mark.asyncio @@ -95,3 +98,29 @@ async def test_async_log_success_event_should_handle_float_duration(): mock_hook.assert_called_once() call_kwargs = mock_hook.call_args assert call_kwargs.kwargs["duration"] == 1.5 + + +@pytest.mark.asyncio +async def test_service_failure_hook_should_schedule_async_failure_hook(): + service_logger = ServiceLogging(mock_testing=True) + + with patch.object( + service_logger, "async_service_failure_hook", new_callable=AsyncMock + ) as mock_hook: + service_logger.service_failure_hook( + service=ServiceTypes.REDIS, + duration=0.123, + error=Exception("boom"), + call_type="health_check", + ) + + # Let the event loop run scheduled tasks + await asyncio.sleep(0) + + mock_hook.assert_awaited_once() + call_kwargs = mock_hook.call_args + assert call_kwargs.kwargs["service"] == ServiceTypes.REDIS + assert call_kwargs.kwargs["duration"] == 0.123 + assert isinstance(call_kwargs.kwargs["error"], Exception) + assert str(call_kwargs.kwargs["error"]) == "boom" + assert call_kwargs.kwargs["call_type"] == "health_check"