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")