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

litellm/proxy/common_request_processing.py

@@ -237,57 +237,59 @@ 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.
+class DDSpanTagger:
+    """Best-effort helpers for tagging the active Datadog APM span with LiteLLM request metadata."""
 
-    This enables searching APM traces by LiteLLM call id returned in
-    `x-litellm-call-id`.
-    """
-    if not litellm_call_id:
-        return
+    @staticmethod
+    def tag_call_id(litellm_call_id: Optional[str]) -> None:
+        """
+        Attach LiteLLM call id to the active Datadog APM span.
 
-    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,
-        )
+        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.
 
-def _add_dd_apm_tags_for_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
 
-    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``.
 
-    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,
-        )
+        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,
+            )
 
 
 def _override_openai_response_model(
@@ -652,8 +654,8 @@ 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"))
-        _add_dd_apm_tags_for_request(
+        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"),
         )

tests/test_litellm/proxy/test_common_request_processing.py

@@ -11,10 +11,9 @@ import litellm
 from litellm._uuid import uuid
 from litellm.integrations.opentelemetry import UserAPIKeyAuth
 from litellm.proxy.common_request_processing import (
+    DDSpanTagger,
     ProxyBaseLLMRequestProcessing,
     ProxyConfig,
-    _add_dd_apm_tags_for_litellm_call_id,
-    _add_dd_apm_tags_for_request,
     _extract_error_from_sse_chunk,
     _get_cost_breakdown_from_logging_obj,
     _override_openai_response_model,
@@ -89,7 +88,7 @@ class TestProxyBaseLLMRequestProcessing:
             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"
@@ -1390,8 +1389,8 @@ class TestStreamingOverheadHeader:
         assert custom_headers["x-litellm-overhead-duration-ms"] == "55.3"
 
 
-class TestAddDdApmTagsForRequest:
-    """Tests for _add_dd_apm_tags_for_request - key/model DD span tagging."""
+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
@@ -1408,7 +1407,7 @@ class TestAddDdApmTagsForRequest:
         with patch(
             "litellm.proxy.common_request_processing.set_active_span_tag"
         ) as mock_set_tag:
-            _add_dd_apm_tags_for_request(
+            DDSpanTagger.tag_request(
                 user_api_key_dict=user_key,
                 requested_model="gpt-4o",
             )
@@ -1424,7 +1423,7 @@ class TestAddDdApmTagsForRequest:
         with patch(
             "litellm.proxy.common_request_processing.set_active_span_tag"
         ) as mock_set_tag:
-            _add_dd_apm_tags_for_request(
+            DDSpanTagger.tag_request(
                 user_api_key_dict=user_key,
                 requested_model=None,
             )
@@ -1438,7 +1437,7 @@ class TestAddDdApmTagsForRequest:
         with patch(
             "litellm.proxy.common_request_processing.set_active_span_tag"
         ) as mock_set_tag:
-            _add_dd_apm_tags_for_request(
+            DDSpanTagger.tag_request(
                 user_api_key_dict=user_key,
                 requested_model="claude-3-5-sonnet",
             )