handle pii

Author: antonpirker

sentry_sdk/integrations/openai_agents/patches/models.py

@@ -26,9 +26,9 @@ def wrapped_get_model(cls, agent, run_config):
 
         @wraps(original_get_response)
         async def wrapped_get_response(*args, **kwargs):
-            with ai_client_span(agent, model, run_config, kwargs) as span:
+            with ai_client_span(agent, kwargs) as span:
                 result = await original_get_response(*args, **kwargs)
-                update_ai_client_span(span, agent, model, run_config, kwargs, result)
+                update_ai_client_span(span, agent, kwargs, result)
 
             return result
 

sentry_sdk/integrations/openai_agents/spans/ai_client.py

@@ -6,12 +6,13 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from agents import Agent, Model, RunConfig
+    from agents import Agent
+    from typing import Any
     from sentry_sdk import Span
 
 
-def ai_client_span(agent, model, run_config, get_response_kwargs):
-    # type: (Agent, Model, RunConfig, dict) -> Span
+def ai_client_span(agent, get_response_kwargs):
+    # type: (Agent, dict[str, Any]) -> Span
     # TODO-anton: implement other types of operations. Now "chat" is hardcoded.
     span = sentry_sdk.start_span(
         op=OP.GEN_AI_CHAT,
@@ -23,7 +24,8 @@ def ai_client_span(agent, model, run_config, get_response_kwargs):
     return span
 
 
-def update_ai_client_span(span, agent, model, run_config, get_response_kwargs, result):
+def update_ai_client_span(span, agent, get_response_kwargs, result):
+    # type: (Span, Agent, dict[str, Any], Any) -> None
     _set_agent_data(span, agent)
     _set_usage_data(span, result.usage)
     _set_input_data(span, get_response_kwargs)

sentry_sdk/integrations/openai_agents/spans/execute_tool.py

@@ -1,5 +1,6 @@
 import sentry_sdk
 from sentry_sdk.consts import OP, SPANDATA
+from sentry_sdk.scope import should_send_default_pii
 
 from ..utils import _set_agent_data
 
@@ -18,19 +19,23 @@ def execute_tool_span(tool, *args, **kwargs):
     )
 
     span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "execute_tool")
-    span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool.name)
-    span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool.description)
-
-    input = args[1]
-    span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, input)
 
     if tool.__class__.__name__ == "FunctionTool":
         span.set_data(SPANDATA.GEN_AI_TOOL_TYPE, "function")
 
+    span.set_data(SPANDATA.GEN_AI_TOOL_NAME, tool.name)
+    span.set_data(SPANDATA.GEN_AI_TOOL_DESCRIPTION, tool.description)
+
+    if should_send_default_pii():
+        input = args[1]
+        span.set_data(SPANDATA.GEN_AI_TOOL_INPUT, input)
+
     return span
 
 
 def update_execute_tool_span(span, agent, tool, result):
     # type: (Span, agents.Agent, agents.Tool, Any) -> None
     _set_agent_data(span, agent)
-    span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, result)
+
+    if should_send_default_pii():
+        span.set_data(SPANDATA.GEN_AI_TOOL_OUTPUT, result)

sentry_sdk/integrations/openai_agents/utils.py

@@ -1,9 +1,10 @@
 from functools import wraps
 
 import sentry_sdk
+from sentry_sdk.consts import SPANDATA
 from sentry_sdk.integrations import DidNotEnable
+from sentry_sdk.scope import should_send_default_pii
 from sentry_sdk.utils import event_from_exception
-from sentry_sdk.consts import SPANDATA
 
 from typing import TYPE_CHECKING
 
@@ -154,6 +155,9 @@ def _set_usage_data(span, usage):
 
 def _set_input_data(span, get_response_kwargs):
     # type: (sentry_sdk.Span, dict[str, Any]) -> None
+    if not should_send_default_pii():
+        return
+
     messages_by_role = {
         "system": [],
         "user": [],
@@ -185,6 +189,9 @@ def _set_input_data(span, get_response_kwargs):
 
 def _set_output_data(span, result):
     # type: (sentry_sdk.Span, Any) -> None
+    if not should_send_default_pii():
+        return
+
     output_messages = {
         "response": [],
         "tool": [],