feat(tracing): Gate create_span gen_ai attributes via data_collection (#7295)

Route input/output attribute collection for AI_AGENT, AI_CHAT, and
AI_TOOL span templates through the `data_collection.gen_ai` option when
enabled, falling back to `send_default_pii` otherwise. Prompts recorded
via `prompt`/`system_prompt` kwargs stay ungated when `data_collection`
is not configured, preserving pre-existing behavior.

Fixes PY-2749
Fixes #7293

Author: ericapisani

sentry_sdk/tracing_utils.py

@@ -43,7 +43,16 @@
 
 if TYPE_CHECKING:
     from types import FrameType
-    from typing import Any, Dict, Generator, Iterator, Optional, Tuple, Union
+    from typing import (
+        Any,
+        Dict,
+        Generator,
+        Iterator,
+        Literal,
+        Optional,
+        Tuple,
+        Union,
+    )
 
     from sentry_sdk._types import Attributes
     from sentry_sdk.utils import ParsedUrl
@@ -1080,7 +1089,6 @@ def create_span_decorator(
         use cases.
     :type template: :py:class:`sentry_sdk.consts.SPANTEMPLATE`
     """
-    from sentry_sdk.scope import should_send_default_pii
 
     def span_decorator(f: "Any") -> "Any":
         """
@@ -1110,20 +1118,21 @@ async def async_wrapper(*args: "Any", **kwargs: "Any") -> "Any":
             span_op = op or _get_span_op(template)
             function_name = name or qualname_from_function(f) or ""
             span_name = _get_span_name(template, function_name, kwargs)
-            send_pii = should_send_default_pii()
+            collect_inputs = _should_collect_gen_ai("inputs")
+            collect_outputs = _should_collect_gen_ai("outputs")
 
             with current_span.start_child(
                 op=span_op,
                 name=span_name,
             ) as span:
                 span.update_data(attributes or {})
                 _set_input_attributes(
-                    span, template, send_pii, function_name, f, args, kwargs
+                    span, template, collect_inputs, function_name, f, args, kwargs
                 )
 
                 result = await f(*args, **kwargs)
 
-                _set_output_attributes(span, template, send_pii, result)
+                _set_output_attributes(span, template, collect_outputs, result)
 
                 return result
 
@@ -1155,20 +1164,21 @@ def sync_wrapper(*args: "Any", **kwargs: "Any") -> "Any":
             span_op = op or _get_span_op(template)
             function_name = name or qualname_from_function(f) or ""
             span_name = _get_span_name(template, function_name, kwargs)
-            send_pii = should_send_default_pii()
+            collect_inputs = _should_collect_gen_ai("inputs")
+            collect_outputs = _should_collect_gen_ai("outputs")
 
             with current_span.start_child(
                 op=span_op,
                 name=span_name,
             ) as span:
                 span.update_data(attributes or {})
                 _set_input_attributes(
-                    span, template, send_pii, function_name, f, args, kwargs
+                    span, template, collect_inputs, function_name, f, args, kwargs
                 )
 
                 result = f(*args, **kwargs)
 
-                _set_output_attributes(span, template, send_pii, result)
+                _set_output_attributes(span, template, collect_outputs, result)
 
                 return result
 
@@ -1373,9 +1383,22 @@ def _get_span_op(template: "Union[str, SPANTEMPLATE]") -> str:
     return str(op)
 
 
+_AI_TEMPLATES = frozenset(
+    {SPANTEMPLATE.AI_AGENT, SPANTEMPLATE.AI_CHAT, SPANTEMPLATE.AI_TOOL}
+)
+
+
+def _should_collect_gen_ai(kind: 'Literal["inputs", "outputs"]') -> bool:
+    client = sentry_sdk.get_client()
+    if has_data_collection_enabled(client.options):
+        return bool(client.options["data_collection"]["gen_ai"][kind])
+
+    return client.should_send_default_pii()
+
+
 def _get_input_attributes(
     template: "Union[str, SPANTEMPLATE]",
-    send_pii: bool,
+    collect_inputs: bool,
     args: "tuple[Any, ...]",
     kwargs: "dict[str, Any]",
 ) -> "dict[str, Any]":
@@ -1384,7 +1407,7 @@ def _get_input_attributes(
     """
     attributes: "dict[str, Any]" = {}
 
-    if template in [SPANTEMPLATE.AI_AGENT, SPANTEMPLATE.AI_TOOL, SPANTEMPLATE.AI_CHAT]:
+    if template in _AI_TEMPLATES:
         mapping = {
             "model": (SPANDATA.GEN_AI_REQUEST_MODEL, str),
             "model_name": (SPANDATA.GEN_AI_REQUEST_MODEL, str),
@@ -1404,22 +1427,25 @@ def _set_from_key(key: str, value: "Any") -> None:
                 if value is not None and isinstance(value, data_type):
                     attributes[attribute] = value
 
-        for key, value in list(kwargs.items()):
-            if key == "prompt" and isinstance(value, str):
-                attributes.setdefault(SPANDATA.GEN_AI_REQUEST_MESSAGES, []).append(
-                    {"role": "user", "content": value}
-                )
-                continue
+        # Pre-data collection, prompts were always recorded here, so they stay
+        # ungated until `send_default_pii` is removed.
+        collect_messages = True
+        if has_data_collection_enabled(sentry_sdk.get_client().options):
+            collect_messages = collect_inputs
 
-            if key == "system_prompt" and isinstance(value, str):
-                attributes.setdefault(SPANDATA.GEN_AI_REQUEST_MESSAGES, []).append(
-                    {"role": "system", "content": value}
-                )
+        roles = {"prompt": "user", "system_prompt": "system"}
+
+        for key, value in list(kwargs.items()):
+            if key in roles:
+                if collect_messages and isinstance(value, str):
+                    attributes.setdefault(SPANDATA.GEN_AI_REQUEST_MESSAGES, []).append(
+                        {"role": roles[key], "content": value}
+                    )
                 continue
 
             _set_from_key(key, value)
 
-    if template == SPANTEMPLATE.AI_TOOL and send_pii:
+    if template == SPANTEMPLATE.AI_TOOL and collect_inputs:
         attributes[SPANDATA.GEN_AI_TOOL_INPUT] = safe_repr(
             {"args": args, "kwargs": kwargs}
         )
@@ -1462,14 +1488,14 @@ def _set_from_keys(attribute: str, keys: "tuple[str, ...]") -> None:
 
 
 def _get_output_attributes(
-    template: "Union[str, SPANTEMPLATE]", send_pii: bool, result: "Any"
+    template: "Union[str, SPANTEMPLATE]", collect_outputs: bool, result: "Any"
 ) -> "dict[str, Any]":
     """
     Get output attributes for the given span template.
     """
     attributes: "dict[str, Any]" = {}
 
-    if template in [SPANTEMPLATE.AI_AGENT, SPANTEMPLATE.AI_TOOL, SPANTEMPLATE.AI_CHAT]:
+    if template in _AI_TEMPLATES:
         with capture_internal_exceptions():
             # Usage from result, result.usage, and result.metadata.usage
             usage_candidates = [result]
@@ -1495,7 +1521,7 @@ def _get_output_attributes(
                 attributes[SPANDATA.GEN_AI_RESPONSE_MODEL] = model_name
 
     # Tool output
-    if template == SPANTEMPLATE.AI_TOOL and send_pii:
+    if template == SPANTEMPLATE.AI_TOOL and collect_outputs:
         attributes[SPANDATA.GEN_AI_TOOL_OUTPUT] = safe_repr(result)
 
     return attributes
@@ -1504,7 +1530,7 @@ def _get_output_attributes(
 def _set_input_attributes(
     span: "Span",
     template: "Union[str, SPANTEMPLATE]",
-    send_pii: bool,
+    collect_inputs: bool,
     name: str,
     f: "Any",
     args: "tuple[Any, ...]",
@@ -1515,7 +1541,7 @@ def _set_input_attributes(
 
     :param span: The span to set attributes on.
     :param template: The template to use to set attributes on the span.
-    :param send_pii: Whether to send PII data.
+    :param collect_inputs: Whether gen_ai inputs may be collected.
     :param f: The wrapped function.
     :param args: The arguments to the wrapped function.
     :param kwargs: The keyword arguments to the wrapped function.
@@ -1541,22 +1567,25 @@ def _set_input_attributes(
         if docstring is not None:
             attributes[SPANDATA.GEN_AI_TOOL_DESCRIPTION] = docstring
 
-    attributes.update(_get_input_attributes(template, send_pii, args, kwargs))
+    attributes.update(_get_input_attributes(template, collect_inputs, args, kwargs))
     span.update_data(attributes or {})
 
 
 def _set_output_attributes(
-    span: "Span", template: "Union[str, SPANTEMPLATE]", send_pii: bool, result: "Any"
+    span: "Span",
+    template: "Union[str, SPANTEMPLATE]",
+    collect_outputs: bool,
+    result: "Any",
 ) -> None:
     """
     Set span output attributes based on the given span template.
 
     :param span: The span to set attributes on.
     :param template: The template to use to set attributes on the span.
-    :param send_pii: Whether to send PII data.
+    :param collect_outputs: Whether gen_ai outputs may be collected.
     :param result: The result of the wrapped function.
     """
-    span.update_data(_get_output_attributes(template, send_pii, result) or {})
+    span.update_data(_get_output_attributes(template, collect_outputs, result) or {})
 
 
 def _should_continue_trace(baggage: "Optional[Baggage]") -> bool:

tests/tracing/test_decorator.py

@@ -814,32 +814,112 @@ def my_agent(*args, **kwargs):
         with sentry_sdk.start_transaction(name="test-transaction"):
             my_agent(22, 33, arg1=44, arg2=55)
 
-        (_, tool_span, _) = (item.payload for item in items)
+        (_, tool_span, chat_span) = (item.payload for item in items)
+        tool_data = tool_span["attributes"]
+        chat_data = chat_span["attributes"]
+    else:
+        events = capture_events()
 
-        if send_default_pii:
-            assert (
-                tool_span["attributes"]["gen_ai.tool.input"]
-                == "{'args': (1, 2), 'kwargs': {'tool_arg1': '3', 'tool_arg2': '4'}}"
-            )
-            assert tool_span["attributes"]["gen_ai.tool.output"] == "'tool_output'"
-        else:
-            assert "gen_ai.tool.input" not in tool_span["attributes"]
-            assert "gen_ai.tool.output" not in tool_span["attributes"]
+        with sentry_sdk.start_transaction(name="test-transaction"):
+            my_agent(22, 33, arg1=44, arg2=55)
+
+        (event,) = events
+        (_, tool_span, chat_span) = event["spans"]
+        tool_data = tool_span["data"]
+        chat_data = chat_span["data"]
+
+    if send_default_pii:
+        assert (
+            tool_data["gen_ai.tool.input"]
+            == "{'args': (1, 2), 'kwargs': {'tool_arg1': '3', 'tool_arg2': '4'}}"
+        )
+        assert tool_data["gen_ai.tool.output"] == "'tool_output'"
+    else:
+        assert "gen_ai.tool.input" not in tool_data
+        assert "gen_ai.tool.output" not in tool_data
+
+    # Without `data_collection`, prompts are recorded regardless of `send_default_pii`.
+    assert chat_data["gen_ai.request.messages"] == (
+        "[{'role': 'user', 'content': 'What is the weather in Tokyo?'}, "
+        "{'role': 'system', 'content': 'You are a helpful assistant that can answer "
+        "questions about the weather.'}]"
+    )
+
+
+@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
+@pytest.mark.parametrize("collect_outputs", [True, False])
+@pytest.mark.parametrize("collect_inputs", [True, False])
+def test_span_templates_ai_data_collection(
+    sentry_init,
+    capture_events,
+    capture_items,
+    collect_inputs,
+    collect_outputs,
+    stream_gen_ai_spans,
+):
+    @sentry_sdk.trace(template=SPANTEMPLATE.AI_TOOL)
+    def my_tool(arg1, arg2, **kwargs):
+        """This is a tool function."""
+        return "tool_output"
+
+    @sentry_sdk.trace(template=SPANTEMPLATE.AI_CHAT)
+    def my_chat(model=None, **kwargs):
+        return "chat_output"
+
+    @sentry_sdk.trace(template=SPANTEMPLATE.AI_AGENT)
+    def my_agent(*args, **kwargs):
+        my_tool(1, 2, tool_arg1="3", tool_arg2="4")
+        my_chat(
+            model="my-gpt-4o-mini",
+            prompt="What is the weather in Tokyo?",
+            system_prompt="You are a helpful assistant.",
+        )
+        return "agent_output"
+
+    sentry_init(
+        traces_sample_rate=1.0,
+        stream_gen_ai_spans=stream_gen_ai_spans,
+        _experiments={
+            "data_collection": {
+                "gen_ai": {"inputs": collect_inputs, "outputs": collect_outputs}
+            }
+        },
+    )
+
+    if stream_gen_ai_spans:
+        items = capture_items("span")
+
+        with sentry_sdk.start_transaction(name="test-transaction"):
+            my_agent(22, 33, arg1=44, arg2=55)
+
+        (_, tool_span, chat_span) = (item.payload for item in items)
+        tool_data = tool_span["attributes"]
+        chat_data = chat_span["attributes"]
     else:
         events = capture_events()
 
         with sentry_sdk.start_transaction(name="test-transaction"):
             my_agent(22, 33, arg1=44, arg2=55)
 
         (event,) = events
-        (_, tool_span, _) = event["spans"]
+        (_, tool_span, chat_span) = event["spans"]
+        tool_data = tool_span["data"]
+        chat_data = chat_span["data"]
 
-        if send_default_pii:
-            assert (
-                tool_span["data"]["gen_ai.tool.input"]
-                == "{'args': (1, 2), 'kwargs': {'tool_arg1': '3', 'tool_arg2': '4'}}"
-            )
-            assert tool_span["data"]["gen_ai.tool.output"] == "'tool_output'"
-        else:
-            assert "gen_ai.tool.input" not in tool_span["data"]
-            assert "gen_ai.tool.output" not in tool_span["data"]
+    if collect_inputs:
+        assert (
+            tool_data["gen_ai.tool.input"]
+            == "{'args': (1, 2), 'kwargs': {'tool_arg1': '3', 'tool_arg2': '4'}}"
+        )
+        assert chat_data["gen_ai.request.messages"] == (
+            "[{'role': 'user', 'content': 'What is the weather in Tokyo?'}, "
+            "{'role': 'system', 'content': 'You are a helpful assistant.'}]"
+        )
+    else:
+        assert "gen_ai.tool.input" not in tool_data
+        assert "gen_ai.request.messages" not in chat_data
+
+    if collect_outputs:
+        assert tool_data["gen_ai.tool.output"] == "'tool_output'"
+    else:
+        assert "gen_ai.tool.output" not in tool_data