feat(data-collection): split generative-AI input/output collection

sentry_sdk/integrations/anthropic.py

@@ -16,7 +16,10 @@
 )
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration, _check_minimum_version
-from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.scope import (
+    should_collect_gen_ai_inputs,
+    should_collect_gen_ai_outputs,
+)
 from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing import Span
 from sentry_sdk.tracing_utils import (
@@ -146,7 +149,9 @@ class AnthropicIntegration(Integration):
     identifier = "anthropic"
     origin = f"auto.ai.{identifier}"
 
-    def __init__(self: "AnthropicIntegration", include_prompts: bool = True) -> None:
+    def __init__(
+        self: "AnthropicIntegration", include_prompts: "Optional[bool]" = None
+    ) -> None:
         self.include_prompts = include_prompts
 
     @staticmethod
@@ -393,8 +398,7 @@ def _set_common_input_data(
     if (
         messages is not None
         and len(messages) > 0  # type: ignore
-        and should_send_default_pii()
-        and integration.include_prompts
+        and should_collect_gen_ai_inputs(integration.include_prompts)
     ):
         if isinstance(system, str) or isinstance(system, Iterable):
             set_on_span(
@@ -586,7 +590,7 @@ def _set_output_data(
         set_on_span(SPANDATA.GEN_AI_RESPONSE_ID, response_id)
     if finish_reason is not None:
         set_on_span(SPANDATA.GEN_AI_RESPONSE_FINISH_REASONS, [finish_reason])
-    if should_send_default_pii() and integration.include_prompts:
+    if should_collect_gen_ai_outputs(integration.include_prompts):
         output_messages: "dict[str, list[Any]]" = {
             "response": [],
             "tool": [],

sentry_sdk/integrations/cohere.py

@@ -10,13 +10,16 @@
 from sentry_sdk.tracing_utils import has_span_streaming_enabled
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Iterator, Union
+    from typing import Any, Callable, Iterator, Optional, Union
 
     from sentry_sdk.tracing import Span
 
 import sentry_sdk
 from sentry_sdk.integrations import DidNotEnable, Integration
-from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.scope import (
+    should_collect_gen_ai_inputs,
+    should_collect_gen_ai_outputs,
+)
 from sentry_sdk.utils import capture_internal_exceptions, event_from_exception, reraise
 
 try:
@@ -73,7 +76,9 @@ class CohereIntegration(Integration):
     identifier = "cohere"
     origin = f"auto.ai.{identifier}"
 
-    def __init__(self: "CohereIntegration", include_prompts: bool = True) -> None:
+    def __init__(
+        self: "CohereIntegration", include_prompts: "Optional[bool]" = None
+    ) -> None:
         self.include_prompts = include_prompts
 
     @staticmethod
@@ -179,7 +184,7 @@ def new_chat(*args: "Any", **kwargs: "Any") -> "Any":
             reraise(*exc_info)
 
         with capture_internal_exceptions():
-            if should_send_default_pii() and integration.include_prompts:
+            if should_collect_gen_ai_inputs(integration.include_prompts):
                 set_data_normalized(
                     span,
                     SPANDATA.AI_INPUT_MESSAGES,
@@ -215,8 +220,9 @@ def new_iterator() -> "Iterator[StreamedChatResponse]":
                                 collect_chat_response_fields(
                                     span,
                                     x.response,
-                                    include_pii=should_send_default_pii()
-                                    and integration.include_prompts,
+                                    include_pii=should_collect_gen_ai_outputs(
+                                        integration.include_prompts
+                                    ),
                                 )
                             yield x
                     _end_span(span)
@@ -226,8 +232,9 @@ def new_iterator() -> "Iterator[StreamedChatResponse]":
                 collect_chat_response_fields(
                     span,
                     res,
-                    include_pii=should_send_default_pii()
-                    and integration.include_prompts,
+                    include_pii=should_collect_gen_ai_outputs(
+                        integration.include_prompts
+                    ),
                 )
                 _end_span(span)
             else:
@@ -265,8 +272,8 @@ def new_embed(*args: "Any", **kwargs: "Any") -> "Any":
             )
 
         with span_ctx as span:
-            if "texts" in kwargs and (
-                should_send_default_pii() and integration.include_prompts
+            if "texts" in kwargs and should_collect_gen_ai_inputs(
+                integration.include_prompts
             ):
                 if isinstance(kwargs["texts"], str):
                     set_data_normalized(span, SPANDATA.AI_TEXTS, [kwargs["texts"]])

sentry_sdk/integrations/google_genai/__init__.py

@@ -5,6 +5,7 @@
     Callable,
     Iterator,
     List,
+    Optional,
 )
 
 import sentry_sdk
@@ -41,7 +42,10 @@ class GoogleGenAIIntegration(Integration):
     identifier = IDENTIFIER
     origin = ORIGIN
 
-    def __init__(self: "GoogleGenAIIntegration", include_prompts: bool = True) -> None:
+    def __init__(
+        self: "GoogleGenAIIntegration",
+        include_prompts: "Optional[bool]" = None,
+    ) -> None:
         self.include_prompts = include_prompts
 
     @staticmethod

sentry_sdk/integrations/google_genai/streaming.py

@@ -2,7 +2,7 @@
 
 from sentry_sdk.ai.utils import set_data_normalized
 from sentry_sdk.consts import SPANDATA
-from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.scope import should_collect_gen_ai_outputs
 from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.utils import (
     safe_serialize,
@@ -107,10 +107,8 @@ def set_span_data_for_streaming_response(
         span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
     )
 
-    if (
-        should_send_default_pii()
-        and integration.include_prompts
-        and accumulated_response.get("text")
+    if should_collect_gen_ai_outputs(integration.include_prompts) and (
+        accumulated_response.get("text")
     ):
         set_on_span(
             SPANDATA.GEN_AI_RESPONSE_TEXT,

sentry_sdk/integrations/google_genai/utils.py

@@ -27,7 +27,10 @@
     truncate_and_annotate_messages,
 )
 from sentry_sdk.consts import OP, SPANDATA
-from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.scope import (
+    should_collect_gen_ai_inputs,
+    should_collect_gen_ai_outputs,
+)
 from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing_utils import (
     has_span_streaming_enabled,
@@ -899,7 +902,7 @@ def set_span_data_for_request(
     config: "Optional[GenerateContentConfig]" = kwargs.get("config")
 
     # Set input messages/prompts if PII is allowed
-    if should_send_default_pii() and integration.include_prompts:
+    if should_collect_gen_ai_inputs(integration.include_prompts):
         messages = []
 
         # Add system instruction if present
@@ -977,7 +980,7 @@ def set_span_data_for_response(
     set_on_span = (
         span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
     )
-    if should_send_default_pii() and integration.include_prompts:
+    if should_collect_gen_ai_outputs(integration.include_prompts):
         response_texts = _extract_response_text(response)
         if response_texts:
             # Format as JSON string array as per documentation
@@ -1063,7 +1066,7 @@ def set_span_data_for_embed_request(
 ) -> None:
     """Set span data for embedding request."""
     # Include input contents if PII is allowed
-    if should_send_default_pii() and integration.include_prompts:
+    if should_collect_gen_ai_inputs(integration.include_prompts):
         if contents:
             # For embeddings, contents is typically a list of strings/texts
             input_texts = []

sentry_sdk/integrations/huggingface_hub.py

@@ -12,7 +12,10 @@
 )
 from sentry_sdk.consts import OP, SPANDATA
 from sentry_sdk.integrations import DidNotEnable, Integration
-from sentry_sdk.scope import should_send_default_pii
+from sentry_sdk.scope import (
+    should_collect_gen_ai_inputs,
+    should_collect_gen_ai_outputs,
+)
 from sentry_sdk.traces import StreamedSpan
 from sentry_sdk.tracing_utils import has_span_streaming_enabled
 from sentry_sdk.utils import (
@@ -22,7 +25,7 @@
 )
 
 if TYPE_CHECKING:
-    from typing import Any, Callable, Iterable, Union
+    from typing import Any, Callable, Iterable, Optional, Union
 
     from sentry_sdk.tracing import Span
 
@@ -37,7 +40,8 @@ class HuggingfaceHubIntegration(Integration):
     origin = f"auto.ai.{identifier}"
 
     def __init__(
-        self: "HuggingfaceHubIntegration", include_prompts: bool = True
+        self: "HuggingfaceHubIntegration",
+        include_prompts: "Optional[bool]" = None,
     ) -> None:
         self.include_prompts = include_prompts
 
@@ -114,7 +118,7 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
             _set_span_data_attribute(span, SPANDATA.GEN_AI_REQUEST_MODEL, model)
 
         # Input attributes
-        if should_send_default_pii() and integration.include_prompts:
+        if should_collect_gen_ai_inputs(integration.include_prompts):
             set_data_normalized(
                 span, SPANDATA.GEN_AI_REQUEST_MESSAGES, prompt, unpack=False
             )
@@ -206,7 +210,7 @@ def new_huggingface_task(*args: "Any", **kwargs: "Any") -> "Any":
                     finish_reason,
                 )
 
-            if should_send_default_pii() and integration.include_prompts:
+            if should_collect_gen_ai_outputs(integration.include_prompts):
                 if tool_calls is not None and len(tool_calls) > 0:
                     set_data_normalized(
                         span,
@@ -280,7 +284,7 @@ def new_details_iterator() -> "Iterable[Any]":
                                 finish_reason,
                             )
 
-                        if should_send_default_pii() and integration.include_prompts:
+                        if should_collect_gen_ai_outputs(integration.include_prompts):
                             if len(response_text_buffer) > 0:
                                 text_response = "".join(response_text_buffer)
                                 if text_response:
@@ -359,7 +363,7 @@ def new_iterator() -> "Iterable[str]":
                                 finish_reason,
                             )
 
-                        if should_send_default_pii() and integration.include_prompts:
+                        if should_collect_gen_ai_outputs(integration.include_prompts):
                             if tool_calls is not None and len(tool_calls) > 0:
                                 set_data_normalized(
                                     span,