diff --git a/litellm/__init__.py b/litellm/__init__.py index 523ee36b865..69e5f78c997 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -17,6 +17,7 @@ from typing import ( TYPE_CHECKING, ) from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams +from litellm.types.integrations.datadog import DatadogInitParams from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.caching.caching import Cache, DualCache, RedisCache, InMemoryCache from litellm.caching.llm_caching_handler import LLMClientCache @@ -343,6 +344,7 @@ suppress_debug_info = False dynamodb_table_name: Optional[str] = None s3_callback_params: Optional[Dict] = None datadog_llm_observability_params: Optional[Union[DatadogLLMObsInitParams, Dict]] = None +datadog_params: Optional[Union[DatadogInitParams, Dict]] = None aws_sqs_callback_params: Optional[Dict] = None generic_logger_headers: Optional[Dict] = None default_key_generate_params: Optional[Dict] = None diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 1fa651ec71c..f5a0883bea4 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -19,7 +19,7 @@ import os import traceback import uuid from datetime import datetime as datetimeObj -from typing import Any, List, Optional, Union +from typing import Any, Dict, List, Optional, Union import httpx from httpx import Response @@ -71,6 +71,13 @@ class DataDogLogger( raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>") if os.getenv("DD_SITE", None) is None: raise Exception("DD_SITE is not set in .env, set 'DD_SITE=<>") + + ######################################################### + # Handle datadog_params set as litellm.datadog_params + ######################################################### + dict_datadog_params = self._get_datadog_params() + kwargs.update(dict_datadog_params) + self.async_client = get_async_httpx_client( llm_provider=httpxSpecialProvider.LoggingCallback ) @@ -101,6 +108,21 @@ class DataDogLogger( ) raise e + def _get_datadog_params(self) -> Dict: + """ + Get the datadog_params from litellm.datadog_params + + These are params specific to initializing the DataDogLogger e.g. turn_off_message_logging + """ + dict_datadog_params: Dict = {} + if litellm.datadog_params is not None: + if isinstance(litellm.datadog_params, DatadogInitParams): + dict_datadog_params = litellm.datadog_params.model_dump() + elif isinstance(litellm.datadog_params, Dict): + # only allow params that are of DatadogInitParams + dict_datadog_params = DatadogInitParams(**litellm.datadog_params).model_dump() + return dict_datadog_params + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): """ Async Log success events to Datadog @@ -458,6 +480,7 @@ class DataDogLogger( else: clean_metadata[key] = value + # Build the initial payload payload = { "id": id, diff --git a/litellm/types/integrations/datadog.py b/litellm/types/integrations/datadog.py index fce60b9c4b5..b7411843947 100644 --- a/litellm/types/integrations/datadog.py +++ b/litellm/types/integrations/datadog.py @@ -3,6 +3,8 @@ from typing import Optional from typing_extensions import TypedDict +from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams + DD_MAX_BATCH_SIZE = 1000 @@ -25,6 +27,14 @@ 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 DatadogInitParams(StandardCustomLoggerInitParams): + """ + Params for initializing a DataDog logger on litellm + """ + + pass + + class DatadogProxyFailureHookJsonMessage(TypedDict, total=False): exception: str error_class: str diff --git a/tests/logging_callback_tests/test_datadog.py b/tests/logging_callback_tests/test_datadog.py index 4b56684ade3..db6f193327c 100644 --- a/tests/logging_callback_tests/test_datadog.py +++ b/tests/logging_callback_tests/test_datadog.py @@ -24,7 +24,9 @@ from litellm.types.utils import ( StandardLoggingModelInformation, StandardLoggingMetadata, StandardLoggingHiddenParams, + LiteLLMCommonStrings, ) +from litellm.types.integrations.datadog import DatadogInitParams verbose_logger.setLevel(logging.DEBUG) @@ -576,4 +578,54 @@ def test_get_datadog_tags(): # Test with None request_tags standard_logging_obj["request_tags"] = None tags_none_request = DataDogLogger._get_datadog_tags(standard_logging_obj) - assert "request_tag:" not in tags_none_request \ No newline at end of file + assert "request_tag:" not in tags_none_request + + +@pytest.mark.asyncio +async def test_datadog_message_redaction(): + """ + Test that DataDog logger correctly initializes with turn_off_message_logging=True + from litellm.datadog_params + """ + try: + # Test using litellm.datadog_params pattern + litellm.datadog_params = DatadogInitParams(turn_off_message_logging=True) + + os.environ["DD_SITE"] = "https://fake.datadoghq.com" + os.environ["DD_API_KEY"] = "anything" + + # Mock the periodic flush to avoid async issues + with patch("asyncio.create_task"): + dd_logger = DataDogLogger() + + # Verify that turn_off_message_logging was set correctly from litellm.datadog_params + assert hasattr(dd_logger, 'turn_off_message_logging'), "DataDogLogger should have turn_off_message_logging attribute" + assert dd_logger.turn_off_message_logging is True, f"Expected turn_off_message_logging=True, got {dd_logger.turn_off_message_logging}" + + # Test the redaction method inherited from CustomLogger + model_call_details = { + "standard_logging_object": { + "messages": [{"role": "user", "content": "This is sensitive information that should be redacted"}], + "response": {"choices": [{"message": {"content": "This is a sensitive response that should be redacted"}}]} + } + } + + # Apply redaction using the inherited method + redacted_details = dd_logger.redact_standard_logging_payload_from_model_call_details(model_call_details) + redacted_str = LiteLLMCommonStrings.redacted_by_litellm.value + + # Verify that messages are redacted + redacted_standard_obj = redacted_details["standard_logging_object"] + assert redacted_standard_obj["messages"][0]["content"] == redacted_str, f"Messages not redacted. Got: {redacted_standard_obj['messages'][0]['content']}" + + # Verify that response is redacted + assert redacted_standard_obj["response"]["choices"][0]["message"]["content"] == redacted_str, f"Response not redacted. Got: {redacted_standard_obj['response']['choices'][0]['message']['content']}" + + print("✅ DataDog message redaction test passed") + + except Exception as e: + pytest.fail(f"Test failed with exception: {str(e)}") + finally: + # Clean up + litellm.datadog_params = None + litellm.callbacks = [] \ No newline at end of file