fix: implement sync service failure hook

This commit is contained in:
Arc-hit 2026-05-01 14:41:45 +05:30
parent eab0075353
commit e8e0ed2927
2 changed files with 67 additions and 3 deletions

View file

@ -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,

View file

@ -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"