From c94956268cfafe2bb975d517964d902c9c63957f Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 26 Nov 2024 23:30:33 -0800 Subject: [PATCH 1/3] fix JSONDecodeError --- litellm/proxy/utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 46d554190..933db8a76 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -907,7 +907,11 @@ class ProxyLogging: parent_otel_span=parent_otel_span, token=_hash_token_if_needed(token=api_key), ) - 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( original_exception=original_exception, request_data=request_data, From c09ca785a4c519c1172b94e8b102cb324d67d194 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 26 Nov 2024 23:42:29 -0800 Subject: [PATCH 2/3] handle async_log_proxy_authentication_errors --- litellm/proxy/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 933db8a76..1fe944a72 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -891,7 +891,7 @@ class ProxyLogging: original_exception: Exception, request: Request, parent_otel_span: Optional[Any], - api_key: str, + api_key: Optional[str], ): """ Handler for Logging Authentication Errors on LiteLLM Proxy @@ -905,7 +905,7 @@ class ProxyLogging: user_api_key_dict = UserAPIKeyAuth( parent_otel_span=parent_otel_span, - token=_hash_token_if_needed(token=api_key), + token=_hash_token_if_needed(token=api_key or ""), ) try: request_data = await request.json() From cc834637070af4ae1c98c60d32d57c5b63d2cc07 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 27 Nov 2024 11:58:09 -0800 Subject: [PATCH 3/3] fix test_async_log_proxy_authentication_errors_get_request --- tests/proxy_unit_tests/test_proxy_server.py | 123 ++++++++++++++++++++ 1 file changed, 123 insertions(+) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index f0bcc7190..bde5ca050 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2195,3 +2195,126 @@ async def test_async_log_proxy_authentication_errors(): assert ( mock_logger.user_api_key_dict_logged.token is not None ) # 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