diff --git a/.circleci/config.yml b/.circleci/config.yml index 5e441ade02..46d2ea2c6e 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1380,7 +1380,6 @@ jobs: - run: python ./tests/code_coverage_tests/recursive_detector.py - run: python ./tests/code_coverage_tests/test_router_strategy_async.py - run: python ./tests/code_coverage_tests/litellm_logging_code_coverage.py - # - run: python ./tests/code_coverage_tests/bedrock_pricing.py - run: python ./tests/documentation_tests/test_env_keys.py - run: python ./tests/documentation_tests/test_router_settings.py - run: python ./tests/documentation_tests/test_api_docs.py diff --git a/docs/my-website/docs/observability/datadog.md b/docs/my-website/docs/observability/datadog.md index 7cd98d7269..08ebf8b28c 100644 --- a/docs/my-website/docs/observability/datadog.md +++ b/docs/my-website/docs/observability/datadog.md @@ -9,8 +9,14 @@ LiteLLM Supports logging to the following Datdog Integrations: - `datadog_llm_observability` [Datadog LLM Observability](https://www.datadoghq.com/product/llm-observability/) - `ddtrace-run` [Datadog Tracing](#datadog-tracing) - - +## Datadog Logs + +| Feature | Details | +|---------|---------| +| **What is logged** | [StandardLoggingPayload](../proxy/logging_spec) | +| **Events** | Success + Failure | +| **Product Link** | [Datadog Logs](https://docs.datadoghq.com/logs/) | + We will use the `--config` to set `litellm.callbacks = ["datadog"]` this will log all successful LLM calls to DataDog @@ -26,8 +32,16 @@ litellm_settings: service_callback: ["datadog"] # logs redis, postgres failures on datadog ``` - - + +## Datadog LLM Observability + +**Overview** + +| Feature | Details | +|---------|---------| +| **What is logged** | [StandardLoggingPayload](../proxy/logging_spec) | +| **Events** | Success + Failure | +| **Product Link** | [Datadog LLM Observability](https://www.datadoghq.com/product/llm-observability/) | ```yaml model_list: @@ -38,8 +52,7 @@ litellm_settings: callbacks: ["datadog_llm_observability"] # logs llm success logs on datadog ``` - - + **Step 2**: Set Required env variables for datadog @@ -80,7 +93,53 @@ Expected output on Datadog -#### Datadog Tracing +### Redacting Messages and Responses + +This section covers how to redact sensitive data from messages and responses in the logged payload on Datadog LLM Observability. + + +When redaction is enabled, the actual message content and response text will be excluded from Datadog logs while preserving metadata like token counts, latency, and model information. + +**Step 1**: Configure redaction in your `config.yaml` + +```yaml showLineNumbers title="config.yaml" +model_list: + - model_name: gpt-3.5-turbo + litellm_params: + model: gpt-3.5-turbo +litellm_settings: + callbacks: ["datadog_llm_observability"] # logs llm success logs on datadog + + # Params to apply only for "datadog_llm_observability" callback + datadog_llm_observability_params: + turn_off_message_logging: true # redacts input messages and output responses +``` + +**Step 2**: Send a chat completion request + +```shell +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Content-Type: application/json' \ + --data '{ + "model": "gpt-3.5-turbo", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ] +}' +``` + +**Step 3**: Verify redaction in Datadog LLM Observability + +On the Datadog LLM Observability page, you should see that both input messages and output responses are redacted, while metadata (token counts, timing, model info) remains visible. + + + + + +### Datadog Tracing Use `ddtrace-run` to enable [Datadog Tracing](https://ddtrace.readthedocs.io/en/stable/installation_quickstart.html) on litellm proxy @@ -104,7 +163,7 @@ docker run \ --config /app/config.yaml --detailed_debug ``` -### Set DD variables (`DD_SERVICE` etc) +## Set DD variables (`DD_SERVICE` etc) LiteLLM supports customizing the following Datadog environment variables diff --git a/docs/my-website/docs/proxy/logging.md b/docs/my-website/docs/proxy/logging.md index e956b0970d..5d3f841722 100644 --- a/docs/my-website/docs/proxy/logging.md +++ b/docs/my-website/docs/proxy/logging.md @@ -1539,6 +1539,9 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \ ## [Datadog](../observability/datadog) +👉 Go here for using [Datadog LLM Observability](../observability/datadog) with LiteLLM Proxy + + ## Lunary #### Step1: Install dependencies and set your environment variables Install the dependencies diff --git a/docs/my-website/img/dd_llm_obs.png b/docs/my-website/img/dd_llm_obs.png new file mode 100644 index 0000000000..be7c7c7717 Binary files /dev/null and b/docs/my-website/img/dd_llm_obs.png differ diff --git a/litellm/__init__.py b/litellm/__init__.py index 412f552e9c..68d94aabb1 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -5,7 +5,8 @@ warnings.filterwarnings("ignore", message=".*conflict with protected namespace.* ### INIT VARIABLES #################### import threading import os -from typing import Callable, List, Optional, Dict, Union, Any, Literal, get_args +from typing import Callable, List, Optional, Dict, Union, Any, Literal, get_args, TYPE_CHECKING +from litellm.types.integrations.datadog_llm_obs import DatadogLLMObsInitParams 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 @@ -297,6 +298,7 @@ model_cost_map_url: str = "https://raw.githubusercontent.com/BerriAI/litellm/mai 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 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/custom_logger.py b/litellm/integrations/custom_logger.py index 6755990bdf..cdc1200547 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -34,11 +34,11 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth from litellm.types.mcp import ( + MCPDuringCallRequestObject, + MCPDuringCallResponseObject, MCPPostCallResponseObject, MCPPreCallRequestObject, MCPPreCallResponseObject, - MCPDuringCallRequestObject, - MCPDuringCallResponseObject, ) from litellm.types.router import PreRoutingHookResponse @@ -57,8 +57,21 @@ else: class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class # Class variables or attributes - def __init__(self, message_logging: bool = True, **kwargs) -> None: + def __init__( + self, + turn_off_message_logging: bool = False, + + # deprecated param, use `turn_off_message_logging` instead + message_logging: bool = True, + **kwargs + ) -> None: + """ + Args: + turn_off_message_logging: bool - if True, the message logging will be turned off. Message and response will be redacted from StandardLoggingPayload. + message_logging: bool - deprecated param, use `turn_off_message_logging` instead + """ self.message_logging = message_logging + self.turn_off_message_logging = turn_off_message_logging pass def log_pre_api_call(self, model, messages, kwargs): @@ -534,3 +547,49 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac if LITELLM_METADATA_FIELD in request_kwargs: return LITELLM_METADATA_FIELD return OLD_LITELLM_METADATA_FIELD + + def redact_standard_logging_payload_from_model_call_details( + self, model_call_details: Dict + ) -> Dict: + """ + Only redacts messages and responses when self.turn_off_message_logging is True + + + By default, self.turn_off_message_logging is False and this does nothing. + + Return a redacted deepcopy of the provided logging payload. + + This is useful for logging payloads that contain sensitive information. + """ + from copy import copy + + from litellm import Choices, Message, ModelResponse + from litellm.types.utils import LiteLLMCommonStrings + turn_off_message_logging: bool = getattr(self, "turn_off_message_logging", False) + + if turn_off_message_logging is False: + return model_call_details + + # Only make a shallow copy of the top-level dict to avoid deepcopy issues + # with complex objects like AuthenticationError that may be present + model_call_details_copy = copy(model_call_details) + redacted_str = LiteLLMCommonStrings.redacted_by_litellm.value + standard_logging_object = model_call_details.get("standard_logging_object") + if standard_logging_object is None: + return model_call_details_copy + + # Make a copy of just the standard_logging_object to avoid modifying the original + standard_logging_object_copy = copy(standard_logging_object) + + if standard_logging_object_copy.get("messages") is not None: + standard_logging_object_copy["messages"] = [Message(content=redacted_str).model_dump()] + + if standard_logging_object_copy.get("response") is not None: + model_response = ModelResponse( + choices=[Choices(message=Message(content=redacted_str))] + ) + model_response_dict = model_response.model_dump() + standard_logging_object_copy["response"] = model_response_dict + + model_call_details_copy["standard_logging_object"] = standard_logging_object_copy + return model_call_details_copy diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 8cee33968b..2577ed3ddf 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -58,18 +58,40 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger): asyncio.create_task(self.periodic_flush()) self.flush_lock = asyncio.Lock() self.log_queue: List[LLMObsPayload] = [] + + ######################################################### + # Handle datadog_llm_observability_params set as litellm.datadog_llm_observability_params + ######################################################### + dict_datadog_llm_obs_params = self._get_datadog_llm_obs_params() + kwargs.update(dict_datadog_llm_obs_params) CustomBatchLogger.__init__(self, **kwargs, flush_lock=self.flush_lock) except Exception as e: verbose_logger.exception(f"DataDogLLMObs: Error initializing - {str(e)}") raise e + def _get_datadog_llm_obs_params(self) -> Dict: + """ + Get the datadog_llm_observability_params from litellm.datadog_llm_observability_params + + These are params specific to initializing the DataDogLLMObsLogger e.g. turn_off_message_logging + """ + dict_datadog_llm_obs_params: Dict = {} + if litellm.datadog_llm_observability_params is not None: + if isinstance(litellm.datadog_llm_observability_params, DatadogLLMObsInitParams): + dict_datadog_llm_obs_params = litellm.datadog_llm_observability_params.model_dump() + elif isinstance(litellm.datadog_llm_observability_params, Dict): + # only allow params that are of DatadogLLMObsInitParams + dict_datadog_llm_obs_params = DatadogLLMObsInitParams(**litellm.datadog_llm_observability_params).model_dump() + return dict_datadog_llm_obs_params + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: verbose_logger.debug( f"DataDogLLMObs: Logging success event for model {kwargs.get('model', 'unknown')}" ) payload = self.create_llm_obs_payload( - kwargs, response_obj, start_time, end_time + kwargs, start_time, end_time ) verbose_logger.debug(f"DataDogLLMObs: Payload: {payload}") self.log_queue.append(payload) @@ -128,7 +150,7 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger): verbose_logger.exception(f"DataDogLLMObs: Error sending batch - {str(e)}") def create_llm_obs_payload( - self, kwargs: Dict, response_obj: Any, start_time: datetime, end_time: datetime + self, kwargs: Dict, start_time: datetime, end_time: datetime ) -> LLMObsPayload: standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( "standard_logging_object" @@ -138,6 +160,7 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger): messages = standard_logging_payload["messages"] messages = self._ensure_string_content(messages=messages) + response_obj = standard_logging_payload.get("response") metadata = kwargs.get("litellm_params", {}).get("metadata", {}) @@ -146,7 +169,10 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger): messages ) ) - output_meta = OutputMeta(messages=self._get_response_messages(response_obj)) + output_meta = OutputMeta(messages=self._get_response_messages( + response_obj=response_obj, + call_type=standard_logging_payload.get("call_type") + )) meta = Meta( kind=self._get_datadog_span_kind(standard_logging_payload.get("call_type")), @@ -198,14 +224,16 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger): return 0.0 - def _get_response_messages(self, response_obj: Any) -> List[Any]: + def _get_response_messages( + self, response_obj: Any, call_type: Optional[str] + ) -> List[Any]: """ Get the messages from the response object for now this handles logging /chat/completions responses """ - if isinstance(response_obj, litellm.ModelResponse): - return [response_obj["choices"][0]["message"].json()] + if call_type in [CallTypes.completion.value, CallTypes.acompletion.value]: + return [response_obj["choices"][0]["message"]] return [] def _get_datadog_span_kind(self, call_type: Optional[str]) -> Literal["llm", "tool", "task", "embedding", "retrieval"]: diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index defdaa0b01..12af18804d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -79,9 +79,7 @@ from litellm.types.llms.openai import ( ResponseCompletedEvent, ResponsesAPIResponse, ) -from litellm.types.mcp import ( - MCPPostCallResponseObject, -) +from litellm.types.mcp import MCPPostCallResponseObject from litellm.types.rerank import RerankResponse from litellm.types.router import CustomPricingLiteLLMParams from litellm.types.utils import ( @@ -169,10 +167,10 @@ try: from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import ( SMTPEmailLogger, ) + from litellm_enterprise.integrations.prometheus import PrometheusLogger from litellm_enterprise.litellm_core_utils.litellm_logging import ( StandardLoggingPayloadSetup as EnterpriseStandardLoggingPayloadSetup, ) - from litellm_enterprise.integrations.prometheus import PrometheusLogger EnterpriseStandardLoggingPayloadSetupVAR: Optional[ @@ -947,7 +945,8 @@ class Logging(LiteLLMLoggingBaseClass): if additional_args.get("request_str", None) is not None: # print the sagemaker / bedrock client request curl_command = "\nRequest Sent from LiteLLM:\n" - curl_command += additional_args.get("request_str", None) + request_str = additional_args.get("request_str", "") + curl_command += request_str elif api_base == "": curl_command = str(self.model_call_details) return curl_command @@ -2267,15 +2266,23 @@ class Logging(LiteLLMLoggingBaseClass): start_time=start_time, end_time=end_time, ) + if isinstance(callback, CustomLogger): # custom logger class + model_call_details: Dict = self.model_call_details + ################################## + # call redaction hook for custom logger + model_call_details = callback.redact_standard_logging_payload_from_model_call_details( + model_call_details=model_call_details + ) + ################################## if self.stream is True: if ( "async_complete_streaming_response" - in self.model_call_details + in model_call_details ): await callback.async_log_success_event( - kwargs=self.model_call_details, - response_obj=self.model_call_details[ + kwargs=model_call_details, + response_obj=model_call_details[ "async_complete_streaming_response" ], start_time=start_time, @@ -2283,14 +2290,14 @@ class Logging(LiteLLMLoggingBaseClass): ) else: await callback.async_log_stream_event( # [TODO]: move this to being an async log stream event function - kwargs=self.model_call_details, + kwargs=model_call_details, response_obj=result, start_time=start_time, end_time=end_time, ) else: await callback.async_log_success_event( - kwargs=self.model_call_details, + kwargs=model_call_details, response_obj=result, start_time=start_time, end_time=end_time, @@ -3211,13 +3218,14 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 _in_memory_loggers.append(_literalai_logger) return _literalai_logger # type: ignore elif logging_integration == "prometheus": - for callback in _in_memory_loggers: - if isinstance(callback, PrometheusLogger): - return callback # type: ignore + if PrometheusLogger is not None: + for callback in _in_memory_loggers: + if isinstance(callback, PrometheusLogger): + return callback # type: ignore - _prometheus_logger = PrometheusLogger() - _in_memory_loggers.append(_prometheus_logger) - return _prometheus_logger # type: ignore + _prometheus_logger = PrometheusLogger() + _in_memory_loggers.append(_prometheus_logger) + return _prometheus_logger # type: ignore elif logging_integration == "datadog": for callback in _in_memory_loggers: if isinstance(callback, DataDogLogger): @@ -3533,6 +3541,7 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915 f"[Non-Blocking Error] Error initializing custom logger: {e}" ) return None + return None def get_custom_logger_compatible_class( # noqa: PLR0915 @@ -3574,9 +3583,10 @@ def get_custom_logger_compatible_class( # noqa: PLR0915 if isinstance(callback, LiteralAILogger): return callback elif logging_integration == "prometheus": - for callback in _in_memory_loggers: - if isinstance(callback, PrometheusLogger): - return callback + if PrometheusLogger is not None: + for callback in _in_memory_loggers: + if isinstance(callback, PrometheusLogger): + return callback elif logging_integration == "datadog": for callback in _in_memory_loggers: if isinstance(callback, DataDogLogger): diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 4b71df0117..724a46e609 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -11,10 +11,6 @@ def _get_salt_key(): salt_key = os.getenv("LITELLM_SALT_KEY", None) if salt_key is None: - verbose_proxy_logger.debug( - "LITELLM_SALT_KEY is None using master_key to encrypt/decrypt secrets stored in DB" - ) - salt_key = master_key return salt_key diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 68d7c6786f..7aababa79d 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -3,5 +3,4 @@ model_list: litellm_params: model: vertex_ai/* -litellm_settings: - callbacks: ["datadog_llm_observability"] + diff --git a/litellm/types/integrations/custom_logger.py b/litellm/types/integrations/custom_logger.py new file mode 100644 index 0000000000..96952404b7 --- /dev/null +++ b/litellm/types/integrations/custom_logger.py @@ -0,0 +1,10 @@ +from typing import Optional + +from pydantic import BaseModel + + +class StandardCustomLoggerInitParams(BaseModel): + """ + Params for initializing a CustomLogger. + """ + turn_off_message_logging: Optional[bool] = False \ No newline at end of file diff --git a/litellm/types/integrations/datadog_llm_obs.py b/litellm/types/integrations/datadog_llm_obs.py index b0336f8a42..25685db483 100644 --- a/litellm/types/integrations/datadog_llm_obs.py +++ b/litellm/types/integrations/datadog_llm_obs.py @@ -3,9 +3,10 @@ Payloads for Datadog LLM Observability Service (LLMObs) API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=example#api-standards """ - from typing import Any, Dict, List, Literal, Optional, TypedDict +from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams + class InputMeta(TypedDict): messages: List[ @@ -55,3 +56,10 @@ class DDSpanAttributes(TypedDict): class DDIntakePayload(TypedDict): type: str attributes: DDSpanAttributes + + +class DatadogLLMObsInitParams(StandardCustomLoggerInitParams): + """ + Params for initializing a DatadogLLMObs logger on litellm + """ + pass \ No newline at end of file diff --git a/tests/logging_callback_tests/test_datadog_llm_obs.py b/tests/logging_callback_tests/test_datadog_llm_obs.py index 0fc5506601..ebe4543c5a 100644 --- a/tests/logging_callback_tests/test_datadog_llm_obs.py +++ b/tests/logging_callback_tests/test_datadog_llm_obs.py @@ -102,35 +102,3 @@ async def test_datadog_llm_obs_logging(): await asyncio.sleep(6) - -@pytest.mark.asyncio -async def test_create_llm_obs_payload(): - datadog_llm_obs_logger = DataDogLLMObsLogger() - standard_logging_payload = create_standard_logging_payload() - payload = datadog_llm_obs_logger.create_llm_obs_payload( - kwargs={ - "model": "gpt-4", - "messages": [{"role": "user", "content": "Hello"}], - "standard_logging_object": standard_logging_payload, - }, - response_obj=litellm.ModelResponse( - id="test_id", - choices=[{"message": {"content": "Hi there!"}}], - created=12, - model="gpt-4", - ), - start_time=datetime.now(), - end_time=datetime.now() + timedelta(seconds=1), - ) - - print("dd created payload", payload) - - assert payload["name"] == "litellm_llm_call" - assert payload["meta"]["kind"] == "llm" - assert payload["meta"]["input"]["messages"] == [ - {"role": "user", "content": "Hello, world!"} - ] - assert payload["meta"]["output"]["messages"][0]["content"] == "Hi there!" - assert payload["metrics"]["input_tokens"] == 20 - assert payload["metrics"]["output_tokens"] == 10 - assert payload["metrics"]["total_tokens"] == 30 diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py index 0842918c89..18d3efdddd 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py @@ -1,8 +1,9 @@ +import asyncio import json import os import sys import uuid -from datetime import datetime +from datetime import datetime, timedelta from typing import Dict, Optional from unittest.mock import MagicMock, Mock, patch @@ -10,9 +11,14 @@ import pytest # Adds the grandparent directory to sys.path to allow importing project modules sys.path.insert(0, os.path.abspath("../..")) - +import litellm +from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger -from litellm.types.integrations.datadog_llm_obs import LLMMetrics, LLMObsPayload +from litellm.types.integrations.datadog_llm_obs import ( + DatadogLLMObsInitParams, + LLMMetrics, + LLMObsPayload, +) from litellm.types.utils import ( StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -212,3 +218,100 @@ class TestDataDogLLMObsLogger: assert logger._get_datadog_span_kind(None) == "llm" + +class TestDataDogLLMObsLogger(DataDogLLMObsLogger): + """Test suite for DataDog LLM Observability Logger""" + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_standard_logging_payload = kwargs.get("standard_logging_object") + + +class TestS3Logger(CustomLogger): + """Test suite for S3 Logger""" + def __init__(self, **kwargs): + super().__init__(**kwargs) + self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.logged_standard_logging_payload = kwargs.get("standard_logging_object") + + +@pytest.mark.asyncio +async def test_dd_llms_obs_redaction(mock_env_vars): + # init DD with turn_off_message_logging=True + litellm._turn_on_debug() + from litellm.types.utils import LiteLLMCommonStrings + litellm.datadog_llm_observability_params = DatadogLLMObsInitParams(turn_off_message_logging=True) + dd_llms_obs_logger = TestDataDogLLMObsLogger() + test_s3_logger = TestS3Logger() + litellm.callbacks = [ + dd_llms_obs_logger, + test_s3_logger + ] + + # call litellm + await litellm.acompletion( + model="gpt-4o", + mock_response="Hi there!", + messages=[{"role": "user", "content": "Hello, world!"}] + ) + + # sleep 1 second for logging to complete + await asyncio.sleep(1) + + ################# + # test validation + # 1. both loggers logged a standard_logging_payload + # 2. DD LLM Obs standard_logging_payload has messages and response redacted + # 3. S3 standard_logging_payload does not have messages and response redacted + + assert dd_llms_obs_logger.logged_standard_logging_payload is not None + assert test_s3_logger.logged_standard_logging_payload is not None + + print("logged DD LLM Obs payload", json.dumps(dd_llms_obs_logger.logged_standard_logging_payload, indent=4, default=str)) + print("\n\nlogged S3 payload", json.dumps(test_s3_logger.logged_standard_logging_payload, indent=4, default=str)) + + assert dd_llms_obs_logger.logged_standard_logging_payload["messages"][0]["content"] == LiteLLMCommonStrings.redacted_by_litellm.value + assert dd_llms_obs_logger.logged_standard_logging_payload["response"]["choices"][0]["message"]["content"] == LiteLLMCommonStrings.redacted_by_litellm.value + + assert test_s3_logger.logged_standard_logging_payload["messages"] == [{"role": "user", "content": "Hello, world!"}] + assert test_s3_logger.logged_standard_logging_payload["response"]["choices"][0]["message"]["content"] == "Hi there!" + + +@pytest.fixture +def mock_env_vars(): + """Mock environment variables for DataDog""" + with patch.dict(os.environ, { + "DD_API_KEY": "test_api_key", + "DD_SITE": "us5.datadoghq.com" + }): + yield + +@pytest.mark.asyncio +async def test_create_llm_obs_payload(mock_env_vars): + datadog_llm_obs_logger = DataDogLLMObsLogger() + standard_logging_payload = create_standard_logging_payload_with_cache() + payload = datadog_llm_obs_logger.create_llm_obs_payload( + kwargs={ + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "standard_logging_object": standard_logging_payload, + }, + start_time=datetime.now(), + end_time=datetime.now() + timedelta(seconds=1), + ) + + print("dd created payload", payload) + + assert payload["name"] == "litellm_llm_call" + assert payload["meta"]["kind"] == "llm" + assert payload["meta"]["input"]["messages"] == [ + {"role": "user", "content": "Hello, world!"} + ] + assert payload["meta"]["output"]["messages"][0]["content"] == "Hi there!" + assert payload["metrics"]["input_tokens"] == 10 + assert payload["metrics"]["output_tokens"] == 20 + assert payload["metrics"]["total_tokens"] == 30