feat(llmobs): allow span processor to return None to omit spans (#14260)

Allow users to omit any traces they want by returning none in the
user_span_processor method.

Motivation:
I am trying to prevent some auto traces from cluttering my Datadog
observability dashboards. These traces create noise and make it harder
to focus on the more critical traces. My LLM observability overview is
filled with what is clustered as empty input which is incorrect. e.g.
all embedding traces are are just a spammer.

To do that, I am using the new span_processor, and the new way to omit
specific spans, will be to return null by the relevant span processor.

I have added a test for that.

The only risk is that there is no indication that span was omitted but I
was afraid that a debug log would be too spammy. But please let me know
if that will help. Also there are some telemetric collected about sent
span which might also need to be changed, will be great to hear your
thoughts

Co-authored-by: amirbenami <amir@linxsecurity.io>

Author: Kyle-Verhoog

ddtrace/llmobs/_llmobs.py

@@ -145,7 +145,10 @@ class LLMObsSpan:
     Passed to the `span_processor` function in the `enable` or `register_processor` methods.
 
     Example::
-        def span_processor(span: LLMObsSpan) -> LLMObsSpan:
+        def span_processor(span: LLMObsSpan) -> Optional[LLMObsSpan]:
+            # Modify input/output
+            if span.get_tag("omit_span") == "1":
+                return None
             if span.get_tag("no_input") == "1":
                 span.input = []
             return span
@@ -178,7 +181,7 @@ class LLMObs(Service):
     def __init__(
         self,
         tracer: Optional[Tracer] = None,
-        span_processor: Optional[Callable[[LLMObsSpan], LLMObsSpan]] = None,
+        span_processor: Optional[Callable[[LLMObsSpan], Optional[LLMObsSpan]]] = None,
     ) -> None:
         super(LLMObs, self).__init__()
         self.tracer = tracer or ddtrace.tracer
@@ -230,6 +233,8 @@ def _submit_llmobs_span(self, span: Span) -> None:
         span_event = None
         try:
             span_event = self._llmobs_span_event(span)
+            if span_event is None:
+                return
             self._llmobs_span_writer.enqueue(span_event)
         except (KeyError, TypeError, ValueError):
             log.error(
@@ -241,7 +246,7 @@ def _submit_llmobs_span(self, span: Span) -> None:
             if self._evaluator_runner:
                 self._evaluator_runner.enqueue(span_event, span)
 
-    def _llmobs_span_event(self, span: Span) -> LLMObsSpanEvent:
+    def _llmobs_span_event(self, span: Span) -> Optional[LLMObsSpanEvent]:
         """Span event object structure."""
         span_kind = span._get_ctx_item(SPAN_KIND)
         if not span_kind:
@@ -321,8 +326,12 @@ def _llmobs_span_event(self, span: Span) -> LLMObsSpanEvent:
             try:
                 llmobs_span._tags = cast(Dict[str, str], span._get_ctx_item(TAGS))
                 user_llmobs_span = self._user_span_processor(llmobs_span)
+                if user_llmobs_span is None:
+                    return None
                 if not isinstance(user_llmobs_span, LLMObsSpan):
-                    raise TypeError("User span processor must return an LLMObsSpan, got %r" % type(user_llmobs_span))
+                    raise TypeError(
+                        "User span processor must return an LLMObsSpan or None, 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)
@@ -488,7 +497,7 @@ def enable(
         project_name: Optional[str] = None,
         env: Optional[str] = None,
         service: Optional[str] = None,
-        span_processor: Optional[Callable[[LLMObsSpan], LLMObsSpan]] = None,
+        span_processor: Optional[Callable[[LLMObsSpan], Optional[LLMObsSpan]]] = None,
         _tracer: Optional[Tracer] = None,
         _auto: bool = False,
     ) -> None:
@@ -505,8 +514,8 @@ def enable(
         :param str project_name: Your project name used for experiments.
         :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.
+        :param Callable[[LLMObsSpan], Optional[LLMObsSpan]] span_processor: A function that takes an LLMObsSpan and
+            returns an LLMObsSpan or None. If None is returned, the span will be omitted and not sent to LLMObs.
         """
         if cls.enabled:
             log.debug("%s already enabled", cls.__name__)
@@ -741,14 +750,16 @@ def experiment(
         )
 
     @classmethod
-    def register_processor(cls, processor: Optional[Callable[[LLMObsSpan], LLMObsSpan]] = None) -> None:
+    def register_processor(cls, processor: Optional[Callable[[LLMObsSpan], Optional[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.
+        You can also return None to omit the span entirely from being sent to LLMObs.
 
         To deregister the processor, call `register_processor(None)`.
 
-        :param processor: A function that takes an LLMObsSpan and returns an LLMObsSpan.
+        :param processor: A function that takes an LLMObsSpan and returns an LLMObsSpan or None.
+                         If None is returned, the span will be omitted and not sent to LLMObs.
         """
         cls._instance._user_span_processor = processor
 

releasenotes/notes/llmobs-span-proc-drop-170c0be459e84722.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    LLM Observability: add ability to drop spans by having a ``SpanProcessor`` return ``None``.

tests/llmobs/test_llmobs.py

@@ -1,6 +1,7 @@
 import asyncio
 import os
 from textwrap import dedent
+from typing import Optional
 
 import pytest
 
@@ -167,6 +168,26 @@ def test_processor_bad_return_type(self, llmobs, llmobs_enable_opts, llmobs_even
         assert llmobs_events[0]["meta"]["input"] == {"messages": [{"content": "value", "role": ""}]}
         assert llmobs_events[0]["meta"]["output"] == {"messages": [{"content": "value", "role": ""}]}
 
+    def _omit_span_processor(span: LLMObsSpan) -> Optional[LLMObsSpan]:
+        if span.get_tag("omit_span") == "true":
+            return None
+        return span
+
+    @pytest.mark.parametrize("llmobs_enable_opts", [dict(span_processor=_omit_span_processor)])
+    def test_processor_omit_span(self, llmobs, llmobs_enable_opts, llmobs_events):
+        """Test that a processor that returns None omits the span from being sent."""
+        # Create a span that should be omitted
+        with llmobs.llm() as llm_span:
+            llmobs.annotate(llm_span, input_data="omit me", output_data="response", tags={"omit_span": "true"})
+
+        # Create a span that should be kept
+        with llmobs.llm() as llm_span:
+            llmobs.annotate(llm_span, input_data="keep me", output_data="response", tags={"omit_span": "false"})
+
+        # Only the second span should be in the events
+        assert len(llmobs_events) == 1
+        assert llmobs_events[0]["meta"]["input"]["messages"][0]["content"] == "keep me"
+
     def test_ddtrace_run_register_processor(self, ddtrace_run_python_code_in_subprocess, llmobs_backend):
         """Users using ddtrace-run can register a processor to be called on each LLMObs span."""
         env = os.environ.copy()