From cbeaebf826b4884d2fecf81765d6f73849c86fa3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Thu, 26 Feb 2026 16:42:23 -0800 Subject: [PATCH] _add_dd_apm_tags_for_litellm_call_id (#22219) --- litellm/litellm_core_utils/dd_tracing.py | 47 ++++++++++++++++++- litellm/proxy/common_request_processing.py | 24 +++++++++- .../proxy/test_common_request_processing.py | 15 ++++++ 3 files changed, 83 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/dd_tracing.py b/litellm/litellm_core_utils/dd_tracing.py index ce784ecf6a8..ae4f46c38bd 100644 --- a/litellm/litellm_core_utils/dd_tracing.py +++ b/litellm/litellm_core_utils/dd_tracing.py @@ -5,7 +5,7 @@ If the ddtrace package is not installed, the tracer will be a no-op. """ from contextlib import contextmanager -from typing import TYPE_CHECKING, Any, Union +from typing import TYPE_CHECKING, Any, Optional, Union from litellm.secret_managers.main import get_secret_bool @@ -76,3 +76,48 @@ if should_use_dd_tracer: tracer = NullTracer() else: tracer = NullTracer() + + +def get_active_span() -> Optional[Any]: + """ + Return the active Datadog span, checking current span first and then root span. + """ + try: + current_span_fn = getattr(tracer, "current_span", None) + if callable(current_span_fn): + current_span = current_span_fn() + if current_span is not None: + return current_span + + current_root_span_fn = getattr(tracer, "current_root_span", None) + if callable(current_root_span_fn): + return current_root_span_fn() + except Exception: + return None + return None + + +def set_active_span_tag(tag_key: str, tag_value: str) -> bool: + """ + Best-effort helper to set a tag on the active Datadog span. + + Returns: + bool: True if a span tag was set, False otherwise. + """ + if not tag_key or tag_value is None: + return False + + span = get_active_span() + if span is None: + return False + + try: + if hasattr(span, "set_tag_str"): + span.set_tag_str(tag_key, str(tag_value)) + return True + if hasattr(span, "set_tag"): + span.set_tag(tag_key, str(tag_value)) + return True + except Exception: + return False + return False diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 40fae4e4a56..1269f58213a 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 tracer +from litellm.litellm_core_utils.dd_tracing import set_active_span_tag, 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, @@ -245,6 +245,26 @@ 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, @@ -642,6 +662,7 @@ 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")) ### AUTO STREAM USAGE TRACKING ### # If always_include_stream_usage is enabled and this is a streaming request @@ -658,7 +679,6 @@ class ProxyBaseLLMRequestProcessing: and "include_usage" not in self.data["stream_options"] ): self.data["stream_options"]["include_usage"] = True - ### CALL HOOKS ### - modify/reject incoming data before calling the model ## LOGGING OBJECT ## - initialize logging object for logging success/failure events for call diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index bf794478f10..9b905d24fd1 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -13,6 +13,7 @@ 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, @@ -79,6 +80,20 @@ class TestProxyBaseLLMRequestProcessing: pytest.fail("litellm_call_id is not a valid UUID") assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"] + 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) + monkeypatch.setattr( + litellm.proxy.common_request_processing, + "set_active_span_tag", + mock_set_active_span_tag, + ) + + _add_dd_apm_tags_for_litellm_call_id("test-call-id") + + mock_set_active_span_tag.assert_called_once_with( + "litellm.call_id", "test-call-id" + ) + @pytest.mark.asyncio async def test_should_apply_hierarchical_router_settings_as_override( self, monkeypatch