mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix: implement sync service failure hook
This commit is contained in:
parent
eab0075353
commit
e8e0ed2927
2 changed files with 67 additions and 3 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue