(feat) log proxy auth errors on datadog (#6931)

* add new dd type for auth errors

* add async_log_proxy_authentication_errors

* fix comment

* use async_log_proxy_authentication_errors

* test_datadog_post_call_failure_hook

* test_async_log_proxy_authentication_errors
This commit is contained in:
Ishaan Jaff 2024-11-26 20:26:57 -08:00 committed by GitHub
parent aea68cbeb6
commit 4bc06392db
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 241 additions and 9 deletions

View file

@ -344,3 +344,81 @@ async def test_datadog_logging():
await asyncio.sleep(5)
except Exception as e:
print(e)
@pytest.mark.asyncio
async def test_datadog_post_call_failure_hook():
"""Test logging proxy failures (e.g., authentication errors) to DataDog"""
try:
from litellm.integrations.datadog.datadog import DataDogLogger
os.environ["DD_SITE"] = "https://fake.datadoghq.com"
os.environ["DD_API_KEY"] = "anything"
dd_logger = DataDogLogger()
# Create a mock for the async_client's post method
mock_post = AsyncMock()
mock_post.return_value.status_code = 202
mock_post.return_value.text = "Accepted"
dd_logger.async_client.post = mock_post
# Create a test exception
class AuthenticationError(Exception):
def __init__(self):
self.status_code = 401
super().__init__("Invalid API key")
test_exception = AuthenticationError()
# Create test request data and user API key dict
request_data = {
"model": "gpt-4",
"messages": [{"role": "user", "content": "Hello"}],
}
user_api_key_dict = UserAPIKeyAuth(
api_key="fake_key", user_id="test_user", team_id="test_team"
)
# Call the failure hook
await dd_logger.async_post_call_failure_hook(
request_data=request_data,
original_exception=test_exception,
user_api_key_dict=user_api_key_dict,
)
# Wait for the periodic flush
await asyncio.sleep(6)
# Assert that the mock was called
assert mock_post.called, "HTTP request was not made"
# Get the arguments of the last call
args, kwargs = mock_post.call_args
# Verify endpoint
assert kwargs["url"].endswith("/api/v2/logs"), "Incorrect DataDog endpoint"
# Decode and verify payload
body = kwargs["data"]
with gzip.open(io.BytesIO(body), "rb") as f:
body = f.read().decode("utf-8")
body = json.loads(body)
assert len(body) == 1, "Expected one log entry"
log_entry = body[0]
assert log_entry["status"] == "error", "Expected error status"
assert log_entry["service"] == "litellm-server"
# Verify message content
message = json.loads(log_entry["message"])
print("logged message", json.dumps(message, indent=2))
assert message["exception"] == "Invalid API key"
assert message["error_class"] == "AuthenticationError"
assert message["status_code"] == 401
assert "traceback" in message
assert message["user_api_key_dict"]["api_key"] == "fake_key"
except Exception as e:
pytest.fail(f"Test failed with exception: {str(e)}")