chore(llmobs): [MLOB-3828] link input and output guardrail spans for openai agents (#14488)

This PR implements span linking for Open AI Agents guardrail spans. 

For input guardrails, a span link is created from the guardrail span to
the first LLM span which is a sibling of the input guardrail (share the
same llmobs parent). This implementation prevents a link from the
guardrail span to the openai_agents.response span from being created
while still linking the guardrail to the Triage Agent (LLM) span.


<img width="708" height="264" alt="image"
src="https://github.com/user-attachments/assets/35df597a-9d27-464e-b302-9b241e2b4f4d"
/>


For output guardrails, a span link is created from the last finished LLM
span (for a given agent run) to the output guardrail span.


<img width="2766" height="178" alt="image"
src="https://github.com/user-attachments/assets/35773ad9-d512-4be5-8db9-8a559ab29047"
/>


With this implementation, we get a trace like the following which better
represents the guardrail's functionality of checking the user input to
the first agent and the final output of the last agent.

<img width="488" height="1296" alt="image"
src="https://github.com/user-attachments/assets/4aac3564-d9bb-491e-9155-5847cec1f8ce"
/>

This PR also removes the experimental `_record_object` method from the
LLMObs service since it is not being used anywhere.

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

Author: ncybul

ddtrace/llmobs/_constants.py

@@ -93,6 +93,10 @@
 DISPATCH_ON_LLM_TOOL_CHOICE = "on_llm_tool_choice"
 DISPATCH_ON_TOOL_CALL = "on_tool_call"
 
+DISPATCH_ON_GUARDRAIL_SPAN_START = "on_guardrail_span_start"
+DISPATCH_ON_LLM_SPAN_FINISH = "on_llm_span_finish"
+DISPATCH_ON_OPENAI_AGENT_SPAN_FINISH = "on_openai_agent_span_finish"
+
 # Tool call arguments are used to lookup the associated tool call info.
 # When there are no tool call args, we use this as a place-holder lookup key
 OAI_HANDOFF_TOOL_ARG = "{}"

ddtrace/llmobs/_integrations/openai_agents.py

@@ -12,7 +12,9 @@
 from ddtrace.internal.utils import get_argument_value
 from ddtrace.internal.utils.formats import format_trace_id
 from ddtrace.llmobs._constants import AGENT_MANIFEST
+from ddtrace.llmobs._constants import DISPATCH_ON_GUARDRAIL_SPAN_START
 from ddtrace.llmobs._constants import DISPATCH_ON_LLM_TOOL_CHOICE
+from ddtrace.llmobs._constants import DISPATCH_ON_OPENAI_AGENT_SPAN_FINISH
 from ddtrace.llmobs._constants import DISPATCH_ON_TOOL_CALL
 from ddtrace.llmobs._constants import DISPATCH_ON_TOOL_CALL_OUTPUT_USED
 from ddtrace.llmobs._constants import INPUT_MESSAGES
@@ -79,6 +81,10 @@ def trace(
         elif oai_span:
             self.oai_to_llmobs_span[oai_span.span_id] = llmobs_span
             self._llmobs_update_trace_info_input(oai_span, llmobs_span)
+
+            if oai_span.span_type == "guardrail":
+                core.dispatch(DISPATCH_ON_GUARDRAIL_SPAN_START, (llmobs_span,))
+
         return llmobs_span
 
     def _llmobs_set_tags(
@@ -131,6 +137,7 @@ def _llmobs_set_tags(
             self._llmobs_set_handoff_attributes(span, oai_span)
         elif span_type == "agent":
             self._llmobs_set_agent_attributes(span, oai_span)
+            core.dispatch(DISPATCH_ON_OPENAI_AGENT_SPAN_FINISH, ())
         elif span_type == "custom":
             custom_data = oai_span.formatted_custom_data
             if custom_data:

ddtrace/llmobs/_llmobs.py

@@ -48,7 +48,10 @@
 from ddtrace.llmobs._constants import ANNOTATIONS_CONTEXT_ID
 from ddtrace.llmobs._constants import DECORATOR
 from ddtrace.llmobs._constants import DEFAULT_PROJECT_NAME
+from ddtrace.llmobs._constants import DISPATCH_ON_GUARDRAIL_SPAN_START
+from ddtrace.llmobs._constants import DISPATCH_ON_LLM_SPAN_FINISH
 from ddtrace.llmobs._constants import DISPATCH_ON_LLM_TOOL_CHOICE
+from ddtrace.llmobs._constants import DISPATCH_ON_OPENAI_AGENT_SPAN_FINISH
 from ddtrace.llmobs._constants import DISPATCH_ON_TOOL_CALL
 from ddtrace.llmobs._constants import DISPATCH_ON_TOOL_CALL_OUTPUT_USED
 from ddtrace.llmobs._constants import EXPERIMENT_CSV_FIELD_MAX_SIZE
@@ -91,7 +94,6 @@
 from ddtrace.llmobs._experiment import JSONType
 from ddtrace.llmobs._utils import AnnotationContext
 from ddtrace.llmobs._utils import LinkTracker
-from ddtrace.llmobs._utils import ToolCallTracker
 from ddtrace.llmobs._utils import _get_ml_app
 from ddtrace.llmobs._utils import _get_nearest_llmobs_ancestor
 from ddtrace.llmobs._utils import _get_session_id
@@ -219,8 +221,6 @@ def __init__(
         self._annotations: List[Tuple[str, str, Dict[str, Any]]] = []
         self._annotation_context_lock = forksafe.RLock()
 
-        self._tool_call_tracker = ToolCallTracker()
-
     def _on_span_start(self, span: Span) -> None:
         if self.enabled and span.span_type == SpanTypes.LLM:
             self._activate_llmobs_span(span)
@@ -256,6 +256,9 @@ def _llmobs_span_event(self, span: Span) -> Optional[LLMObsSpanEvent]:
         if not span_kind:
             raise KeyError("Span kind not found in span context")
 
+        if span_kind == "llm":
+            core.dispatch(DISPATCH_ON_LLM_SPAN_FINISH, (span,))
+
         llmobs_span = LLMObsSpan()
         _dd_attrs = {
             "span_id": str(span.span_id),
@@ -501,9 +504,13 @@ def _stop_service(self) -> None:
         core.reset_listeners("threading.submit", self._current_trace_context)
         core.reset_listeners("threading.execution", self._llmobs_context_provider.activate)
 
-        core.reset_listeners(DISPATCH_ON_LLM_TOOL_CHOICE, self._tool_call_tracker.on_llm_tool_choice)
-        core.reset_listeners(DISPATCH_ON_TOOL_CALL, self._tool_call_tracker.on_tool_call)
-        core.reset_listeners(DISPATCH_ON_TOOL_CALL_OUTPUT_USED, self._tool_call_tracker.on_tool_call_output_used)
+        core.reset_listeners(DISPATCH_ON_LLM_TOOL_CHOICE, self._link_tracker.on_llm_tool_choice)
+        core.reset_listeners(DISPATCH_ON_TOOL_CALL, self._link_tracker.on_tool_call)
+        core.reset_listeners(DISPATCH_ON_TOOL_CALL_OUTPUT_USED, self._link_tracker.on_tool_call_output_used)
+
+        core.reset_listeners(DISPATCH_ON_GUARDRAIL_SPAN_START, self._link_tracker.on_guardrail_span_start)
+        core.reset_listeners(DISPATCH_ON_LLM_SPAN_FINISH, self._link_tracker.on_llm_span_finish)
+        core.reset_listeners(DISPATCH_ON_OPENAI_AGENT_SPAN_FINISH, self._link_tracker.on_openai_agent_span_finish)
 
         forksafe.unregister(self._child_after_fork)
 
@@ -614,9 +621,13 @@ def enable(
             core.on("threading.submit", cls._instance._current_trace_context, "llmobs_ctx")
             core.on("threading.execution", cls._instance._llmobs_context_provider.activate)
 
-            core.on(DISPATCH_ON_LLM_TOOL_CHOICE, cls._instance._tool_call_tracker.on_llm_tool_choice)
-            core.on(DISPATCH_ON_TOOL_CALL, cls._instance._tool_call_tracker.on_tool_call)
-            core.on(DISPATCH_ON_TOOL_CALL_OUTPUT_USED, cls._instance._tool_call_tracker.on_tool_call_output_used)
+            core.on(DISPATCH_ON_LLM_TOOL_CHOICE, cls._instance._link_tracker.on_llm_tool_choice)
+            core.on(DISPATCH_ON_TOOL_CALL, cls._instance._link_tracker.on_tool_call)
+            core.on(DISPATCH_ON_TOOL_CALL_OUTPUT_USED, cls._instance._link_tracker.on_tool_call_output_used)
+
+            core.on(DISPATCH_ON_GUARDRAIL_SPAN_START, cls._instance._link_tracker.on_guardrail_span_start)
+            core.on(DISPATCH_ON_LLM_SPAN_FINISH, cls._instance._link_tracker.on_llm_span_finish)
+            core.on(DISPATCH_ON_OPENAI_AGENT_SPAN_FINISH, cls._instance._link_tracker.on_openai_agent_span_finish)
 
             atexit.register(cls.disable)
             telemetry_writer.product_activated(TELEMETRY_APM_PRODUCT.LLMOBS, True)
@@ -811,41 +822,6 @@ def disable(cls) -> None:
 
         log.debug("%s disabled", cls.__name__)
 
-    def _record_object(self, span, obj, input_or_output):
-        if obj is None:
-            return
-        span_links = []
-        for span_link in self._link_tracker.get_span_links_from_object(obj):
-            try:
-                if span_link["attributes"]["from"] == "input" and input_or_output == "output":
-                    continue
-            except KeyError:
-                log.debug("failed to read span link: ", span_link)
-                continue
-            span_links.append(
-                {
-                    "trace_id": span_link["trace_id"],
-                    "span_id": span_link["span_id"],
-                    "attributes": {
-                        "from": span_link["attributes"]["from"],
-                        "to": input_or_output,
-                    },
-                }
-            )
-        self._tag_span_links(span, span_links)
-        self._link_tracker.add_span_links_to_object(
-            obj,
-            [
-                {
-                    "trace_id": self.export_span(span)["trace_id"],
-                    "span_id": self.export_span(span)["span_id"],
-                    "attributes": {
-                        "from": input_or_output,
-                    },
-                }
-            ],
-        )
-
     def _tag_span_links(self, span, span_links):
         if not span_links:
             return

ddtrace/llmobs/_utils.py

@@ -5,6 +5,7 @@
 from typing import Dict
 from typing import List
 from typing import Optional
+from typing import Set
 from typing import Tuple
 from typing import Union
 
@@ -88,23 +89,6 @@ def validate_prompt(prompt: dict) -> Dict[str, Union[str, dict, List[str]]]:
     return validated_prompt
 
 
-class LinkTracker:
-    def __init__(self, object_span_links=None):
-        self._object_span_links = object_span_links or {}
-
-    def get_object_id(self, obj):
-        return f"{type(obj).__name__}_{id(obj)}"
-
-    def add_span_links_to_object(self, obj, span_links):
-        obj_id = self.get_object_id(obj)
-        if obj_id not in self._object_span_links:
-            self._object_span_links[obj_id] = []
-        self._object_span_links[obj_id] += span_links
-
-    def get_span_links_from_object(self, obj):
-        return self._object_span_links.get(self.get_object_id(obj), [])
-
-
 class AnnotationContext:
     def __init__(self, _register_annotator, _deregister_annotator):
         self._register_annotator = _register_annotator
@@ -298,12 +282,20 @@ class TrackedToolCall:
     tool_kind: str = "function"  # one of "function", "handoff"
 
 
-class ToolCallTracker:
-    """Used to track tool data and their associated llm/tool spans for span linking."""
+class LinkTracker:
+    """
+    This class is used to create span links across integrations.
+
+    The primary use cases are:
+    - Linking LLM spans to their associated tool spans and vice versa
+    - Linking LLM spans to their associated guardrail spans and vice versa
+    """
 
     def __init__(self):
         self._tool_calls: Dict[str, TrackedToolCall] = {}  # maps tool id's to tool call data
         self._lookup_tool_id: Dict[Tuple[str, str], str] = {}  # maps (tool_name, arguments) to tool id's
+        self._active_guardrail_spans: Set[Span] = set()
+        self._last_llm_span: Optional[Span] = None
 
     def on_llm_tool_choice(
         self, tool_id: str, tool_name: str, arguments: str, llm_span_context: Dict[str, str]
@@ -379,3 +371,47 @@ def on_tool_call_output_used(self, tool_id: str, llm_span: Span) -> None:
             "output",
             "input",
         )
+
+    def on_llm_span_finish(self, span: Span) -> None:
+        """
+        Called when an LLM span event is created. If the LLM span is the first LLM span,
+        it will consume all active guardrail links.
+        """
+        self._last_llm_span = span
+        spans_to_remove = set()
+        for guardrail_span in self._active_guardrail_spans:
+            # some guardrail spans may have LLM spans as children which we don't want to link to
+            if _get_nearest_llmobs_ancestor(guardrail_span) == _get_nearest_llmobs_ancestor(span):
+                add_span_link(
+                    span,
+                    str(guardrail_span.span_id),
+                    format_trace_id(guardrail_span.trace_id),
+                    "output",
+                    "input",
+                )
+                spans_to_remove.add(guardrail_span)
+        self._active_guardrail_spans -= spans_to_remove
+
+    def on_guardrail_span_start(self, span: Span) -> None:
+        """
+        Called when a guardrail span starts. This is used to track the active guardrail
+        spans and link the output of the last LLM span to the input of the guardrail span.
+        """
+        self._active_guardrail_spans.add(span)
+        if self._last_llm_span is not None and _get_nearest_llmobs_ancestor(span) == _get_nearest_llmobs_ancestor(
+            self._last_llm_span
+        ):
+            add_span_link(
+                span,
+                str(self._last_llm_span.span_id),
+                format_trace_id(self._last_llm_span.trace_id),
+                "output",
+                "input",
+            )
+
+    def on_openai_agent_span_finish(self) -> None:
+        """
+        Called when an OpenAI agent span finishes. This is used to reset the last LLM span
+        since output guardrails are only linked to the last LLM span for a particular agent.
+        """
+        self._last_llm_span = None