chore(llmobs): span linking for oai agents sdk (#13072)

Add span linking between tool & llm spans for the openai agents sdk.

We use the core dispatch api since span linking requires
cross-integration communication in the case where someone selects "chat
completions" as the llm api to use for the agents sdk.

Signals are dispatched
- when LLM spans finish (chat completions api) in the oai integration
- when LLM spans finish (responses api) in the agents sdk integration
- when tool calls/handoffs finish in the agents sdk integration

`ToolCallTracker` in `ddtrace.llmobs._utils` contains the functions that
handles these signals to add span links.

### Links created
**[LLM output -> tool input]** for the case where an LLM span chooses a
tool and that tool is later executed via the agents sdk. We do this by
mapping the tool name & arguments to it's tool id. When the tool call is
triggered, we have access to it's name and arguments. From there, we can
look up it's tool id and the LLM span that is used to generate that
argument. We pop the tool name/arg from the lookup dictionary after it's
used.

**[Tool output -> LLM input]** for the case where a tool's output is fed
back into a later LLM call, either in the same agent or another agent.
We can tell this since the tool_id is present in the LLM's input
messages. We then use this tool id to lookup the tool span.

So the general lifecycle is:
1. An llm chooses a tool. A save the tool id, tool name, and tool
arguments and correlate it with the LLM span
2. The tool is run. 
- We look at the argument and name of the tool and use it to look up the
LLM span that chose this tool. We then delete the name/arg from the
lookup dict. We then
- We save the span/trace id of the tool and correlate it with the
tool_id
4. The tool output is used as input for an LLM span. We have access to
the tool id here, and lookup the span/trace id of the tool to link it to
the LLM span

#### A note on handoffs
Hand-offs are implemented as tool calls in the agents SDK, so the span
linking logic is largely the same. Two notes
- there are no arguments for handoffs, so we use a dummy default lookup
key for [LLM output -> tool input] linking step
- the tool_id representing a handoff may be continually used as input
for an LLM call since the list of messages is kept and added to across
agent runs. However, it realistically should only be linked to the first
LLM call of the agent being handed-off to since. Unlike other tool
calls, a handoff is only an orchestration step and it doesn't provide
extra context actually "used" in downstream llm generations
- There are two brittle parts of hand-off linking that relies on some
implementation details internal to the agents sdk
    - We are re-constructing the raw tool name used for hand-offs 
`handoff_tool_name =
"transfer_to_{}".format("_".join(oai_span.to_agent.split(" ")).lower())`
- We are using `{}` as the placeholder for the hand-off tool call
argument. This is what's generated by the LLM when it chooses a handoff.

We can improve on this by inferring these values when an LLM chooses a
handoff tool, but this requires a bit more exploring

## Checklist
- [ ] 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
- [ ] 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)

---------

Co-authored-by: lievan <[email protected]>

Author: lievan

ddtrace/llmobs/_constants.py

@@ -69,3 +69,11 @@
 NAME = "_ml_obs.name"
 DECORATOR = "_ml_obs.decorator"
 INTEGRATION = "_ml_obs.integration"
+
+DISPATCH_ON_TOOL_CALL_OUTPUT_USED = "on_tool_call_output_used"
+DISPATCH_ON_LLM_TOOL_CHOICE = "on_llm_tool_choice"
+DISPATCH_ON_TOOL_CALL = "on_tool_call"
+
+# 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.py

@@ -5,8 +5,12 @@
 from typing import Optional
 from typing import Tuple
 
+from ddtrace.internal import core
 from ddtrace.internal.constants import COMPONENT
+from ddtrace.internal.utils.formats import format_trace_id
 from ddtrace.internal.utils.version import parse_version
+from ddtrace.llmobs._constants import DISPATCH_ON_LLM_TOOL_CHOICE
+from ddtrace.llmobs._constants import DISPATCH_ON_TOOL_CALL_OUTPUT_USED
 from ddtrace.llmobs._constants import INPUT_DOCUMENTS
 from ddtrace.llmobs._constants import INPUT_MESSAGES
 from ddtrace.llmobs._constants import INPUT_TOKENS_METRIC_KEY
