feat: add include_prompts for Pydantic AI integration

Author: constantinius

sentry_sdk/integrations/pydantic_ai/__init__.py

@@ -16,6 +16,18 @@
 
 class PydanticAIIntegration(Integration):
     identifier = "pydantic_ai"
+    origin = f"auto.ai.{identifier}"
+
+    def __init__(self, include_prompts=True):
+        # type: (bool) -> None
+        """
+        Initialize the Pydantic AI integration.
+
+        Args:
+            include_prompts: Whether to include prompts and messages in span data.
+                Requires send_default_pii=True. Defaults to True.
+        """
+        self.include_prompts = include_prompts
 
     @staticmethod
     def setup_once():

sentry_sdk/integrations/pydantic_ai/spans/execute_tool.py

@@ -1,10 +1,9 @@
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
-from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.utils import safe_serialize
 
 from ..consts import SPAN_ORIGIN
-from ..utils import _set_agent_data
+from ..utils import _set_agent_data, _should_send_prompts
 
 from typing import TYPE_CHECKING
 
@@ -27,7 +26,7 @@ def execute_tool_span(tool_name, tool_args, agent):
 
     _set_agent_data(span, agent)
 
-    if should_send_default_pii() and tool_args is not None:
+    if _should_send_prompts() and tool_args is not None:
         span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, safe_serialize(tool_args))
 
     return span
@@ -39,5 +38,5 @@ def update_execute_tool_span(span, result):
     if not span:
         return
 
-    if should_send_default_pii() and result is not None:
+    if _should_send_prompts() and result is not None:
         span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, safe_serialize(result))

sentry_sdk/integrations/pydantic_ai/spans/invoke_agent.py

@@ -1,10 +1,9 @@
 import sentry_sdk
 from sentry_sdk.ai.utils import get_start_span_function, set_data_normalized
 from sentry_sdk.consts import OP, SPANDATA
-from sentry_sdk.scope import should_send_default_pii
 
 from ..consts import SPAN_ORIGIN
-from ..utils import _set_agent_data, _set_model_data
+from ..utils import _set_agent_data, _set_model_data, _should_send_prompts
 
 from typing import TYPE_CHECKING
 
@@ -65,8 +64,8 @@ def invoke_agent_span(user_prompt, agent, model, model_settings):
             # If we can't extract tools, just skip it
             pass
 
-    # Add user prompt and system prompts if available and PII is allowed
-    if should_send_default_pii():
+    # Add user prompt and system prompts if available and prompts are enabled
+    if _should_send_prompts():
         messages = []
 
         # Add system prompts (both instructions and system_prompt)
@@ -135,7 +134,7 @@ def invoke_agent_span(user_prompt, agent, model, model_settings):
 def update_invoke_agent_span(span, output):
     # type: (sentry_sdk.tracing.Span, Any) -> None
     """Update and close the invoke agent span."""
-    if span and should_send_default_pii() and output:
+    if span and _should_send_prompts() and output:
         output_text = str(output) if not isinstance(output, str) else output
         set_data_normalized(
             span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_text, unpack=False

sentry_sdk/integrations/pydantic_ai/utils.py

@@ -19,6 +19,27 @@
     raise DidNotEnable("pydantic-ai not installed")
 
 
+def _should_send_prompts():
+    # type: () -> bool
+    """
+    Check if prompts should be sent to Sentry.
+
+    This checks both send_default_pii and the include_prompts integration setting.
+    """
+    if not should_send_default_pii():
+        return False
+
+    # Get the integration instance from the client
+    integration = sentry_sdk.get_client().get_integration(
+        pydantic_ai.__name__.split(".")[0]
+    )
+
+    if integration is None:
+        return False
+
+    return getattr(integration, "include_prompts", True)
+
+
 def _capture_exception(exc):
     # type: (Any) -> None
     set_span_errored()
@@ -154,7 +175,7 @@ def _set_usage_data(span, usage):
 def _set_input_messages(span, messages):
     # type: (sentry_sdk.tracing.Span, Any) -> None
     """Set input messages data on a span."""
-    if not should_send_default_pii():
+    if not _should_send_prompts():
         return
 
     if not messages:
@@ -247,7 +268,7 @@ def _set_input_messages(span, messages):
 def _set_output_data(span, response):
     # type: (sentry_sdk.tracing.Span, Any) -> None
     """Set output data on a span."""
-    if not should_send_default_pii():
+    if not _should_send_prompts():
         return
 
     if not response:

tests/integrations/pydantic_ai/__init__.py

@@ -0,0 +1,3 @@
+import pytest
+
+pytest.importorskip("pydantic_ai")