Merge pull request BerriAI#18319 from BerriAI/litellm_feat_datadog_log_trace_linking

feat: datadog log trace linking

Author: uc4w6c

litellm/integrations/datadog/datadog.py

@@ -33,6 +33,7 @@
     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 @@ def _create_datadog_logging_payload_helper(
             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 @@ def _create_v0_logging_payload(
         )
         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

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):

tests/logging_callback_tests/test_datadog.py

@@ -26,6 +26,7 @@
 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"
+        assert dd_logger.DD_API_KEY == "fake-api-key"