@@ -158,6 +162,9 @@ def _llmobs_set_meta_tags_from_chat(span: Span, kwargs: Dict[str, Any], messages
         """Extract prompt/response tags from a chat completion and set them as temporary "_ml_obs.meta.*" tags."""
         input_messages = []
         for m in kwargs.get("messages", []):
+            tool_call_id = m.get("tool_call_id")
+            if tool_call_id:
+                core.dispatch(DISPATCH_ON_TOOL_CALL_OUTPUT_USED, (tool_call_id, span))
             input_messages.append({"content": str(_get_attr(m, "content", "")), "role": str(_get_attr(m, "role", ""))})
         parameters = {k: v for k, v in kwargs.items() if k not in ("model", "messages", "tools", "functions")}
         span._set_ctx_items({INPUT_MESSAGES: input_messages, METADATA: parameters})
@@ -199,13 +206,28 @@ def _llmobs_set_meta_tags_from_chat(span: Span, kwargs: Dict[str, Any], messages
                 continue
             tool_calls = _get_attr(choice_message, "tool_calls", []) or []
             for tool_call in tool_calls:
+                tool_args = getattr(tool_call.function, "arguments", "")
+                tool_name = getattr(tool_call.function, "name", "")
+                tool_id = getattr(tool_call, "id", "")
                 tool_call_info = {
-                    "name": getattr(tool_call.function, "name", ""),
-                    "arguments": json.loads(getattr(tool_call.function, "arguments", "")),
-                    "tool_id": getattr(tool_call, "id", ""),
-                    "type": getattr(tool_call, "type", ""),
+                    "name": tool_name,
+                    "arguments": json.loads(tool_args),
+                    "tool_id": tool_id,
+                    "type": "function",
                 }
                 tool_calls_info.append(tool_call_info)
+                core.dispatch(
+                    DISPATCH_ON_LLM_TOOL_CHOICE,
+                    (
+                        tool_id,
+                        tool_name,
+                        tool_args,
+                        {
+                            "trace_id": format_trace_id(span.trace_id),
+                            "span_id": str(span.span_id),
+                        },
+                    ),
+                )
             if tool_calls_info:
                 output_messages.append({"content": content, "role": role, "tool_calls": tool_calls_info})
                 continue

ddtrace/llmobs/_integrations/openai_agents.py

@@ -6,15 +6,20 @@
 from typing import Union
 import weakref
 
+from ddtrace.internal import core
 from ddtrace.internal.logger import get_logger
 from ddtrace.internal.utils.formats import format_trace_id
+from ddtrace.llmobs._constants import DISPATCH_ON_LLM_TOOL_CHOICE
+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
 from ddtrace.llmobs._constants import INPUT_VALUE
 from ddtrace.llmobs._constants import METADATA
 from ddtrace.llmobs._constants import METRICS
 from ddtrace.llmobs._constants import MODEL_NAME
 from ddtrace.llmobs._constants import MODEL_PROVIDER
 from ddtrace.llmobs._constants import NAME
+from ddtrace.llmobs._constants import OAI_HANDOFF_TOOL_ARG
 from ddtrace.llmobs._constants import OUTPUT_MESSAGES
 from ddtrace.llmobs._constants import OUTPUT_VALUE
 from ddtrace.llmobs._constants import PARENT_ID_KEY
@@ -201,11 +206,30 @@ def _llmobs_set_response_attributes(self, span: Span, oai_span: OaiSpanAdapter)
             span._set_ctx_item(MODEL_PROVIDER, "openai")
 
         if oai_span.input:
-            messages = oai_span.llmobs_input_messages()
+            messages, tool_call_ids = oai_span.llmobs_input_messages()
+
+            for tool_call_id in tool_call_ids:
+                core.dispatch(DISPATCH_ON_TOOL_CALL_OUTPUT_USED, (tool_call_id, span))
+
             span._set_ctx_item(INPUT_MESSAGES, messages)
 
         if oai_span.response and oai_span.response.output:
-            messages = oai_span.llmobs_output_messages()
+            messages, tool_call_outputs = oai_span.llmobs_output_messages()
+
+            for tool_id, tool_name, tool_args in tool_call_outputs:
+                core.dispatch(
+                    DISPATCH_ON_LLM_TOOL_CHOICE,
+                    (
+                        tool_id,
+                        tool_name,
+                        tool_args,
+                        {
+                            "trace_id": format_trace_id(span.trace_id),
+                            "span_id": str(span.span_id),
+                        },
+                    ),
+                )
+
             span._set_ctx_item(OUTPUT_MESSAGES, messages)
 
         metadata = oai_span.llmobs_metadata
@@ -219,12 +243,20 @@ def _llmobs_set_response_attributes(self, span: Span, oai_span: OaiSpanAdapter)
     def _llmobs_set_tool_attributes(self, span: Span, oai_span: OaiSpanAdapter) -> None:
         span._set_ctx_item(INPUT_VALUE, oai_span.input or "")
         span._set_ctx_item(OUTPUT_VALUE, oai_span.output or "")
+        core.dispatch(
+            DISPATCH_ON_TOOL_CALL,
+            (oai_span.name, oai_span.input, "function", span),
+        )
 
     def _llmobs_set_handoff_attributes(self, span: Span, oai_span: OaiSpanAdapter) -> None:
         handoff_tool_name = "transfer_to_{}".format("_".join(oai_span.to_agent.split(" ")).lower())
         span.name = handoff_tool_name
         span._set_ctx_item("input_value", oai_span.from_agent or "")
         span._set_ctx_item("output_value", oai_span.to_agent or "")
+        core.dispatch(
+            DISPATCH_ON_TOOL_CALL,
+            (handoff_tool_name, OAI_HANDOFF_TOOL_ARG, "handoff", span),
+        )
 
     def _llmobs_set_agent_attributes(self, span: Span, oai_span: OaiSpanAdapter) -> None:
         if oai_span.llmobs_metadata:

ddtrace/llmobs/_integrations/utils.py

@@ -13,6 +13,7 @@
 
 from ddtrace.internal.logger import get_logger
 from ddtrace.llmobs._constants import INPUT_TOKENS_METRIC_KEY
+from ddtrace.llmobs._constants import OAI_HANDOFF_TOOL_ARG
 from ddtrace.llmobs._constants import OUTPUT_TOKENS_METRIC_KEY
 from ddtrace.llmobs._constants import TOTAL_TOKENS_METRIC_KEY
 from ddtrace.llmobs._utils import _get_attr
@@ -486,19 +487,22 @@ def get_error_data(self) -> Optional[Dict[str, Any]]:
             return None
         return self.error.get("data")
 
-    def llmobs_input_messages(self) -> List[Dict[str, Any]]:
+    def llmobs_input_messages(self) -> Tuple[List[Dict[str, Any]], List[str]]:
         """Returns processed input messages for LLM Obs LLM spans.
 
         Returns:
-            A list of processed messages
+            - A list of processed messages
+            - A list of tool call IDs for span linking purposes
         """
         messages = self.input
         processed: List[Dict[str, Any]] = []
+        tool_call_ids: List[str] = []
+
         if not messages:
-            return processed
+            return processed, tool_call_ids
 
         if isinstance(messages, str):
-            return [{"content": messages, "role": "user"}]
+            return [{"content": messages, "role": "user"}], tool_call_ids
 
         for item in messages:
             processed_item: Dict[str, Union[str, List[Dict[str, str]]]] = {}
@@ -541,6 +545,7 @@ def llmobs_input_messages(self) -> List[Dict[str, Any]]:
                         output = json.loads(output)
                     except json.JSONDecodeError:
                         output = {"output": output}
+                tool_call_ids.append(item["call_id"])
                 processed_item["tool_calls"] = [
                     {
                         "tool_id": item["call_id"],
@@ -550,21 +555,23 @@ def llmobs_input_messages(self) -> List[Dict[str, Any]]:
             if processed_item:
                 processed.append(processed_item)
 
-        return processed
+        return processed, tool_call_ids
 
-    def llmobs_output_messages(self) -> List[Dict[str, Any]]:
+    def llmobs_output_messages(self) -> Tuple[List[Dict[str, Any]], List[Tuple[str, str, str]]]:
         """Returns processed output messages for LLM Obs LLM spans.
 
         Returns:
-            A list of processed messages
+            - A list of processed messages
+            - A list of tool call data (name, id, args) for span linking purposes
         """
         if not self.response or not self.response.output:
-            return []
+            return [], []
 
         messages: List[Any] = self.response.output
         processed: List[Dict[str, Any]] = []
+        tool_call_outputs: List[Tuple[str, str, str]] = []
         if not messages:
-            return processed
+            return processed, tool_call_outputs
 
         if not isinstance(messages, list):
             messages = [messages]
@@ -581,6 +588,13 @@ def llmobs_output_messages(self) -> List[Dict[str, Any]]:
                 message.update({"role": getattr(item, "role", "assistant"), "content": text})
             # Handle tool calls
             elif hasattr(item, "call_id") and hasattr(item, "arguments"):
+                tool_call_outputs.append(
+                    (
+                        item.call_id,
+                        getattr(item, "name", ""),
+                        item.arguments if item.arguments else OAI_HANDOFF_TOOL_ARG,
+                    )
+                )
                 message.update(
                     {
                         "tool_calls": [
@@ -599,7 +613,7 @@ def llmobs_output_messages(self) -> List[Dict[str, Any]]:
                 message.update({"content": str(item)})
             processed.append(message)
 
-        return processed
+        return processed, tool_call_outputs
 
     def llmobs_trace_input(self) -> Optional[str]:
         """Converts Response span data to an input value for top level trace.
@@ -611,7 +625,7 @@ def llmobs_trace_input(self) -> Optional[str]:
             return None
 
         try:
-            messages = self.llmobs_input_messages()
+            messages, _ = self.llmobs_input_messages()
             if messages and len(messages) > 0:
                 return messages[0].get("content")
         except (AttributeError, IndexError):

ddtrace/llmobs/_llmobs.py

@@ -36,6 +36,9 @@
 from ddtrace.llmobs._constants import AGENTLESS_BASE_URL
 from ddtrace.llmobs._constants import ANNOTATIONS_CONTEXT_ID
 from ddtrace.llmobs._constants import DECORATOR
+from ddtrace.llmobs._constants import DISPATCH_ON_LLM_TOOL_CHOICE
+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_DOCUMENTS
 from ddtrace.llmobs._constants import INPUT_MESSAGES
 from ddtrace.llmobs._constants import INPUT_PROMPT
@@ -61,6 +64,7 @@
 from ddtrace.llmobs._evaluators.runner import EvaluatorRunner
 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_session_id
 from ddtrace.llmobs._utils import _get_span_name
@@ -122,6 +126,8 @@ def __init__(self, tracer=None):
         self._annotations = []
         self._annotation_context_lock = forksafe.RLock()
 
+        self._tool_call_tracker = ToolCallTracker()
+
     def _on_span_start(self, span):
         if self.enabled and span.span_type == SpanTypes.LLM:
             self._activate_llmobs_span(span)
@@ -303,6 +309,10 @@ 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)
+
         forksafe.unregister(self._child_after_fork)
 
     @classmethod
@@ -399,6 +409,10 @@ 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)
+
             atexit.register(cls.disable)
             telemetry_writer.product_activated(TELEMETRY_APM_PRODUCT.LLMOBS, True)