From 41bbb3a6a55e139ba26684604bc5392763c6f50a Mon Sep 17 00:00:00 2001 From: Yuta Saito Date: Mon, 22 Dec 2025 06:50:31 +0900 Subject: [PATCH] feat: datadog log trace linking --- litellm/integrations/datadog/datadog.py | 52 +++++++ litellm/types/integrations/datadog.py | 23 +-- tests/logging_callback_tests/test_datadog.py | 140 +++++++++++++++++-- 3 files changed, 199 insertions(+), 16 deletions(-) diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py index 08bf5034f0..503e8d8c87 100644 --- a/litellm/integrations/datadog/datadog.py +++ b/litellm/integrations/datadog/datadog.py @@ -33,6 +33,7 @@ from litellm.integrations.datadog.datadog_handler import ( get_datadog_source, get_datadog_tags, ) +from litellm.litellm_core_utils.dd_tracing import tracer from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, get_async_httpx_client, @@ -337,6 +338,7 @@ class DataDogLogger( service=get_datadog_service(), status=status, ) + self._add_trace_context_to_payload(dd_payload=dd_payload) return dd_payload def create_datadog_logging_payload( @@ -574,6 +576,56 @@ class DataDogLogger( ) return dd_payload + def _add_trace_context_to_payload( + self, + dd_payload: DatadogPayload, + ) -> None: + """Attach Datadog APM trace context if one is active.""" + + try: + trace_context = self._get_active_trace_context() + if trace_context is None: + return + + dd_payload["dd.trace_id"] = trace_context["trace_id"] + span_id = trace_context.get("span_id") + if span_id is not None: + dd_payload["dd.span_id"] = span_id + except Exception: + verbose_logger.exception( + "Datadog: Failed to attach trace context to payload" + ) + + def _get_active_trace_context(self) -> Optional[Dict[str, str]]: + try: + current_span = None + current_span_fn = getattr(tracer, "current_span", None) + if callable(current_span_fn): + current_span = current_span_fn() + + if current_span is None: + current_root_span_fn = getattr(tracer, "current_root_span", None) + if callable(current_root_span_fn): + current_span = current_root_span_fn() + + if current_span is None: + return None + + trace_id = getattr(current_span, "trace_id", None) + if trace_id is None: + return None + + span_id = getattr(current_span, "span_id", None) + trace_context: Dict[str, str] = {"trace_id": str(trace_id)} + if span_id is not None: + trace_context["span_id"] = str(span_id) + return trace_context + except Exception: + verbose_logger.exception( + "Datadog: Failed to retrieve active trace context from tracer" + ) + return None + async def async_health_check(self) -> IntegrationHealthCheckStatus: """ Check if the service is healthy diff --git a/litellm/types/integrations/datadog.py b/litellm/types/integrations/datadog.py index b741184394..89faac2783 100644 --- a/litellm/types/integrations/datadog.py +++ b/litellm/types/integrations/datadog.py @@ -1,7 +1,7 @@ from enum import Enum from typing import Optional -from typing_extensions import TypedDict +from typing_extensions import NotRequired, TypedDict from litellm.types.integrations.custom_logger import StandardCustomLoggerInitParams @@ -14,13 +14,20 @@ class DataDogStatus(str, Enum): ERROR = "error" -class DatadogPayload(TypedDict, total=False): - ddsource: str - ddtags: str - hostname: str - message: str - service: str - status: str +DatadogPayload = TypedDict( + "DatadogPayload", + { + "ddsource": str, + "ddtags": str, + "hostname": str, + "message": str, + "service": str, + "status": str, + "dd.trace_id": NotRequired[str], + "dd.span_id": NotRequired[str], + }, + total=False, +) class DD_ERRORS(Enum): diff --git a/tests/logging_callback_tests/test_datadog.py b/tests/logging_callback_tests/test_datadog.py index 7fa4dfe722..fc4b3ff3cf 100644 --- a/tests/logging_callback_tests/test_datadog.py +++ b/tests/logging_callback_tests/test_datadog.py @@ -26,6 +26,7 @@ import litellm from litellm import completion from litellm._logging import verbose_logger from litellm.integrations.datadog.datadog import * +import litellm.integrations.datadog.datadog as datadog_module from datetime import datetime, timedelta from litellm.types.utils import ( StandardLoggingPayload, @@ -90,6 +91,24 @@ def create_standard_logging_payload() -> StandardLoggingPayload: ) +class _DummySpan: + def __init__(self, trace_id=None, span_id=None): + self.trace_id = trace_id + self.span_id = span_id + + +class _DummyTracer: + def __init__(self, current_span=None, current_root_span=None): + self._current_span = current_span + self._current_root_span = current_root_span + + def current_span(self): + return self._current_span + + def current_root_span(self): + return self._current_root_span + + @pytest.mark.asyncio async def test_create_datadog_logging_payload(): """Test creating a DataDog logging payload from a standard logging object""" @@ -219,20 +238,35 @@ async def test_datadog_logging_http_request(): # Get the expected fields and their types from DatadogPayload expected_fields = DatadogPayload.__annotations__ - # Assert that all elements in body have the fields of DatadogPayload with correct types + required_fields = { + "ddsource": str, + "ddtags": str, + "hostname": str, + "message": str, + "service": str, + "status": str, + } + optional_fields = set(expected_fields.keys()) - set(required_fields.keys()) + + # Assert that all elements in body have the required fields with correct types for log in body: assert isinstance(log, dict), "Each log should be a dictionary" - for field, expected_type in expected_fields.items(): + for field, expected_type in required_fields.items(): assert field in log, f"Field '{field}' is missing from the log" assert isinstance( log[field], expected_type ), f"Field '{field}' has incorrect type. Expected {expected_type}, got {type(log[field])}" - # Additional assertion to ensure no extra fields are present - for log in body: - assert set(log.keys()) == set( - expected_fields.keys() - ), f"Log contains unexpected fields: {set(log.keys()) - set(expected_fields.keys())}" + for optional_field in optional_fields: + if optional_field in log: + assert isinstance( + log[optional_field], str + ), f"Optional field '{optional_field}' must be a string" + + unexpected_fields = set(log.keys()) - set(expected_fields.keys()) + assert ( + not unexpected_fields + ), f"Log contains unexpected fields: {unexpected_fields}" # Parse the 'message' field as JSON and check its structure message = json.loads(body[0]["message"]) @@ -256,6 +290,96 @@ async def test_datadog_logging_http_request(): pytest.fail(f"Test failed with exception: {str(e)}") +@pytest.mark.asyncio +async def test_add_trace_context_uses_current_span(monkeypatch): + monkeypatch.setenv("DD_SITE", "https://fake.datadoghq.com") + monkeypatch.setenv("DD_API_KEY", "anything") + tracer = _DummyTracer(current_span=_DummySpan(trace_id=123, span_id=456)) + monkeypatch.setattr(datadog_module, "tracer", tracer) + + dd_logger = DataDogLogger() + payload = DatadogPayload( + ddsource="litellm", + ddtags="env:test", + hostname="host", + message="{}", + service="svc", + status="info", + ) + + dd_logger._add_trace_context_to_payload(payload) + assert payload["dd.trace_id"] == "123" + assert payload["dd.span_id"] == "456" + + +@pytest.mark.asyncio +async def test_add_trace_context_falls_back_to_root_span(monkeypatch): + monkeypatch.setenv("DD_SITE", "https://fake.datadoghq.com") + monkeypatch.setenv("DD_API_KEY", "anything") + tracer = _DummyTracer( + current_span=None, + current_root_span=_DummySpan(trace_id=789, span_id=None), + ) + monkeypatch.setattr(datadog_module, "tracer", tracer) + + dd_logger = DataDogLogger() + payload = DatadogPayload( + ddsource="litellm", + ddtags="env:test", + hostname="host", + message="{}", + service="svc", + status="info", + ) + + dd_logger._add_trace_context_to_payload(payload) + assert payload["dd.trace_id"] == "789" + assert "dd.span_id" not in payload + + +@pytest.mark.asyncio +async def test_add_trace_context_handles_missing_tracer(monkeypatch): + monkeypatch.setenv("DD_SITE", "https://fake.datadoghq.com") + monkeypatch.setenv("DD_API_KEY", "anything") + monkeypatch.setattr(datadog_module, "tracer", object()) + + dd_logger = DataDogLogger() + payload = DatadogPayload( + ddsource="litellm", + ddtags="env:test", + hostname="host", + message="{}", + service="svc", + status="info", + ) + + dd_logger._add_trace_context_to_payload(payload) + assert "dd.trace_id" not in payload + assert "dd.span_id" not in payload + + +@pytest.mark.asyncio +async def test_add_trace_context_ignores_span_without_trace_id(monkeypatch): + monkeypatch.setenv("DD_SITE", "https://fake.datadoghq.com") + monkeypatch.setenv("DD_API_KEY", "anything") + tracer = _DummyTracer(current_span=_DummySpan(trace_id=None, span_id=555)) + monkeypatch.setattr(datadog_module, "tracer", tracer) + + dd_logger = DataDogLogger() + payload = DatadogPayload( + ddsource="litellm", + ddtags="env:test", + hostname="host", + message="{}", + service="svc", + status="info", + ) + + dd_logger._add_trace_context_to_payload(payload) + assert "dd.trace_id" not in payload + assert "dd.span_id" not in payload + + @pytest.mark.asyncio async def test_datadog_log_redis_failures(): """ @@ -701,4 +825,4 @@ def test_datadog_ignores_ddtrace_agent_host(): ) # Verify API key is set correctly - assert dd_logger.DD_API_KEY == "fake-api-key" \ No newline at end of file + assert dd_logger.DD_API_KEY == "fake-api-key"