feat(proxy): add key_alias, key_hash, requested_model DD APM span tags (#22710)

* feat(proxy): add key_alias, key_hash, requested_model tags to DD APM spans

* refactor(proxy): consolidate DD APM tag helpers into DDSpanTagger class

* refactor(proxy): move DDSpanTagger to its own file litellm/proxy/dd_span_tagger.py

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 set_active_span_tag, tracer
+from litellm.litellm_core_utils.dd_tracing import 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,
@@ -41,6 +41,7 @@ from litellm.proxy.common_utils.callback_utils import (
     get_logging_caching_headers,
     get_remaining_tokens_and_requests_from_request_data,
 )
+from litellm.proxy.dd_span_tagger import DDSpanTagger
 from litellm.proxy.route_llm_request import route_request
 from litellm.proxy.utils import ProxyLogging
 from litellm.router import Router
@@ -245,26 +246,6 @@ 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,
@@ -662,7 +643,11 @@ 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"))
+        DDSpanTagger.tag_call_id(self.data.get("litellm_call_id"))
+        DDSpanTagger.tag_request(
+            user_api_key_dict=user_api_key_dict,
+            requested_model=self.data.get("model"),
+        )
 
         ### AUTO STREAM USAGE TRACKING ###
         # If always_include_stream_usage is enabled and this is a streaming request

litellm/proxy/dd_span_tagger.py

@@ -0,0 +1,60 @@
+from typing import Optional
+
+from litellm._logging import verbose_proxy_logger
+from litellm.litellm_core_utils.dd_tracing import set_active_span_tag
+from litellm.proxy._types import UserAPIKeyAuth
+
+
+class DDSpanTagger:
+    """Best-effort helpers for tagging the active Datadog APM span with LiteLLM request metadata."""
+
+    @staticmethod
+    def tag_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:
+            verbose_proxy_logger.debug(
+                "Failed to tag active ddtrace span with litellm.call_id",
+                exc_info=True,
+            )
+
+    @staticmethod
+    def tag_request(
+        user_api_key_dict: UserAPIKeyAuth,
+        requested_model: Optional[str],
+    ) -> None:
+        """
+        Attach key and model tags to the active Datadog APM span.
+
+        Tags set (all best-effort, skipped when value is absent):
+        - ``litellm.key_alias``      — human-readable alias for the API key
+        - ``litellm.key_hash``       — hashed API key (safe to log; never the raw secret)
+        - ``litellm.requested_model``— model name as sent by the client
+
+        Use cases:
+        - Trace all requests from a specific user/key: filter by ``litellm.key_alias`` or
+          ``litellm.key_hash``.
+        - Trace all requests for a specific model: filter by ``litellm.requested_model``.
+
+        Note: key_alias / key_hash are not available for unauthenticated (e.g. 401) requests.
+        """
+        try:
+            if user_api_key_dict.key_alias:
+                set_active_span_tag("litellm.key_alias", str(user_api_key_dict.key_alias))
+            if user_api_key_dict.token:
+                set_active_span_tag("litellm.key_hash", str(user_api_key_dict.token))
+            if requested_model:
+                set_active_span_tag("litellm.requested_model", str(requested_model))
+        except Exception:
+            verbose_proxy_logger.debug(
+                "Failed to tag active ddtrace span with key/model tags",
+                exc_info=True,
+            )

tests/test_litellm/proxy/test_common_request_processing.py

@@ -1,7 +1,7 @@
 import copy
 import datetime
 from typing import AsyncGenerator
-from unittest.mock import AsyncMock, MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from fastapi import Request, status
@@ -13,13 +13,13 @@ 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,
     _parse_event_data_for_error,
     create_response,
 )
+from litellm.proxy.dd_span_tagger import DDSpanTagger
 from litellm.proxy.utils import ProxyLogging
 
 
@@ -82,13 +82,15 @@ class TestProxyBaseLLMRequestProcessing:
 
     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)
+        import litellm.proxy.dd_span_tagger
+
         monkeypatch.setattr(
-            litellm.proxy.common_request_processing,
+            litellm.proxy.dd_span_tagger,
             "set_active_span_tag",
             mock_set_active_span_tag,
         )
 
-        _add_dd_apm_tags_for_litellm_call_id("test-call-id")
+        DDSpanTagger.tag_call_id("test-call-id")
 
         mock_set_active_span_tag.assert_called_once_with(
             "litellm.call_id", "test-call-id"
@@ -1564,3 +1566,59 @@ class TestStreamingOverheadHeader:
             "It was missing — this is the streaming overhead header regression."
         )
         assert custom_headers["x-litellm-overhead-duration-ms"] == "55.3"
+
+
+class TestDDSpanTaggerTagRequest:
+    """Tests for DDSpanTagger.tag_request - key/model DD span tagging."""
+
+    def _make_user_api_key_dict(self, key_alias=None, token=None):
+        from litellm.proxy._types import UserAPIKeyAuth
+
+        d = UserAPIKeyAuth()
+        d.key_alias = key_alias
+        d.token = token
+        return d
+
+    def test_tags_key_alias_and_model(self):
+        """key_alias and requested_model are set on the span when present."""
+        user_key = self._make_user_api_key_dict(key_alias="my-prod-key", token="hashed123")
+
+        with patch(
+            "litellm.proxy.dd_span_tagger.set_active_span_tag"
+        ) as mock_set_tag:
+            DDSpanTagger.tag_request(
+                user_api_key_dict=user_key,
+                requested_model="gpt-4o",
+            )
+
+        mock_set_tag.assert_any_call("litellm.key_alias", "my-prod-key")
+        mock_set_tag.assert_any_call("litellm.key_hash", "hashed123")
+        mock_set_tag.assert_any_call("litellm.requested_model", "gpt-4o")
+
+    def test_no_tags_when_key_absent(self):
+        """No key tags are set when key_alias and token are None (e.g. 401 path)."""
+        user_key = self._make_user_api_key_dict(key_alias=None, token=None)
+
+        with patch(
+            "litellm.proxy.dd_span_tagger.set_active_span_tag"
+        ) as mock_set_tag:
+            DDSpanTagger.tag_request(
+                user_api_key_dict=user_key,
+                requested_model=None,
+            )
+
+        mock_set_tag.assert_not_called()
+
+    def test_only_model_tagged_when_no_key_info(self):
+        """requested_model is tagged even when there's no key info."""
+        user_key = self._make_user_api_key_dict(key_alias=None, token=None)
+
+        with patch(
+            "litellm.proxy.dd_span_tagger.set_active_span_tag"
+        ) as mock_set_tag:
+            DDSpanTagger.tag_request(
+                user_api_key_dict=user_key,
+                requested_model="claude-3-5-sonnet",
+            )
+
+        mock_set_tag.assert_called_once_with("litellm.requested_model", "claude-3-5-sonnet")