Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

(fix) handle json decode errors for DD exception logging #6934

Merged
merged 3 commits into from
Nov 27, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 7 additions & 3 deletions litellm/proxy/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -891,7 +891,7 @@ async def async_log_proxy_authentication_errors(
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
Expand All @@ -905,9 +905,13 @@ async def async_log_proxy_authentication_errors(

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 ""),
)
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,
Expand Down
123 changes: 123 additions & 0 deletions tests/proxy_unit_tests/test_proxy_server.py
Original file line number Diff line number Diff line change
Expand Up @@ -2195,3 +2195,126 @@ async def async_post_call_failure_hook(
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