feat(llmobs): add span processor (#13426)

Add capability to add a span processor. The processor can be used to
mutate or redact sensitive data contained in inputs and outputs from LLM
calls.

```python
from ddtrace.llmobs import LLMObsSpan

def my_processor(span: LLMObsSpan):
    for message in span.output:
        message["content"] = ""

LLMObs.enable(span_processor=my_processor)

LLMObs.register_processor(my_processor)
```

Public docs: DataDog/documentation#29365
Shared tests: TODO

Closes: #11179

Author: Kyle-Verhoog

ddtrace/llmobs/__init__.py

@@ -7,6 +7,7 @@
 """
 
 from ._llmobs import LLMObs
+from ._llmobs import LLMObsSpan
 
 
-__all__ = ["LLMObs"]
+__all__ = ["LLMObs", "LLMObsSpan"]

ddtrace/llmobs/_llmobs.py

@@ -1,11 +1,16 @@
+from dataclasses import dataclass
+from dataclasses import field
 import json
 import os
 import time
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import List
+from typing import Literal
 from typing import Optional
 from typing import Tuple
+from typing import TypedDict
 from typing import Union
 from typing import cast
 
@@ -101,14 +106,50 @@
 }
 
 
+@dataclass
+class LLMObsSpan:
+    """LLMObs span object.
+
+    Passed to the `span_processor` function in the `enable` or `register_processor` methods.
+
+    Example::
+        def span_processor(span: LLMObsSpan) -> LLMObsSpan:
+            if span.get_tag("no_input") == "1":
+                span.input = []
+            return span
+    """
+
+    class Message(TypedDict):
+        content: str
+        role: str
+
+    input: List[Message] = field(default_factory=list)
+    output: List[Message] = field(default_factory=list)
+    _tags: Dict[str, str] = field(default_factory=dict)
+
+    def get_tag(self, key: str) -> Optional[str]:
+        """Get a tag from the span.
+
+        :param str key: The key of the tag to get.
+        :return: The value of the tag or None if the tag does not exist.
+        :rtype: Optional[str]
+        """
+        return self._tags.get(key)
+
+
 class LLMObs(Service):
     _instance = None  # type: LLMObs
     enabled = False
 
-    def __init__(self, tracer: Optional[Tracer] = None):
+    def __init__(
+        self,
+        tracer: Optional[Tracer] = None,
+        span_processor: Optional[Callable[[LLMObsSpan], LLMObsSpan]] = None,
+    ) -> None:
         super(LLMObs, self).__init__()
         self.tracer = tracer or ddtrace.tracer
         self._llmobs_context_provider = LLMObsContextProvider()
+        self._user_span_processor = span_processor
         agentless_enabled = config._llmobs_agentless_enabled if config._llmobs_agentless_enabled is not None else True
         self._llmobs_span_writer = LLMObsSpanWriter(
             interval=float(os.getenv("_DD_LLMOBS_WRITER_INTERVAL", 1.0)),
@@ -160,33 +201,46 @@ def _submit_llmobs_span(self, span: Span) -> None:
             if self._evaluator_runner:
                 self._evaluator_runner.enqueue(span_event, span)
 
-    @classmethod
-    def _llmobs_span_event(cls, span: Span) -> LLMObsSpanEvent:
+    def _llmobs_span_event(self, span: Span) -> LLMObsSpanEvent:
         """Span event object structure."""
         span_kind = span._get_ctx_item(SPAN_KIND)
         if not span_kind:
             raise KeyError("Span kind not found in span context")
+
+        llmobs_span = LLMObsSpan()
+
         meta: Dict[str, Any] = {"span.kind": span_kind, "input": {}, "output": {}}
         if span_kind in ("llm", "embedding") and span._get_ctx_item(MODEL_NAME) is not None:
             meta["model_name"] = span._get_ctx_item(MODEL_NAME)
             meta["model_provider"] = (span._get_ctx_item(MODEL_PROVIDER) or "custom").lower()
         meta["metadata"] = span._get_ctx_item(METADATA) or {}
 
+        input_type: Literal["value", "messages", ""] = ""
+        output_type: Literal["value", "messages", ""] = ""
+        if span._get_ctx_item(INPUT_VALUE) is not None:
+            input_type = "value"
+            llmobs_span.input = [
+                {"content": safe_json(span._get_ctx_item(INPUT_VALUE), ensure_ascii=False), "role": ""}
+            ]
+
         input_messages = span._get_ctx_item(INPUT_MESSAGES)
         if span_kind == "llm" and input_messages is not None:
-            meta["input"]["messages"] = enforce_message_role(input_messages)
+            input_type = "messages"
+            llmobs_span.input = cast(List[LLMObsSpan.Message], enforce_message_role(input_messages))
 
-        if span._get_ctx_item(INPUT_VALUE) is not None:
-            meta["input"]["value"] = safe_json(span._get_ctx_item(INPUT_VALUE), ensure_ascii=False)
+        if span._get_ctx_item(OUTPUT_VALUE) is not None:
+            output_type = "value"
+            llmobs_span.output = [
+                {"content": safe_json(span._get_ctx_item(OUTPUT_VALUE), ensure_ascii=False), "role": ""}
+            ]
 
         output_messages = span._get_ctx_item(OUTPUT_MESSAGES)
         if span_kind == "llm" and output_messages is not None:
-            meta["output"]["messages"] = enforce_message_role(output_messages)
+            output_type = "messages"
+            llmobs_span.output = cast(List[LLMObsSpan.Message], enforce_message_role(output_messages))
 
         if span_kind == "embedding" and span._get_ctx_item(INPUT_DOCUMENTS) is not None:
             meta["input"]["documents"] = span._get_ctx_item(INPUT_DOCUMENTS)
-        if span._get_ctx_item(OUTPUT_VALUE) is not None:
-            meta["output"]["value"] = safe_json(span._get_ctx_item(OUTPUT_VALUE), ensure_ascii=False)
         if span_kind == "retrieval" and span._get_ctx_item(OUTPUT_DOCUMENTS) is not None:
             meta["output"]["documents"] = span._get_ctx_item(OUTPUT_DOCUMENTS)
         if span._get_ctx_item(INPUT_PROMPT) is not None:
@@ -205,6 +259,32 @@ def _llmobs_span_event(cls, span: Span) -> LLMObsSpanEvent:
                     ERROR_TYPE: span.get_tag(ERROR_TYPE),
                 }
             )
+
+        if self._user_span_processor:
+            error = False
+            try:
+                llmobs_span._tags = cast(Dict[str, str], span._get_ctx_item(TAGS))
+                user_llmobs_span = self._user_span_processor(llmobs_span)
+                if not isinstance(user_llmobs_span, LLMObsSpan):
+                    raise TypeError("User span processor must return an LLMObsSpan, got %r" % type(user_llmobs_span))
+                llmobs_span = user_llmobs_span
+            except Exception as e:
+                log.error("Error in LLMObs span processor (%r): %r", self._user_span_processor, e)
+                error = True
+            finally:
+                telemetry.record_llmobs_user_processor_called(error)
+
+        if llmobs_span.input is not None:
+            if input_type == "messages":
+                meta["input"]["messages"] = llmobs_span.input
+            elif input_type == "value":
+                meta["input"]["value"] = llmobs_span.input[0]["content"]
+        if llmobs_span.output is not None:
+            if output_type == "messages":
+                meta["output"]["messages"] = llmobs_span.output
+            elif output_type == "value":
+                meta["output"]["value"] = llmobs_span.output[0]["content"]
+
         if not meta["input"]:
             meta.pop("input")
         if not meta["output"]:
@@ -233,7 +313,7 @@ def _llmobs_span_event(cls, span: Span) -> LLMObsSpanEvent:
             span._set_ctx_item(SESSION_ID, session_id)
             llmobs_span_event["session_id"] = session_id
 
-        llmobs_span_event["tags"] = cls._llmobs_tags(span, ml_app, session_id)
+        llmobs_span_event["tags"] = self._llmobs_tags(span, ml_app, session_id)
 
         span_links = span._get_ctx_item(SPAN_LINKS)
         if isinstance(span_links, list) and span_links:
@@ -339,6 +419,7 @@ def enable(
         api_key: Optional[str] = None,
         env: Optional[str] = None,
         service: Optional[str] = None,
+        span_processor: Optional[Callable[[LLMObsSpan], LLMObsSpan]] = None,
         _tracer: Optional[Tracer] = None,
         _auto: bool = False,
     ) -> None:
@@ -352,6 +433,8 @@ def enable(
         :param str api_key: Your datadog api key.
         :param str env: Your environment name.
         :param str service: Your service name.
+        :param Callable[[LLMObsSpan], LLMObsSpan] span_processor: A function that takes an LLMObsSpan and returns an
+            LLMObsSpan.
         """
         if cls.enabled:
             log.debug("%s already enabled", cls.__name__)
@@ -379,9 +462,9 @@ def enable(
                 )
 
             config._llmobs_agentless_enabled = should_use_agentless(
-                user_defined_agentless_enabled=agentless_enabled
-                if agentless_enabled is not None
-                else config._llmobs_agentless_enabled
+                user_defined_agentless_enabled=(
+                    agentless_enabled if agentless_enabled is not None else config._llmobs_agentless_enabled
+                )
             )
 
             if config._llmobs_agentless_enabled:
@@ -411,7 +494,7 @@ def enable(
                 cls._patch_integrations()
 
             # override the default _instance with a new tracer
-            cls._instance = cls(tracer=_tracer)
+            cls._instance = cls(tracer=_tracer, span_processor=span_processor)
             cls.enabled = True
             cls._instance.start()
 
@@ -434,6 +517,18 @@ def enable(
         finally:
             telemetry.record_llmobs_enabled(error, config._llmobs_agentless_enabled, config._dd_site, start_ns, _auto)
 
+    @classmethod
+    def register_processor(cls, processor: Optional[Callable[[LLMObsSpan], LLMObsSpan]] = None) -> None:
+        """Register a processor to be called on each LLMObs span.
+
+        This can be used to modify the span before it is sent to LLMObs. For example, you can modify the input/output.
+
+        To deregister the processor, call `register_processor(None)`.
+
+        :param processor: A function that takes an LLMObsSpan and returns an LLMObsSpan.
+        """
+        cls._instance._user_span_processor = processor
+
     @classmethod
     def _integration_is_enabled(cls, integration: str) -> bool:
         if integration not in SUPPORTED_LLMOBS_INTEGRATIONS:

ddtrace/llmobs/_telemetry.py

@@ -31,6 +31,7 @@ class LLMObsTelemetryMetrics:
     USER_FLUSHES = "user_flush"
     INJECT_HEADERS = "inject_distributed_headers"
     ACTIVATE_HEADERS = "activate_distributed_headers"
+    USER_PROCESSOR_CALLED = "user_processor_called"
 
 
 def _find_integration_from_tags(tags):
@@ -156,6 +157,16 @@ def record_llmobs_annotate(span: Optional[Span], error: Optional[str]):
     )
 
 
+def record_llmobs_user_processor_called(error: bool) -> None:
+    tags = [("error", "1" if error else "0")]
+    telemetry_writer.add_count_metric(
+        namespace=TELEMETRY_NAMESPACE.MLOBS,
+        name=LLMObsTelemetryMetrics.USER_PROCESSOR_CALLED,
+        value=1,
+        tags=tuple(tags),
+    )
+
+
 def record_llmobs_submit_evaluation(join_on: Dict[str, Any], metric_type: str, error: Optional[str]):
     _metric_type = metric_type if metric_type in ("categorical", "score") else "other"
     custom_joining_key = str(int(join_on.get("tag") is not None))

releasenotes/notes/llmobs-processor-d5cb47b12bc3bbd1.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    LLM Observability: add processor capability to process span inputs and outputs. See usage documentation [here](https://docs.datadoghq.com/llm_observability/setup/sdk/python/#span-processing).

tests/llmobs/_utils.py

@@ -832,6 +832,8 @@ def _expected_span_link(span_event, link_from, link_to):
 
 
 class TestLLMObsSpanWriter(LLMObsSpanWriter):
+    __test__ = False
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.events = []

tests/llmobs/conftest.py

@@ -2,6 +2,7 @@
 from http.server import HTTPServer
 import json
 import os
+import pprint
 import threading
 import time
 
@@ -224,26 +225,35 @@ def _llmobs_backend():
 
 @pytest.fixture
 def llmobs_backend(_llmobs_backend):
-    _, reqs = _llmobs_backend
+    _url, reqs = _llmobs_backend
 
     class _LLMObsBackend:
+        def url(self):
+            return _url
+
         def wait_for_num_events(self, num, attempts=1000):
             for _ in range(attempts):
                 if len(reqs) == num:
                     return [json.loads(r["body"]) for r in reqs]
                 # time.sleep will yield the GIL so the server can process the request
                 time.sleep(0.001)
             else:
-                raise TimeoutError(f"Expected {num} events, got {len(reqs)}")
+                raise TimeoutError(f"Expected {num} events, got {len(reqs)}: {pprint.pprint(reqs)}")
 
     return _LLMObsBackend()
 
 
+@pytest.fixture
+def llmobs_enable_opts():
+    yield {}
+
+
 @pytest.fixture
 def llmobs(
     ddtrace_global_config,
     monkeypatch,
     tracer,
+    llmobs_enable_opts,
     llmobs_env,
     llmobs_span_writer,
     mock_llmobs_eval_metric_writer,
@@ -256,7 +266,7 @@ def llmobs(
     global_config.update(ddtrace_global_config)
     # TODO: remove once rest of tests are moved off of global config tampering
     with override_global_config(global_config):
-        llmobs_service.enable(_tracer=tracer)
+        llmobs_service.enable(_tracer=tracer, **llmobs_enable_opts)
         llmobs_service._instance._llmobs_span_writer = llmobs_span_writer
         llmobs_service._instance._llmobs_span_writer.start()
         yield llmobs_service