From 7befe3c78f111309d96d5114788330561d350aea Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Tue, 3 Mar 2026 20:22:59 -0800 Subject: [PATCH] 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 ++------ litellm/proxy/dd_span_tagger.py | 60 +++++++++++++++++ .../proxy/test_common_request_processing.py | 66 +++++++++++++++++-- 3 files changed, 129 insertions(+), 26 deletions(-) create mode 100644 litellm/proxy/dd_span_tagger.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 1269f58213a..e91f3fcd270 100644 --- a/litellm/proxy/common_request_processing.py +++ b/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 diff --git a/litellm/proxy/dd_span_tagger.py b/litellm/proxy/dd_span_tagger.py new file mode 100644 index 00000000000..08b7d928d0e --- /dev/null +++ b/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, + ) diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 9b905d24fd1..ba1084eafe0 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/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")