From 4bc8ca12db505e49f5821d58dd663449ee566f8f Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 3 Mar 2026 14:00:17 -0800 Subject: [PATCH] refactor(proxy): consolidate DD APM tag helpers into DDSpanTagger class --- litellm/proxy/common_request_processing.py | 96 ++++++++++--------- .../proxy/test_common_request_processing.py | 15 ++- 2 files changed, 56 insertions(+), 55 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 5f95811b7e8..7ae784920aa 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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"), ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 2cf31b2c8b5..593f0573c99 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/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", )