_add_dd_apm_tags_for_litellm_call_id (#22219)

litellm/litellm_core_utils/dd_tracing.py

@@ -5,7 +5,7 @@ If the ddtrace package is not installed, the tracer will be a no-op.
 """
 
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 from litellm.secret_managers.main import get_secret_bool
 
@@ -76,3 +76,48 @@ if should_use_dd_tracer:
         tracer = NullTracer()
 else:
     tracer = NullTracer()
+
+
+def get_active_span() -> Optional[Any]:
+    """
+    Return the active Datadog span, checking current span first and then root span.
+    """
+    try:
+        current_span_fn = getattr(tracer, "current_span", None)
+        if callable(current_span_fn):
+            current_span = current_span_fn()
+            if current_span is not None:
+                return current_span
+
+        current_root_span_fn = getattr(tracer, "current_root_span", None)
+        if callable(current_root_span_fn):
+            return current_root_span_fn()
+    except Exception:
+        return None
+    return None
+
+
+def set_active_span_tag(tag_key: str, tag_value: str) -> bool:
+    """
+    Best-effort helper to set a tag on the active Datadog span.
+
+    Returns:
+        bool: True if a span tag was set, False otherwise.
+    """
+    if not tag_key or tag_value is None:
+        return False
+
+    span = get_active_span()
+    if span is None:
+        return False
+
+    try:
+        if hasattr(span, "set_tag_str"):
+            span.set_tag_str(tag_key, str(tag_value))
+            return True
+        if hasattr(span, "set_tag"):
+            span.set_tag(tag_key, str(tag_value))
+            return True
+    except Exception:
+        return False
+    return False

litellm/proxy/common_request_processing.py

@@ -29,7 +29,7 @@ from litellm.constants import (
     MAX_PAYLOAD_SIZE_FOR_DEBUG_LOG,
     STREAM_SSE_DATA_PREFIX,
 )
-from litellm.litellm_core_utils.dd_tracing import tracer
+from litellm.litellm_core_utils.dd_tracing import set_active_span_tag, tracer
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.litellm_core_utils.llm_response_utils.get_headers import (
     get_response_headers,
@@ -245,6 +245,26 @@ async def create_response(
     )
 
 
+def _add_dd_apm_tags_for_litellm_call_id(litellm_call_id: Optional[str]) -> None:
+    """
+    Attach LiteLLM call id to the active Datadog APM span.
+
+    This enables searching APM traces by LiteLLM call id returned in
+    `x-litellm-call-id`.
+    """
+    if not litellm_call_id:
+        return
+
+    try:
+        set_active_span_tag("litellm.call_id", str(litellm_call_id))
+    except Exception:
+        # Tagging is best-effort and should never impact request processing.
+        verbose_proxy_logger.debug(
+            "Failed to tag active ddtrace span with litellm.call_id",
+            exc_info=True,
+        )
+
+
 def _override_openai_response_model(
     *,
     response_obj: Any,
@@ -642,6 +662,7 @@ class ProxyBaseLLMRequestProcessing:
         self.data["litellm_call_id"] = request.headers.get(
             "x-litellm-call-id", str(uuid.uuid4())
         )
+        _add_dd_apm_tags_for_litellm_call_id(self.data.get("litellm_call_id"))
 
         ### AUTO STREAM USAGE TRACKING ###
         # If always_include_stream_usage is enabled and this is a streaming request
@@ -658,7 +679,6 @@ class ProxyBaseLLMRequestProcessing:
                 and "include_usage" not in self.data["stream_options"]
             ):
                 self.data["stream_options"]["include_usage"] = True
-
         ### CALL HOOKS ### - modify/reject incoming data before calling the model
 
         ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call

tests/test_litellm/proxy/test_common_request_processing.py

@@ -13,6 +13,7 @@ from litellm.integrations.opentelemetry import UserAPIKeyAuth
 from litellm.proxy.common_request_processing import (
     ProxyBaseLLMRequestProcessing,
     ProxyConfig,
+    _add_dd_apm_tags_for_litellm_call_id,
     _extract_error_from_sse_chunk,
     _get_cost_breakdown_from_logging_obj,
     _override_openai_response_model,
@@ -79,6 +80,20 @@ class TestProxyBaseLLMRequestProcessing:
             pytest.fail("litellm_call_id is not a valid UUID")
         assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"]
 
+    def test_add_dd_apm_tags_for_litellm_call_id_uses_dd_tracing_helper(self, monkeypatch):
+        mock_set_active_span_tag = MagicMock(return_value=True)
+        monkeypatch.setattr(
+            litellm.proxy.common_request_processing,
+            "set_active_span_tag",
+            mock_set_active_span_tag,
+        )
+
+        _add_dd_apm_tags_for_litellm_call_id("test-call-id")
+
+        mock_set_active_span_tag.assert_called_once_with(
+            "litellm.call_id", "test-call-id"
+        )
+
     @pytest.mark.asyncio
     async def test_should_apply_hierarchical_router_settings_as_override(
         self, monkeypatch