diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index a33fb2dfee..2f45be63cb 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -29,6 +29,7 @@ from litellm.llms.custom_httpx.http_handler import ( from litellm.types.integrations.datadog_llm_obs import * from litellm.types.utils import ( CallTypes, + StandardLoggingGuardrailInformation, StandardLoggingPayload, StandardLoggingPayloadErrorInformation, ) @@ -422,8 +423,42 @@ class DataDogLLMObsLogger(DataDogLogger, CustomBatchLogger): "cache_key": standard_logging_payload.get("cache_key", "unknown"), "saved_cache_cost": standard_logging_payload.get("saved_cache_cost", 0), } + + ######################################################### + # Add latency metrics to metadata + ######################################################### + latency_metrics = self._get_latency_metrics(standard_logging_payload) + _metadata.update({"latency_metrics": latency_metrics}) + _standard_logging_metadata: dict = ( dict(standard_logging_payload.get("metadata", {})) or {} ) _metadata.update(_standard_logging_metadata) return _metadata + + def _get_latency_metrics(self, standard_logging_payload: StandardLoggingPayload) -> Dict: + """ + Get the latency metrics from the standard logging payload + """ + latency_metrics: DDLLMObsLatencyMetrics = DDLLMObsLatencyMetrics() + # Add latency metrics to metadata + # Time to first token (convert from seconds to milliseconds for consistency) + time_to_first_token_seconds = self._get_time_to_first_token_seconds(standard_logging_payload) + if time_to_first_token_seconds > 0: + latency_metrics["time_to_first_token_ms"] = time_to_first_token_seconds * 1000 + + # LiteLLM overhead time + hidden_params = standard_logging_payload.get("hidden_params", {}) + litellm_overhead_ms = hidden_params.get("litellm_overhead_time_ms") + if litellm_overhead_ms is not None: + latency_metrics["litellm_overhead_time_ms"] = litellm_overhead_ms + + # Guardrail overhead latency + guardrail_info: Optional[StandardLoggingGuardrailInformation] = standard_logging_payload.get("guardrail_information") + if guardrail_info is not None: + _guardrail_duration_seconds: Optional[float] = guardrail_info.get("duration") + if _guardrail_duration_seconds is not None: + # Convert from seconds to milliseconds for consistency + latency_metrics["guardrail_overhead_time_ms"] = _guardrail_duration_seconds * 1000 + + return dict(latency_metrics) \ No newline at end of file diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 15953f4229..d02671d69d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -16,6 +16,7 @@ import json import sys from typing import Any, AsyncGenerator, List, Literal, Optional, Tuple, Union +import httpx from fastapi import HTTPException import litellm @@ -284,6 +285,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): response: Optional[Union[Any, litellm.ModelResponse]] = None, request_data: Optional[dict] = None ) -> BedrockGuardrailResponse: + from datetime import datetime + start_time = datetime.now() credentials, aws_region_name = self._load_credentials() bedrock_request_data: dict = dict( self.convert_to_bedrock_format( @@ -317,6 +320,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): data=prepared_request.body, # type: ignore headers=prepared_request.headers, # type: ignore ) + ######################################################### + # Add guardrail information to request trace + ######################################################### + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response=response.json(), + request_data=request_data or {}, + guardrail_status=self._get_bedrock_guardrail_response_status(response=response), + start_time=start_time.timestamp(), + end_time=datetime.now().timestamp(), + duration=(datetime.now() - start_time).total_seconds(), + ) + ######################################################### if response.status_code == 200: # check if the response was flagged _json_response = response.json() @@ -338,6 +353,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return bedrock_guardrail_response + def _get_bedrock_guardrail_response_status(self, response: httpx.Response) -> Literal["success", "failure"]: + """ + Get the status of the bedrock guardrail response. + """ + if response.status_code == 200: + return "success" + return "failure" def _get_http_exception_for_blocked_guardrail(self, response: BedrockGuardrailResponse) -> HTTPException: """ @@ -501,10 +523,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): add_guardrail_to_applied_guardrails_header( request_data=data, guardrail_name=self.guardrail_name ) - return data - @log_guardrail_information async def async_moderation_hook( self, data: dict, @@ -561,7 +581,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return data - @log_guardrail_information async def async_post_call_success_hook( self, data: dict, diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index ef15e3f3ba..28ba8cd093 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -6,4 +6,12 @@ model_list: litellm_params: model: bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0 litellm_settings: - callbacks: ["datadog_llm_observability"] \ No newline at end of file + callbacks: ["datadog_llm_observability"] + +guardrails: + - guardrail_name: "bedrock-pre-guard" + litellm_params: + guardrail: bedrock # supported values: "aporia", "bedrock", "lakera" + mode: "during_call" + guardrailIdentifier: ff6ujrregl1q + guardrailVersion: "DRAFT" \ 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 008101cec1..82fb4fe388 100644 --- a/litellm/types/integrations/datadog_llm_obs.py +++ b/litellm/types/integrations/datadog_llm_obs.py @@ -71,4 +71,10 @@ class DatadogLLMObsInitParams(StandardCustomLoggerInitParams): """ Params for initializing a DatadogLLMObs logger on litellm """ - pass \ No newline at end of file + pass + + +class DDLLMObsLatencyMetrics(TypedDict, total=False): + time_to_first_token_ms: float + litellm_overhead_time_ms: float + guardrail_overhead_time_ms: float \ No newline at end of file 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 c53ea270f4..25263e31f9 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_observability.py @@ -20,6 +20,7 @@ from litellm.types.integrations.datadog_llm_obs import ( LLMObsPayload, ) from litellm.types.utils import ( + StandardLoggingGuardrailInformation, StandardLoggingHiddenParams, StandardLoggingMetadata, StandardLoggingModelInformation, @@ -354,7 +355,7 @@ async def test_dd_llms_obs_redaction(mock_env_vars): 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() + dd_llms_obs_logger = TestDataDogLLMObsLoggerForRedaction() test_s3_logger = TestS3Logger() litellm.callbacks = [ dd_llms_obs_logger, @@ -424,3 +425,145 @@ async def test_create_llm_obs_payload(mock_env_vars): assert payload["metrics"]["input_tokens"] == 10 assert payload["metrics"]["output_tokens"] == 20 assert payload["metrics"]["total_tokens"] == 30 + + +def create_standard_logging_payload_with_latency_metrics() -> StandardLoggingPayload: + """Create a StandardLoggingPayload object with latency metrics for testing""" + guardrail_info = StandardLoggingGuardrailInformation( + guardrail_name="test_guardrail", + guardrail_status="success", + start_time=1234567890.0, + end_time=1234567890.5, + duration=0.5, # 500ms + ) + + hidden_params = StandardLoggingHiddenParams( + model_id="model-123", + cache_key="test-cache-key", + api_base="https://api.openai.com", + response_cost="0.05", + litellm_overhead_time_ms=150.0, # 150ms + additional_headers=None, + ) + + return StandardLoggingPayload( + id="test-request-id-latency", + call_type="completion", + response_cost=0.05, + response_cost_failure_debug_info=None, + status="success", + total_tokens=30, + prompt_tokens=10, + completion_tokens=20, + startTime=1234567890.0, + endTime=1234567892.0, + completionStartTime=1234567890.8, # 800ms after start + response_time=2.0, + model_map_information=StandardLoggingModelInformation( + model_map_key="gpt-4", model_map_value=None + ), + model="gpt-4", + model_id="model-123", + model_group="openai-gpt", + api_base="https://api.openai.com", + metadata=StandardLoggingMetadata( + user_api_key_hash="test_hash", + user_api_key_org_id=None, + user_api_key_alias="test_alias", + user_api_key_team_id="test_team", + user_api_key_user_id="test_user", + user_api_key_team_alias="test_team_alias", + spend_logs_metadata=None, + requester_ip_address="127.0.0.1", + requester_metadata=None, + ), + cache_hit=False, + cache_key=None, + saved_cache_cost=0.0, + request_tags=[], + end_user=None, + requester_ip_address="127.0.0.1", + messages=[{"role": "user", "content": "Hello, world!"}], + response={"choices": [{"message": {"content": "Hi there!"}}]}, + error_str=None, + error_information=None, + model_parameters={"stream": True}, + hidden_params=hidden_params, + guardrail_information=guardrail_info, + trace_id="test-trace-id-latency", + custom_llm_provider="openai", + ) + + +def test_latency_metrics_in_metadata(mock_env_vars): + """Test that time to first token, litellm overhead, and guardrail overhead are included in metadata""" + with patch('litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client'), \ + patch('asyncio.create_task'): + logger = DataDogLLMObsLogger() + + standard_payload = create_standard_logging_payload_with_latency_metrics() + + kwargs = { + "standard_logging_object": standard_payload, + "litellm_params": {"metadata": {}} + } + + start_time = datetime.now() + end_time = datetime.now() + + # Test the metadata generation directly + metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) + latency_metadata = metadata.get("latency_metrics", {}) + + # Verify time to first token is included (800ms) + assert "time_to_first_token_ms" in latency_metadata + assert abs(latency_metadata["time_to_first_token_ms"] - 800.0) < 0.001 # 0.8 seconds * 1000 with tolerance for floating-point precision + + # Verify litellm overhead is included (150ms) + assert "litellm_overhead_time_ms" in latency_metadata + assert latency_metadata["litellm_overhead_time_ms"] == 150.0 + + # Verify guardrail overhead is included (500ms) + assert "guardrail_overhead_time_ms" in latency_metadata + assert latency_metadata["guardrail_overhead_time_ms"] == 500.0 # 0.5 seconds * 1000 + + # Verify these metrics are also included in the full payload + payload = logger.create_llm_obs_payload(kwargs, start_time, end_time) + payload_metadata_latency = payload["meta"]["metadata"]["latency_metrics"] + + assert abs(payload_metadata_latency["time_to_first_token_ms"] - 800.0) < 0.001 + assert payload_metadata_latency["litellm_overhead_time_ms"] == 150.0 + assert payload_metadata_latency["guardrail_overhead_time_ms"] == 500.0 + + +def test_latency_metrics_edge_cases(mock_env_vars): + """Test latency metrics with edge cases (missing fields, zero values, etc.)""" + with patch('litellm.integrations.datadog.datadog_llm_obs.get_async_httpx_client'), \ + patch('asyncio.create_task'): + logger = DataDogLLMObsLogger() + + # Test case 1: No latency metrics present + standard_payload = create_standard_logging_payload_with_cache() + metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) + + # Should not have latency fields if data is missing/zero + assert "time_to_first_token_ms" not in metadata # Will be 0, so not included + assert "litellm_overhead_time_ms" not in metadata # Not present in hidden_params + assert "guardrail_overhead_time_ms" not in metadata # No guardrail_information + + # Test case 2: Zero time to first token should not be included + standard_payload = create_standard_logging_payload_with_cache() + standard_payload["startTime"] = 1000.0 + standard_payload["completionStartTime"] = 1000.0 # Same time = 0 difference + metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) + assert "time_to_first_token_ms" not in metadata + + # Test case 3: Missing guardrail duration should not crash + standard_payload = create_standard_logging_payload_with_cache() + standard_payload["guardrail_information"] = StandardLoggingGuardrailInformation( + guardrail_name="test", + guardrail_status="success", + # duration is missing + ) + metadata = logger._get_dd_llm_obs_payload_metadata(standard_payload) + assert "guardrail_overhead_time_ms" not in metadata