forked from phoenix/litellm-mirror
(fix) handle json decode errors for DD exception logging (#6934)
* fix JSONDecodeError * handle async_log_proxy_authentication_errors * fix test_async_log_proxy_authentication_errors_get_request
This commit is contained in:
parent
77f714dc51
commit
a8b8deb793
2 changed files with 130 additions and 3 deletions
|
@ -891,7 +891,7 @@ class ProxyLogging:
|
||||||
original_exception: Exception,
|
original_exception: Exception,
|
||||||
request: Request,
|
request: Request,
|
||||||
parent_otel_span: Optional[Any],
|
parent_otel_span: Optional[Any],
|
||||||
api_key: str,
|
api_key: Optional[str],
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Handler for Logging Authentication Errors on LiteLLM Proxy
|
Handler for Logging Authentication Errors on LiteLLM Proxy
|
||||||
|
@ -905,9 +905,13 @@ class ProxyLogging:
|
||||||
|
|
||||||
user_api_key_dict = UserAPIKeyAuth(
|
user_api_key_dict = UserAPIKeyAuth(
|
||||||
parent_otel_span=parent_otel_span,
|
parent_otel_span=parent_otel_span,
|
||||||
token=_hash_token_if_needed(token=api_key),
|
token=_hash_token_if_needed(token=api_key or ""),
|
||||||
)
|
)
|
||||||
request_data = await request.json()
|
try:
|
||||||
|
request_data = await request.json()
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
# For GET requests or requests without a JSON body
|
||||||
|
request_data = {}
|
||||||
await self._run_post_call_failure_hook_custom_loggers(
|
await self._run_post_call_failure_hook_custom_loggers(
|
||||||
original_exception=original_exception,
|
original_exception=original_exception,
|
||||||
request_data=request_data,
|
request_data=request_data,
|
||||||
|
|
|
@ -2195,3 +2195,126 @@ async def test_async_log_proxy_authentication_errors():
|
||||||
assert (
|
assert (
|
||||||
mock_logger.user_api_key_dict_logged.token is not None
|
mock_logger.user_api_key_dict_logged.token is not None
|
||||||
) # token should be hashed
|
) # token should be hashed
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_log_proxy_authentication_errors_get_request():
|
||||||
|
"""
|
||||||
|
Test if async_log_proxy_authentication_errors correctly handles GET requests
|
||||||
|
that don't have a JSON body
|
||||||
|
"""
|
||||||
|
import json
|
||||||
|
from fastapi import Request
|
||||||
|
from litellm.proxy.utils import ProxyLogging
|
||||||
|
from litellm.caching import DualCache
|
||||||
|
from litellm.integrations.custom_logger import CustomLogger
|
||||||
|
|
||||||
|
class MockCustomLogger(CustomLogger):
|
||||||
|
def __init__(self):
|
||||||
|
self.called = False
|
||||||
|
self.exception_logged = None
|
||||||
|
self.request_data_logged = None
|
||||||
|
self.user_api_key_dict_logged = None
|
||||||
|
|
||||||
|
async def async_post_call_failure_hook(
|
||||||
|
self,
|
||||||
|
request_data: dict,
|
||||||
|
original_exception: Exception,
|
||||||
|
user_api_key_dict: UserAPIKeyAuth,
|
||||||
|
):
|
||||||
|
self.called = True
|
||||||
|
self.exception_logged = original_exception
|
||||||
|
self.request_data_logged = request_data
|
||||||
|
self.user_api_key_dict_logged = user_api_key_dict
|
||||||
|
|
||||||
|
# Create a mock GET request
|
||||||
|
request = Request(scope={"type": "http", "method": "GET"})
|
||||||
|
|
||||||
|
# Mock the json() method to raise JSONDecodeError
|
||||||
|
async def mock_json():
|
||||||
|
raise json.JSONDecodeError("Expecting value", "", 0)
|
||||||
|
|
||||||
|
request.json = mock_json
|
||||||
|
|
||||||
|
# Create a test exception
|
||||||
|
test_exception = Exception("Invalid API Key")
|
||||||
|
|
||||||
|
# Initialize ProxyLogging
|
||||||
|
mock_logger = MockCustomLogger()
|
||||||
|
litellm.callbacks = [mock_logger]
|
||||||
|
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||||
|
|
||||||
|
# Call the method
|
||||||
|
await proxy_logging_obj.async_log_proxy_authentication_errors(
|
||||||
|
original_exception=test_exception,
|
||||||
|
request=request,
|
||||||
|
parent_otel_span=None,
|
||||||
|
api_key="test-key",
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify the mock logger was called with correct parameters
|
||||||
|
assert mock_logger.called == True
|
||||||
|
assert mock_logger.exception_logged == test_exception
|
||||||
|
assert mock_logger.user_api_key_dict_logged is not None
|
||||||
|
assert mock_logger.user_api_key_dict_logged.token is not None
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_async_log_proxy_authentication_errors_no_api_key():
|
||||||
|
"""
|
||||||
|
Test if async_log_proxy_authentication_errors correctly handles requests
|
||||||
|
with no API key provided
|
||||||
|
"""
|
||||||
|
from fastapi import Request
|
||||||
|
from litellm.proxy.utils import ProxyLogging
|
||||||
|
from litellm.caching import DualCache
|
||||||
|
from litellm.integrations.custom_logger import CustomLogger
|
||||||
|
|
||||||
|
class MockCustomLogger(CustomLogger):
|
||||||
|
def __init__(self):
|
||||||
|
self.called = False
|
||||||
|
self.exception_logged = None
|
||||||
|
self.request_data_logged = None
|
||||||
|
self.user_api_key_dict_logged = None
|
||||||
|
|
||||||
|
async def async_post_call_failure_hook(
|
||||||
|
self,
|
||||||
|
request_data: dict,
|
||||||
|
original_exception: Exception,
|
||||||
|
user_api_key_dict: UserAPIKeyAuth,
|
||||||
|
):
|
||||||
|
self.called = True
|
||||||
|
self.exception_logged = original_exception
|
||||||
|
self.request_data_logged = request_data
|
||||||
|
self.user_api_key_dict_logged = user_api_key_dict
|
||||||
|
|
||||||
|
# Create test data
|
||||||
|
test_data = {"model": "gpt-4", "messages": [{"role": "user", "content": "Hello"}]}
|
||||||
|
|
||||||
|
# Create a mock request
|
||||||
|
request = Request(scope={"type": "http", "method": "POST"})
|
||||||
|
request._json = AsyncMock(return_value=test_data)
|
||||||
|
|
||||||
|
# Create a test exception
|
||||||
|
test_exception = Exception("No API Key Provided")
|
||||||
|
|
||||||
|
# Initialize ProxyLogging
|
||||||
|
mock_logger = MockCustomLogger()
|
||||||
|
litellm.callbacks = [mock_logger]
|
||||||
|
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
|
||||||
|
|
||||||
|
# Call the method with api_key=None
|
||||||
|
await proxy_logging_obj.async_log_proxy_authentication_errors(
|
||||||
|
original_exception=test_exception,
|
||||||
|
request=request,
|
||||||
|
parent_otel_span=None,
|
||||||
|
api_key=None,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Verify the mock logger was called with correct parameters
|
||||||
|
assert mock_logger.called == True
|
||||||
|
assert mock_logger.exception_logged == test_exception
|
||||||
|
assert mock_logger.user_api_key_dict_logged is not None
|
||||||
|
assert (
|
||||||
|
mock_logger.user_api_key_dict_logged.token == ""
|
||||||
|
) # Empty token for no API key
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue