diff --git a/sentry_sdk/integrations/openai_agents/spans/ai_client.py b/sentry_sdk/integrations/openai_agents/spans/ai_client.py index 6ab1b495c1..d075d3c9a8 100644 --- a/sentry_sdk/integrations/openai_agents/spans/ai_client.py +++ b/sentry_sdk/integrations/openai_agents/spans/ai_client.py @@ -6,6 +6,7 @@ from sentry_sdk.integrations import DidNotEnable from sentry_sdk.traces import StreamedSpan from sentry_sdk.tracing_utils import has_span_streaming_enabled +from sentry_sdk.utils import has_data_collection_enabled try: from agents import ( @@ -207,7 +208,9 @@ def ai_client_span( elif hasattr(agent, "_sentry_request_model"): model_name = agent._sentry_request_model - span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options) + client_options = sentry_sdk.get_client().options + + span_streaming = has_span_streaming_enabled(client_options) if span_streaming: span = sentry_sdk.traces.start_span( name=f"chat {model_name}", @@ -232,10 +235,19 @@ def ai_client_span( _set_agent_data(span, agent) - set_on_span( - SPANDATA.GEN_AI_TOOL_DEFINITIONS, - json.dumps(_transform_tool_definitions(agent.tools)), - ) + if has_data_collection_enabled(client_options): + if client_options["data_collection"]["gen_ai"]["inputs"]: + set_on_span( + SPANDATA.GEN_AI_TOOL_DEFINITIONS, + json.dumps(_transform_tool_definitions(agent.tools)), + ) + else: + # This is set unconditionally prior to data collection being introduced. + # Remove this block once data collection is fully rolled out + set_on_span( + SPANDATA.GEN_AI_TOOL_DEFINITIONS, + json.dumps(_transform_tool_definitions(agent.tools)), + ) _set_input_data(span, get_response_kwargs) diff --git a/sentry_sdk/integrations/openai_agents/utils.py b/sentry_sdk/integrations/openai_agents/utils.py index 08cc076510..f8f605fe0a 100644 --- a/sentry_sdk/integrations/openai_agents/utils.py +++ b/sentry_sdk/integrations/openai_agents/utils.py @@ -19,7 +19,11 @@ from sentry_sdk.scope import should_send_default_pii from sentry_sdk.traces import StreamedSpan from sentry_sdk.tracing_utils import should_truncate_gen_ai_input -from sentry_sdk.utils import event_from_exception, safe_serialize +from sentry_sdk.utils import ( + event_from_exception, + has_data_collection_enabled, + safe_serialize, +) if TYPE_CHECKING: from typing import Any, Union @@ -114,7 +118,11 @@ def _set_input_data( span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", get_response_kwargs: "dict[str, Any]", ) -> None: - if not should_send_default_pii(): + client = sentry_sdk.get_client() + if has_data_collection_enabled(client.options): + if not client.options["data_collection"]["gen_ai"]["inputs"]: + return + elif not should_send_default_pii(): return request_messages = [] @@ -182,7 +190,6 @@ def _set_input_data( ) normalized_messages = normalize_message_roles(request_messages) - client = sentry_sdk.get_client() scope = sentry_sdk.get_current_scope() messages_data = ( truncate_and_annotate_messages(normalized_messages, span, scope) @@ -201,7 +208,11 @@ def _set_input_data( def _set_output_data( span: "Union[sentry_sdk.tracing.Span, StreamedSpan]", result: "Any" ) -> None: - if not should_send_default_pii(): + client = sentry_sdk.get_client() + if has_data_collection_enabled(client.options): + if not client.options["data_collection"]["gen_ai"]["outputs"]: + return + elif not should_send_default_pii(): return output_messages: "dict[str, list[Any]]" = { diff --git a/tests/integrations/openai_agents/test_openai_agents.py b/tests/integrations/openai_agents/test_openai_agents.py index 704b7a0c49..cf055b762a 100644 --- a/tests/integrations/openai_agents/test_openai_agents.py +++ b/tests/integrations/openai_agents/test_openai_agents.py @@ -1038,6 +1038,307 @@ async def test_invoke_agent_span_data_collection_outputs( assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span_data +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.parametrize( + "data_collection,send_default_pii,expect_input", + [ + pytest.param( + {"gen_ai": {"inputs": True}}, + False, + True, + id="gen-ai-inputs-enabled-overrides-pii-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False}}, + True, + False, + id="gen-ai-inputs-disabled-overrides-pii-enabled", + ), + pytest.param( + {}, + False, + True, + id="gen-ai-omitted-defaults-to-enabled", + ), + pytest.param( + {"gen_ai": {"inputs": False}}, + False, + False, + id="gen-ai-inputs-disabled-and-pii-disabled", + ), + pytest.param( + {"gen_ai": {"outputs": False}}, + False, + True, + id="gen-ai-outputs-disabled-does-not-affect-inputs", + ), + pytest.param( + None, + False, + False, + id="no-data-collection-falls-back-to-send-default-pii", + ), + pytest.param( + None, + True, + True, + id="no-data-collection-pii-enabled-collects", + ), + ], +) +@pytest.mark.asyncio +async def test_data_collection_inputs( + sentry_init, + capture_events, + capture_items, + test_agent, + simple_test_tool, + nonstreaming_responses_model_response, + get_model_response, + data_collection, + send_default_pii, + expect_input, + stream_gen_ai_spans, + span_streaming, +): + client = AsyncOpenAI(api_key="test-key") + model = OpenAIResponsesModel(model="gpt-4", openai_client=client) + agent = test_agent.clone(model=model, tools=[simple_test_tool]) + + response = get_model_response( + nonstreaming_responses_model_response, serialize_pydantic=True + ) + + init_kwargs = { + "integrations": [OpenAIAgentsIntegration()], + "traces_sample_rate": 1.0, + "send_default_pii": send_default_pii, + "stream_gen_ai_spans": stream_gen_ai_spans, + } + if span_streaming: + init_kwargs["disabled_integrations"] = [StdlibIntegration] + init_kwargs["trace_lifecycle"] = "stream" + if data_collection is not None: + init_kwargs["_experiments"] = {"data_collection": data_collection} + + with patch.object( + agent.model._client._client, + "send", + return_value=response, + ) as _: + sentry_init(**init_kwargs) + + if span_streaming: + items = capture_items("span") + elif stream_gen_ai_spans: + items = capture_items("span", "transaction") + else: + events = capture_events() + + result = await agents.Runner.run( + agent, "Test input", run_config=test_run_config + ) + + assert result is not None + + if span_streaming or stream_gen_ai_spans: + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + ai_client_span = next( + span for span in spans if span["attributes"]["sentry.op"] == OP.GEN_AI_CHAT + ) + span_data = ai_client_span["attributes"] + else: + (transaction,) = events + ai_client_span = next( + span for span in transaction["spans"] if span["op"] == OP.GEN_AI_CHAT + ) + span_data = ai_client_span["data"] + + if expect_input: + assert "Test input" in span_data[SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert ( + "helpful test assistant" in span_data[SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS] + ) + else: + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span_data + assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in span_data + + # Tool definitions are only gated once data_collection is configured; without it + # they are set unconditionally, as they were before data_collection existed. + if expect_input or data_collection is None: + assert "simple_test_tool" in span_data[SPANDATA.GEN_AI_TOOL_DEFINITIONS] + else: + assert SPANDATA.GEN_AI_TOOL_DEFINITIONS not in span_data + + # Non-PII data is unaffected by the gate + assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "gpt-4" + assert span_data[SPANDATA.GEN_AI_USAGE_INPUT_TOKENS] == 10 + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.parametrize( + "data_collection,send_default_pii,expect_output", + [ + pytest.param( + {"gen_ai": {"outputs": True}}, + False, + True, + id="gen-ai-outputs-enabled-overrides-pii-disabled", + ), + pytest.param( + {"gen_ai": {"outputs": False}}, + True, + False, + id="gen-ai-outputs-disabled-overrides-pii-enabled", + ), + pytest.param( + {}, + False, + True, + id="gen-ai-omitted-defaults-to-enabled", + ), + pytest.param( + {"gen_ai": {"outputs": False}}, + False, + False, + id="gen-ai-outputs-disabled-and-pii-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False}}, + False, + True, + id="gen-ai-inputs-disabled-does-not-affect-outputs", + ), + pytest.param( + None, + False, + False, + id="no-data-collection-falls-back-to-send-default-pii", + ), + pytest.param( + None, + True, + True, + id="no-data-collection-pii-enabled-collects", + ), + ], +) +@pytest.mark.asyncio +async def test_data_collection_outputs( + sentry_init, + capture_events, + capture_items, + test_agent, + simple_test_tool, + get_model_response, + nonstreaming_responses_tool_call_model_responses, + data_collection, + send_default_pii, + expect_output, + stream_gen_ai_spans, + span_streaming, +): + client = AsyncOpenAI(api_key="test-key") + model = OpenAIResponsesModel(model="gpt-4", openai_client=client) + agent_with_tool = test_agent.clone(tools=[simple_test_tool], model=model) + + usage = ResponseUsage( + input_tokens=10, + input_tokens_details=InputTokensDetails( + cached_tokens=0, + cache_write_tokens=0, + ), + output_tokens=5, + output_tokens_details=OutputTokensDetails( + reasoning_tokens=0, + ), + total_tokens=15, + ) + responses = nonstreaming_responses_tool_call_model_responses( + tool_name="simple_test_tool", + arguments='{"message": "hello"}', + response_model="gpt-4", + response_text="Task completed using the tool", + response_ids=iter(["resp_tool_123", "resp_final_123"]), + usages=iter([usage, usage]), + ) + tool_response = get_model_response(next(responses), serialize_pydantic=True) + final_response = get_model_response(next(responses), serialize_pydantic=True) + + init_kwargs = { + "integrations": [OpenAIAgentsIntegration()], + "traces_sample_rate": 1.0, + "send_default_pii": send_default_pii, + "stream_gen_ai_spans": stream_gen_ai_spans, + } + if span_streaming: + init_kwargs["disabled_integrations"] = [StdlibIntegration] + init_kwargs["trace_lifecycle"] = "stream" + if data_collection is not None: + init_kwargs["_experiments"] = {"data_collection": data_collection} + + with patch.object( + agent_with_tool.model._client._client, + "send", + side_effect=[tool_response, final_response], + ) as _: + sentry_init(**init_kwargs) + + if span_streaming: + items = capture_items("span") + elif stream_gen_ai_spans: + items = capture_items("span", "transaction") + else: + events = capture_events() + + await agents.Runner.run( + agent_with_tool, + "Please use the simple test tool", + run_config=test_run_config, + ) + + if span_streaming or stream_gen_ai_spans: + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + chat_span_data = [ + span["attributes"] + for span in spans + if span["attributes"].get("sentry.op") == OP.GEN_AI_CHAT + ] + else: + (transaction,) = events + chat_span_data = [ + span["data"] + for span in transaction["spans"] + if span["op"] == OP.GEN_AI_CHAT + ] + + assert len(chat_span_data) == 2 + + if expect_output: + assert any( + "simple_test_tool" in data.get(SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, "") + for data in chat_span_data + ) + assert any( + "Task completed using the tool" + in str(data.get(SPANDATA.GEN_AI_RESPONSE_TEXT, "")) + for data in chat_span_data + ) + else: + for data in chat_span_data: + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in data + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in data + + # Non-PII data is unaffected by the gate + for data in chat_span_data: + assert data[SPANDATA.GEN_AI_REQUEST_MODEL] == "gpt-4" + assert data[SPANDATA.GEN_AI_USAGE_OUTPUT_TOKENS] == 5 + + @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio @@ -2510,6 +2811,7 @@ async def test_tool_execution_span( capture_events, capture_items, test_agent, + simple_test_tool, get_model_response, nonstreaming_responses_tool_call_model_responses, stream_gen_ai_spans, @@ -2520,11 +2822,6 @@ async def test_tool_execution_span( Test tool execution span creation with `AgentRunner.run()`. """ - @agents.function_tool - def simple_test_tool(message: str) -> str: - """A simple tool""" - return f"Tool executed with: {message}" - # Create agent with the tool client = AsyncOpenAI(api_key="test-key") model = OpenAIResponsesModel(model="gpt-4", openai_client=client) @@ -3138,6 +3435,7 @@ async def test_run_streamed_tool_execution_span( capture_events, capture_items, test_agent, + simple_test_tool, get_model_response, async_iterator, server_side_event_chunks, @@ -3150,11 +3448,6 @@ async def test_run_streamed_tool_execution_span( Test tool execution span creation with `AgentRunner.run_streamed()`. """ - @agents.function_tool - def simple_test_tool(message: str) -> str: - """A simple tool""" - return f"Tool executed with: {message}" - # Create agent with the tool client = AsyncOpenAI(api_key="test-key") model = OpenAIResponsesModel(model="gpt-4", openai_client=client) @@ -3928,6 +4221,7 @@ async def test_model_behavior_error( capture_events, capture_items, test_agent, + simple_test_tool, stream_gen_ai_spans, span_streaming, ): @@ -3936,11 +4230,6 @@ async def test_model_behavior_error( The mocked API response indicates that "wrong_tool" was called. """ - @agents.function_tool - def simple_test_tool(message: str) -> str: - """A simple tool""" - return f"Tool executed with: {message}" - # Create agent with the tool agent_with_tool = test_agent.clone(tools=[simple_test_tool])