From 7e23fdff9bd956f2ac1c2d892c167cf550d187e2 Mon Sep 17 00:00:00 2001 From: Shin Date: Mon, 9 Feb 2026 19:42:59 +0000 Subject: [PATCH] Fix: Datadog LLM Obs exporter APM trace correlation via monkeypatch Add a monkeypatch utility for DataDogLLMObsLogger to properly correlate LLM Obs traces with APM traces. The issue with the current implementation: 1. Uses tracer from litellm.litellm_core_utils.dd_tracing which is a NullTracer unless USE_DDTRACE=True is set 2. Only sets apm_id but doesn't set parent_id to the APM span ID 3. Falls back to generating a UUID for trace_id instead of using APM trace_id This patch imports directly from ddtrace to get the active APM context. Usage: from litellm.integrations.datadog.datadog_llm_obs_apm_patch import \ patch_datadog_llm_obs_apm_correlation patch_datadog_llm_obs_apm_correlation() import litellm litellm.success_callback = ["datadog_llm_obs"] Ref: LIT-1921 --- .../datadog/datadog_llm_obs_apm_patch.py | 146 ++++++++++++++++++ 1 file changed, 146 insertions(+) create mode 100644 litellm/integrations/datadog/datadog_llm_obs_apm_patch.py diff --git a/litellm/integrations/datadog/datadog_llm_obs_apm_patch.py b/litellm/integrations/datadog/datadog_llm_obs_apm_patch.py new file mode 100644 index 00000000000..f3d9049b579 --- /dev/null +++ b/litellm/integrations/datadog/datadog_llm_obs_apm_patch.py @@ -0,0 +1,146 @@ +""" +Monkeypatch for DataDogLLMObsLogger to properly correlate LLM Obs traces with APM traces. + +The issue is that the current implementation: +1. Uses `tracer` from `litellm.litellm_core_utils.dd_tracing` which is a NullTracer unless USE_DDTRACE=True is set +2. Only sets `apm_id` but doesn't set `parent_id` to the APM span ID +3. Falls back to generating a UUID for trace_id instead of using the APM trace_id + +This patch imports directly from ddtrace to get the active APM context. + +Usage: + # Apply this patch before initializing litellm callbacks + from litellm.integrations.datadog.datadog_llm_obs_apm_patch import patch_datadog_llm_obs_apm_correlation + patch_datadog_llm_obs_apm_correlation() + + # Then set up litellm as usual + import litellm + litellm.success_callback = ["datadog_llm_obs"] +""" + +from datetime import datetime +from typing import Any, Dict, Optional, Tuple + + +def _get_apm_trace_context() -> Tuple[Optional[str], Optional[str]]: + """ + Retrieve the current APM trace ID and span ID directly from ddtrace. + + This bypasses the litellm.litellm_core_utils.dd_tracing module which uses + a NullTracer by default. + + Returns: + Tuple of (trace_id, span_id) as strings, or (None, None) if not available. + """ + try: + # Import directly from ddtrace, not from litellm's wrapper + from ddtrace import tracer + + current_span = tracer.current_span() + if current_span is not None: + trace_id = getattr(current_span, "trace_id", None) + span_id = getattr(current_span, "span_id", None) + if trace_id is not None and span_id is not None: + return str(trace_id), str(span_id) + except ImportError: + # ddtrace not installed + pass + except Exception: + # Any other error, fail silently + pass + + return None, None + + +def patch_datadog_llm_obs_apm_correlation() -> Any: + """ + Monkey-patch DataDogLLMObsLogger to properly correlate with APM traces. + + This patch modifies create_llm_obs_payload to: + 1. Use the APM trace_id instead of generating a UUID + 2. Use the APM span_id as the parent_id + 3. Set the apm_id for trace correlation + + Returns: + The original method (for use with unpatch_datadog_llm_obs if needed). + """ + from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger + from litellm.types.integrations.datadog_llm_obs import LLMObsPayload + + # Store reference to original method + _original_create_llm_obs_payload = DataDogLLMObsLogger.create_llm_obs_payload + + def _patched_create_llm_obs_payload( + self, kwargs: Dict, start_time: datetime, end_time: datetime + ) -> LLMObsPayload: + """ + Patched version that properly correlates with APM traces. + """ + # Get APM trace context BEFORE calling original method + apm_trace_id, apm_span_id = _get_apm_trace_context() + + # Call original method to get the base payload + payload = _original_create_llm_obs_payload(self, kwargs, start_time, end_time) + + if apm_trace_id is not None and apm_span_id is not None: + # Use APM trace_id instead of the UUID that was generated + payload["trace_id"] = apm_trace_id + + # Use the APM span_id as the parent_id to establish the parent-child + # relationship between the APM span and this LLM Obs span + payload["parent_id"] = apm_span_id + + # Set the apm_id for correlation in the Datadog UI + payload["apm_id"] = apm_trace_id + + return payload + + # Apply the patch + DataDogLLMObsLogger.create_llm_obs_payload = _patched_create_llm_obs_payload + + return _original_create_llm_obs_payload # Return original in case you need to restore + + +def patch_datadog_llm_obs_apm_correlation_preserve_litellm_trace() -> Any: + """ + Alternative patch that preserves litellm's trace_id but still correlates with APM. + + Use this if you want to keep litellm's trace_id separate but still have + the LLM Obs span linked to the APM trace via apm_id and parent_id. + + Returns: + The original method (for use with unpatch_datadog_llm_obs if needed). + """ + from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger + from litellm.types.integrations.datadog_llm_obs import LLMObsPayload + + _original_create_llm_obs_payload = DataDogLLMObsLogger.create_llm_obs_payload + + def _patched_create_llm_obs_payload( + self, kwargs: Dict, start_time: datetime, end_time: datetime + ) -> LLMObsPayload: + apm_trace_id, apm_span_id = _get_apm_trace_context() + payload = _original_create_llm_obs_payload(self, kwargs, start_time, end_time) + + if apm_trace_id is not None and apm_span_id is not None: + # Keep litellm's trace_id but set parent and apm correlation + payload["parent_id"] = apm_span_id + payload["apm_id"] = apm_trace_id + + return payload + + DataDogLLMObsLogger.create_llm_obs_payload = _patched_create_llm_obs_payload + + return _original_create_llm_obs_payload + + +def unpatch_datadog_llm_obs(original_method: Any) -> None: + """ + Restore the original create_llm_obs_payload method. + + Args: + original_method: The original method returned by one of the patch functions. + """ + from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger + + DataDogLLMObsLogger.create_llm_obs_payload = original_method