Skip to content

Commit bff8521

Browse files
authored
feat(openai_agents): Gate tool execution span data on data_collection options (#7123)
Respect the data_collection configuration for gen_ai.tool.input and gen_ai.tool.output attributes in execute_tool spans, falling back to send_default_pii for backwards compatibility. This brings tool execution spans in line with other gen_ai span data collection gates. Refs PY-2588
1 parent c7e9e62 commit bff8521

4 files changed

Lines changed: 279 additions & 4 deletions

File tree

sentry_sdk/integrations/openai_agents/patches/runner.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,11 @@
66
from sentry_sdk.integrations import DidNotEnable
77
from sentry_sdk.scope import should_send_default_pii
88
from sentry_sdk.traces import StreamedSpan
9-
from sentry_sdk.utils import capture_internal_exceptions, reraise
9+
from sentry_sdk.utils import (
10+
capture_internal_exceptions,
11+
has_data_collection_enabled,
12+
reraise,
13+
)
1014

1115
from ..spans import (
1216
agent_workflow_span,
@@ -53,7 +57,11 @@ async def on_tool_start(
5357
span.__enter__()
5458
context._sentry_execute_tool_span = span
5559

56-
if not should_send_default_pii():
60+
client = sentry_sdk.get_client()
61+
if has_data_collection_enabled(client.options):
62+
if not client.options["data_collection"]["gen_ai"]["inputs"]:
63+
return
64+
elif not should_send_default_pii():
5765
return
5866

5967
if isinstance(span, StreamedSpan):

sentry_sdk/integrations/openai_agents/patches/tools.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
from functools import wraps
22
from typing import TYPE_CHECKING
33

4+
import sentry_sdk
45
from sentry_sdk.consts import SPANDATA
56
from sentry_sdk.integrations import DidNotEnable
67
from sentry_sdk.scope import should_send_default_pii
78
from sentry_sdk.traces import StreamedSpan
9+
from sentry_sdk.utils import has_data_collection_enabled
810

911
from ..spans import execute_tool_span, update_execute_tool_span
1012

@@ -56,7 +58,11 @@ async def sentry_wrapped_on_invoke_tool(
5658
result = await current_on_invoke(*args, **kwargs)
5759
update_execute_tool_span(span, agent, current_tool, result)
5860

59-
if not should_send_default_pii():
61+
client = sentry_sdk.get_client()
62+
if has_data_collection_enabled(client.options):
63+
if not client.options["data_collection"]["gen_ai"]["inputs"]:
64+
return result
65+
elif not should_send_default_pii():
6066
return result
6167

6268
if isinstance(span, StreamedSpan):

sentry_sdk/integrations/openai_agents/spans/execute_tool.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
from sentry_sdk.scope import should_send_default_pii
66
from sentry_sdk.traces import SpanStatus, StreamedSpan
77
from sentry_sdk.tracing_utils import has_span_streaming_enabled
8+
from sentry_sdk.utils import has_data_collection_enabled
89

910
from ..consts import SPAN_ORIGIN
1011
from ..utils import _set_agent_data
@@ -19,6 +20,7 @@ def execute_tool_span(
1920
tool: "agents.Tool", *args: "Any", **kwargs: "Any"
2021
) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
2122
span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
23+
2224
if span_streaming:
2325
span = sentry_sdk.traces.start_span(
2426
name=f"execute_tool {tool.name}",
@@ -51,6 +53,8 @@ def update_execute_tool_span(
5153
tool: "agents.Tool",
5254
result: "Any",
5355
) -> None:
56+
client = sentry_sdk.get_client()
57+
5458
_set_agent_data(span, agent)
5559

5660
if isinstance(result, str) and result.startswith(
@@ -65,7 +69,10 @@ def update_execute_tool_span(
6569
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
6670
)
6771

68-
if should_send_default_pii():
72+
if has_data_collection_enabled(client.options):
73+
if client.options["data_collection"]["gen_ai"]["outputs"]:
74+
set_on_span(SPANDATA.GEN_AI_TOOL_OUTPUT, result)
75+
elif should_send_default_pii():
6976
set_on_span(SPANDATA.GEN_AI_TOOL_OUTPUT, result)
7077

7178
# Add conversation ID from agent

tests/integrations/openai_agents/test_openai_agents.py

Lines changed: 254 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3071,6 +3071,260 @@ def simple_test_tool(message: str) -> str:
30713071
assert tool_span["data"]["gen_ai.tool.output"] == "Tool executed with: hello"
30723072

30733073

3074+
@pytest.fixture
3075+
def run_tool_agent(
3076+
sentry_init,
3077+
capture_events,
3078+
capture_items,
3079+
test_agent,
3080+
get_model_response,
3081+
nonstreaming_responses_tool_call_model_responses,
3082+
):
3083+
async def inner(tool, span_streaming, run_kwargs=None, **init_kwargs):
3084+
client = AsyncOpenAI(api_key="test-key")
3085+
model = OpenAIResponsesModel(model="gpt-4", openai_client=client)
3086+
agent_with_tool = test_agent.clone(tools=[tool], model=model)
3087+
3088+
responses = nonstreaming_responses_tool_call_model_responses(
3089+
tool_name=tool.name,
3090+
arguments='{"message": "hello"}',
3091+
response_model="gpt-4",
3092+
response_text="Task completed using the tool",
3093+
response_ids=iter(["resp_tool_123", "resp_final_123"]),
3094+
usages=iter(
3095+
[
3096+
ResponseUsage(
3097+
input_tokens=10,
3098+
input_tokens_details=InputTokensDetails(
3099+
cached_tokens=0,
3100+
cache_write_tokens=0,
3101+
),
3102+
output_tokens=5,
3103+
output_tokens_details=OutputTokensDetails(
3104+
reasoning_tokens=0,
3105+
),
3106+
total_tokens=15,
3107+
),
3108+
ResponseUsage(
3109+
input_tokens=15,
3110+
input_tokens_details=InputTokensDetails(
3111+
cached_tokens=0,
3112+
cache_write_tokens=0,
3113+
),
3114+
output_tokens=10,
3115+
output_tokens_details=OutputTokensDetails(
3116+
reasoning_tokens=0,
3117+
),
3118+
total_tokens=25,
3119+
),
3120+
]
3121+
),
3122+
)
3123+
tool_response = get_model_response(
3124+
next(responses),
3125+
serialize_pydantic=True,
3126+
)
3127+
final_response = get_model_response(
3128+
next(responses),
3129+
serialize_pydantic=True,
3130+
)
3131+
3132+
with patch.object(
3133+
agent_with_tool.model._client._client,
3134+
"send",
3135+
side_effect=[tool_response, final_response],
3136+
) as _:
3137+
sentry_init(
3138+
integrations=[OpenAIAgentsIntegration()],
3139+
disabled_integrations=[StdlibIntegration],
3140+
traces_sample_rate=1.0,
3141+
stream_gen_ai_spans=span_streaming,
3142+
trace_lifecycle="stream" if span_streaming else "static",
3143+
**init_kwargs,
3144+
)
3145+
3146+
items = capture_items("span") if span_streaming else None
3147+
events = None if span_streaming else capture_events()
3148+
3149+
await agents.Runner.run(
3150+
agent_with_tool,
3151+
"Please use the tool",
3152+
run_config=test_run_config,
3153+
**(run_kwargs or {}),
3154+
)
3155+
3156+
if span_streaming:
3157+
sentry_sdk.flush()
3158+
tool_span = next(
3159+
item.payload
3160+
for item in items
3161+
if item.payload["attributes"].get("sentry.op") == OP.GEN_AI_EXECUTE_TOOL
3162+
)
3163+
return tool_span, tool_span["attributes"]
3164+
3165+
(transaction,) = events
3166+
tool_span = next(
3167+
span
3168+
for span in transaction["spans"]
3169+
if span["op"] == OP.GEN_AI_EXECUTE_TOOL
3170+
)
3171+
return tool_span, tool_span["data"]
3172+
3173+
return inner
3174+
3175+
3176+
@pytest.fixture
3177+
def simple_test_tool():
3178+
@agents.function_tool
3179+
def simple_test_tool(message: str) -> str:
3180+
"""A simple tool"""
3181+
return f"Tool executed with: {message}"
3182+
3183+
return simple_test_tool
3184+
3185+
3186+
@pytest.mark.parametrize("span_streaming", [True, False])
3187+
@pytest.mark.parametrize(
3188+
"data_collection,send_default_pii,expect_input,expect_output",
3189+
[
3190+
pytest.param(
3191+
{"gen_ai": {"inputs": True, "outputs": True}},
3192+
False,
3193+
True,
3194+
True,
3195+
id="gen-ai-inputs-and-outputs-enabled-overrides-pii-disabled",
3196+
),
3197+
pytest.param(
3198+
{"gen_ai": {"inputs": False, "outputs": False}},
3199+
True,
3200+
False,
3201+
False,
3202+
id="gen-ai-inputs-and-outputs-disabled-overrides-pii-enabled",
3203+
),
3204+
pytest.param(
3205+
{"gen_ai": {"inputs": True, "outputs": False}},
3206+
False,
3207+
True,
3208+
False,
3209+
id="gen-ai-only-inputs-enabled",
3210+
),
3211+
pytest.param(
3212+
{"gen_ai": {"inputs": False, "outputs": True}},
3213+
False,
3214+
False,
3215+
True,
3216+
id="gen-ai-only-outputs-enabled",
3217+
),
3218+
pytest.param(
3219+
{},
3220+
False,
3221+
True,
3222+
True,
3223+
id="gen-ai-omitted-defaults-to-enabled",
3224+
),
3225+
pytest.param(
3226+
{"gen_ai": {"inputs": False, "outputs": False}},
3227+
False,
3228+
False,
3229+
False,
3230+
id="gen-ai-inputs-and-outputs-disabled-and-pii-disabled",
3231+
),
3232+
pytest.param(
3233+
None,
3234+
False,
3235+
False,
3236+
False,
3237+
id="no-gen-ai-data-collection-falls-back-to-send-default-pii",
3238+
),
3239+
pytest.param(
3240+
None,
3241+
True,
3242+
True,
3243+
True,
3244+
id="no-gen-ai-data-collection-pii-enabled-collects",
3245+
),
3246+
],
3247+
)
3248+
@pytest.mark.asyncio
3249+
async def test_tool_execution_span_data_collection(
3250+
run_tool_agent,
3251+
simple_test_tool,
3252+
data_collection,
3253+
send_default_pii,
3254+
expect_input,
3255+
expect_output,
3256+
span_streaming,
3257+
):
3258+
init_kwargs = {"send_default_pii": send_default_pii}
3259+
if data_collection is not None:
3260+
init_kwargs["_experiments"] = {"data_collection": data_collection}
3261+
3262+
_, tool_span_data = await run_tool_agent(
3263+
simple_test_tool,
3264+
span_streaming,
3265+
**init_kwargs,
3266+
)
3267+
3268+
if expect_input:
3269+
assert tool_span_data[SPANDATA.GEN_AI_TOOL_INPUT] == '{"message": "hello"}'
3270+
else:
3271+
assert SPANDATA.GEN_AI_TOOL_INPUT not in tool_span_data
3272+
3273+
if expect_output:
3274+
assert (
3275+
tool_span_data[SPANDATA.GEN_AI_TOOL_OUTPUT] == "Tool executed with: hello"
3276+
)
3277+
else:
3278+
assert SPANDATA.GEN_AI_TOOL_OUTPUT not in tool_span_data
3279+
3280+
3281+
@pytest.mark.parametrize("span_streaming", [True, False])
3282+
@pytest.mark.asyncio
3283+
async def test_tool_execution_error_data_collection(
3284+
run_tool_agent,
3285+
span_streaming,
3286+
):
3287+
@agents.function_tool
3288+
def failing_tool(message: str) -> str:
3289+
"""A tool that fails"""
3290+
raise ValueError("Tool execution failed")
3291+
3292+
tool_span, tool_span_data = await run_tool_agent(
3293+
failing_tool,
3294+
span_streaming,
3295+
_experiments={"data_collection": {"gen_ai": {"outputs": False}}},
3296+
)
3297+
3298+
assert tool_span_data[SPANDATA.GEN_AI_TOOL_NAME] == "failing_tool"
3299+
assert tool_span["status"] == ("error" if span_streaming else "internal_error")
3300+
assert SPANDATA.GEN_AI_TOOL_OUTPUT not in tool_span_data
3301+
3302+
3303+
@pytest.mark.parametrize("span_streaming", [True, False])
3304+
@pytest.mark.skipif(
3305+
parse_version(OPENAI_AGENTS_VERSION) < (0, 4, 0),
3306+
reason="conversation_id support requires openai-agents >= 0.4.0",
3307+
)
3308+
@pytest.mark.asyncio
3309+
async def test_tool_execution_span_non_pii_data_always_set(
3310+
run_tool_agent,
3311+
simple_test_tool,
3312+
span_streaming,
3313+
):
3314+
_, tool_span_data = await run_tool_agent(
3315+
simple_test_tool,
3316+
span_streaming,
3317+
run_kwargs={"conversation_id": "conv_tool_test_456"},
3318+
_experiments={
3319+
"data_collection": {"gen_ai": {"inputs": False, "outputs": False}}
3320+
},
3321+
)
3322+
3323+
assert tool_span_data[SPANDATA.GEN_AI_TOOL_NAME] == "simple_test_tool"
3324+
assert tool_span_data[SPANDATA.GEN_AI_TOOL_DESCRIPTION] == "A simple tool"
3325+
assert tool_span_data[SPANDATA.GEN_AI_CONVERSATION_ID] == "conv_tool_test_456"
3326+
3327+
30743328
@pytest.mark.asyncio
30753329
async def test_hosted_mcp_tool_propagation_header_streamed(
30763330
sentry_init,

0 commit comments

Comments
 (0)