[Feat] DD LLM Observability - Add time to first token, litellm overhead, guardrail overhead latency metrics (#13734)
* fixes for DDLLMObsLatencyMetrics * use _get_latency_metrics * DD LLM Obs - track latency metrics * fixes for bedrock guardrails * DD unit tests * test DD
This commit is contained in:
parent
ef08e18c66
commit
ba1d2e8749
@ -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)
|
||||
@ -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,
|
||||
|
||||
@ -6,4 +6,12 @@ model_list:
|
||||
litellm_params:
|
||||
model: bedrock/us.anthropic.claude-sonnet-4-20250514-v1:0
|
||||
litellm_settings:
|
||||
callbacks: ["datadog_llm_observability"]
|
||||
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"
|
||||
@ -71,4 +71,10 @@ class DatadogLLMObsInitParams(StandardCustomLoggerInitParams):
|
||||
"""
|
||||
Params for initializing a DatadogLLMObs logger on litellm
|
||||
"""
|
||||
pass
|
||||
pass
|
||||
|
||||
|
||||
class DDLLMObsLatencyMetrics(TypedDict, total=False):
|
||||
time_to_first_token_ms: float
|
||||
litellm_overhead_time_ms: float
|
||||
guardrail_overhead_time_ms: float
|
||||
@ -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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user