diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 2a967743a8a..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,61 +246,6 @@ async def create_response( ) -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, - ) - - def _override_openai_response_model( *, response_obj: Any, 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 d64fff27e87..ba1084eafe0 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -11,7 +11,6 @@ import litellm from litellm._uuid import uuid from litellm.integrations.opentelemetry import UserAPIKeyAuth from litellm.proxy.common_request_processing import ( - DDSpanTagger, ProxyBaseLLMRequestProcessing, ProxyConfig, _extract_error_from_sse_chunk, @@ -20,6 +19,7 @@ from litellm.proxy.common_request_processing import ( _parse_event_data_for_error, create_response, ) +from litellm.proxy.dd_span_tagger import DDSpanTagger from litellm.proxy.utils import ProxyLogging @@ -82,8 +82,10 @@ 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, ) @@ -1582,7 +1584,7 @@ class TestDDSpanTaggerTagRequest: user_key = self._make_user_api_key_dict(key_alias="my-prod-key", token="hashed123") with patch( - "litellm.proxy.common_request_processing.set_active_span_tag" + "litellm.proxy.dd_span_tagger.set_active_span_tag" ) as mock_set_tag: DDSpanTagger.tag_request( user_api_key_dict=user_key, @@ -1598,7 +1600,7 @@ class TestDDSpanTaggerTagRequest: user_key = self._make_user_api_key_dict(key_alias=None, token=None) with patch( - "litellm.proxy.common_request_processing.set_active_span_tag" + "litellm.proxy.dd_span_tagger.set_active_span_tag" ) as mock_set_tag: DDSpanTagger.tag_request( user_api_key_dict=user_key, @@ -1612,7 +1614,7 @@ class TestDDSpanTaggerTagRequest: user_key = self._make_user_api_key_dict(key_alias=None, token=None) with patch( - "litellm.proxy.common_request_processing.set_active_span_tag" + "litellm.proxy.dd_span_tagger.set_active_span_tag" ) as mock_set_tag: DDSpanTagger.tag_request( user_api_key_dict=user_key,