From 1f5c972985a788d37a441ee38fe475fe35f736ff Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 26 Nov 2024 19:35:19 -0800 Subject: [PATCH 1/6] add new dd type for auth errors --- litellm/integrations/datadog/datadog.py | 34 +++++++++++++++++++ ...odel_prices_and_context_window_backup.json | 2 +- .../integrations/datadog.py} | 10 +++++- 3 files changed, 44 insertions(+), 2 deletions(-) rename litellm/{integrations/datadog/types.py => types/integrations/datadog.py} (68%) diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 527b6f87d..b7348f50e 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -32,6 +32,8 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.integrations.datadog import * from litellm.types.services import ServiceLoggerPayload from litellm.types.utils import StandardLoggingPayload @@ -364,6 +366,38 @@ class DataDogLogger(CustomBatchLogger): """ return + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + ): + """ + Handles Proxy Errors (not-related to LLM API), ex: Authentication Errors + """ + import json + + _exception_payload = DatadogProxyFailureHookJsonMessage( + exception=str(original_exception), + error_class=str(original_exception.__class__.__name__), + status_code=getattr(original_exception, "status_code", None), + traceback=traceback.format_exc(), + user_api_key_dict=user_api_key_dict.model_dump(), + ) + + json_payload = json.dumps(_exception_payload) + verbose_logger.debug("Datadog: Logger - Logging payload = %s", json_payload) + dd_payload = DatadogPayload( + ddsource=os.getenv("DD_SOURCE", "litellm"), + ddtags="", + hostname="", + message=json_payload, + service="litellm-server", + status=DataDogStatus.ERROR, + ) + + self.log_queue.append(dd_payload) + def _create_v0_logging_payload( self, kwargs: Union[dict, Any], diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a56472f7f..b0d0e7d37 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -2032,7 +2032,6 @@ "tool_use_system_prompt_tokens": 264, "supports_assistant_prefill": true, "supports_prompt_caching": true, - "supports_pdf_input": true, "supports_response_schema": true }, "claude-3-opus-20240229": { @@ -2098,6 +2097,7 @@ "supports_vision": true, "tool_use_system_prompt_tokens": 159, "supports_assistant_prefill": true, + "supports_pdf_input": true, "supports_prompt_caching": true, "supports_response_schema": true }, diff --git a/litellm/integrations/datadog/types.py b/litellm/types/integrations/datadog.py similarity index 68% rename from litellm/integrations/datadog/types.py rename to litellm/types/integrations/datadog.py index 87aa3ce17..79d4eded4 100644 --- a/litellm/integrations/datadog/types.py +++ b/litellm/types/integrations/datadog.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import TypedDict +from typing import Optional, TypedDict class DataDogStatus(str, Enum): @@ -19,3 +19,11 @@ class DatadogPayload(TypedDict, total=False): class DD_ERRORS(Enum): DATADOG_413_ERROR = "Datadog API Error - Payload too large (batch is above 5MB uncompressed). If you want this logged either disable request/response logging or set `DD_BATCH_SIZE=50`" + + +class DatadogProxyFailureHookJsonMessage(TypedDict, total=False): + exception: str + error_class: str + status_code: Optional[int] + traceback: str + user_api_key_dict: dict From cd06f31642e7d58ade73f9ac34d63ec31ce06196 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 26 Nov 2024 19:37:33 -0800 Subject: [PATCH 2/6] add async_log_proxy_authentication_errors --- litellm/proxy/utils.py | 42 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 41 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 2a298af21..c059a8844 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -854,6 +854,20 @@ class ProxyLogging: ), ).start() + await self._run_post_call_failure_hook_custom_loggers( + original_exception=original_exception, + request_data=request_data, + user_api_key_dict=user_api_key_dict, + ) + + return + + async def _run_post_call_failure_hook_custom_loggers( + self, + original_exception: Exception, + request_data: dict, + user_api_key_dict: UserAPIKeyAuth, + ): for callback in litellm.callbacks: try: _callback: Optional[CustomLogger] = None @@ -872,7 +886,33 @@ class ProxyLogging: except Exception as e: raise e - return + async def async_log_proxy_authentication_errors( + self, + original_exception: Exception, + request: Request, + parent_otel_span: Optional[Any], + api_key: str, + ): + """ + Handler for Logging Authentication Errors on LiteLLM Proxy + Why not use post_call_failure_hook? + - `post_call_failure_hook` calls `litellm_logging_obj.async_failure_handler`. This led to the Exception being logged twice + What does this handler do? + - Logs Authentication Errors (like invalid API Key passed) to CustomLogger compatible classes + - calls CustomLogger.async_post_call_failure_hook + """ + + user_api_key_dict = UserAPIKeyAuth( + parent_otel_span=parent_otel_span, + token=_hash_token_if_needed(token=api_key), + ) + request_data = await request.json() + await self._run_post_call_failure_hook_custom_loggers( + original_exception=original_exception, + request_data=request_data, + user_api_key_dict=user_api_key_dict, + ) + pass async def post_call_success_hook( self, From 8d57ccb74eb42bf893ff74a6ed4e476574a263db Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 26 Nov 2024 19:38:02 -0800 Subject: [PATCH 3/6] fix comment --- litellm/proxy/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c059a8844..a155b23cc 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -897,6 +897,7 @@ class ProxyLogging: Handler for Logging Authentication Errors on LiteLLM Proxy Why not use post_call_failure_hook? - `post_call_failure_hook` calls `litellm_logging_obj.async_failure_handler`. This led to the Exception being logged twice + What does this handler do? - Logs Authentication Errors (like invalid API Key passed) to CustomLogger compatible classes - calls CustomLogger.async_post_call_failure_hook From 75dacd88592eac08d78d6c58a5ffd3b5363fdcc9 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 26 Nov 2024 19:39:36 -0800 Subject: [PATCH 4/6] use async_log_proxy_authentication_errors --- litellm/proxy/auth/user_api_key_auth.py | 12 +++++++----- litellm/proxy/utils.py | 2 +- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index d19215245..32f0c95db 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -1197,13 +1197,15 @@ async def user_api_key_auth( # noqa: PLR0915 extra={"requester_ip": requester_ip}, ) - # Log this exception to OTEL - if open_telemetry_logger is not None: - await open_telemetry_logger.async_post_call_failure_hook( # type: ignore + # Log this exception to OTEL, Datadog etc + asyncio.create_task( + proxy_logging_obj.async_log_proxy_authentication_errors( original_exception=e, - request_data={}, - user_api_key_dict=UserAPIKeyAuth(parent_otel_span=parent_otel_span), + request=request, + parent_otel_span=parent_otel_span, + api_key=api_key, ) + ) if isinstance(e, litellm.BudgetExceededError): raise ProxyException( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index a155b23cc..46d554190 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -899,7 +899,7 @@ class ProxyLogging: - `post_call_failure_hook` calls `litellm_logging_obj.async_failure_handler`. This led to the Exception being logged twice What does this handler do? - - Logs Authentication Errors (like invalid API Key passed) to CustomLogger compatible classes + - Logs Authentication Errors (like invalid API Key passed) to CustomLogger compatible classes (OTEL, Datadog etc) - calls CustomLogger.async_post_call_failure_hook """ From b35b269a4639abb9fda2b8b32a51e2007c26d17a Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 26 Nov 2024 19:54:47 -0800 Subject: [PATCH 5/6] test_datadog_post_call_failure_hook --- litellm/integrations/datadog/datadog.py | 1 - tests/logging_callback_tests/test_datadog.py | 78 ++++++++++++++++++++ 2 files changed, 78 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index b7348f50e..42d9a38d6 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -37,7 +37,6 @@ from litellm.types.integrations.datadog import * from litellm.types.services import ServiceLoggerPayload from litellm.types.utils import StandardLoggingPayload -from .types import DD_ERRORS, DatadogPayload, DataDogStatus from .utils import make_json_serializable DD_MAX_BATCH_SIZE = 1000 # max number of logs DD API can accept diff --git a/tests/logging_callback_tests/test_datadog.py b/tests/logging_callback_tests/test_datadog.py index a93156226..43af1b53b 100644 --- a/tests/logging_callback_tests/test_datadog.py +++ b/tests/logging_callback_tests/test_datadog.py @@ -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)}") From 541ebcf2b422503a1eb8970a9da984a713268535 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 26 Nov 2024 20:02:50 -0800 Subject: [PATCH 6/6] test_async_log_proxy_authentication_errors --- tests/proxy_unit_tests/test_proxy_server.py | 70 +++++++++++++++++++++ 1 file changed, 70 insertions(+) diff --git a/tests/proxy_unit_tests/test_proxy_server.py b/tests/proxy_unit_tests/test_proxy_server.py index 64bb67b58..f0bcc7190 100644 --- a/tests/proxy_unit_tests/test_proxy_server.py +++ b/tests/proxy_unit_tests/test_proxy_server.py @@ -2125,3 +2125,73 @@ async def test_proxy_server_prisma_setup_invalid_db(): if _old_db_url: os.environ["DATABASE_URL"] = _old_db_url + + +@pytest.mark.asyncio +async def test_async_log_proxy_authentication_errors(): + """ + Test if async_log_proxy_authentication_errors correctly logs authentication errors through custom loggers + """ + import json + from fastapi import Request + from litellm.proxy.utils import ProxyLogging + from litellm.caching import DualCache + from litellm.integrations.custom_logger import CustomLogger + + # Create a mock custom logger to verify it's called + 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 + print("logged request_data", request_data) + if isinstance(request_data, AsyncMock): + self.request_data_logged = ( + await request_data() + ) # get the actual value from AsyncMock + else: + 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("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.request_data_logged == test_data + assert mock_logger.user_api_key_dict_logged is not None + assert ( + mock_logger.user_api_key_dict_logged.token is not None + ) # token should be hashed