Skip to content

Commit f261e87

Browse files
ericapisaniericapisani
committed
feat(openai-agents): Add data_collection controls for invoke_agent spans
Allow SDK users to control whether agent instructions/inputs and outputs are recorded on invoke_agent spans via `_experiments.data_collection.gen_ai.inputs` and `.outputs`, with `send_default_pii` as a fallback for backwards compatibility. When data_collection config is present, it takes precedence over send_default_pii, giving users granular control over AI telemetry collection. Refs PY-2588
1 parent fe96523 commit f261e87

2 files changed

Lines changed: 358 additions & 4 deletions

File tree

‎sentry_sdk/integrations/openai_agents/spans/invoke_agent.py‎

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
has_span_streaming_enabled,
1515
should_truncate_gen_ai_input,
1616
)
17-
from sentry_sdk.utils import safe_serialize
17+
from sentry_sdk.utils import has_data_collection_enabled, safe_serialize
1818

1919
from ..consts import SPAN_ORIGIN
2020
from ..utils import _set_agent_data, _set_usage_data
@@ -28,7 +28,8 @@
2828
def invoke_agent_span(
2929
context: "agents.RunContextWrapper", agent: "agents.Agent", kwargs: "dict[str, Any]"
3030
) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
31-
span_streaming = has_span_streaming_enabled(sentry_sdk.get_client().options)
31+
client_options = sentry_sdk.get_client().options
32+
span_streaming = has_span_streaming_enabled(client_options)
3233
if span_streaming:
3334
span = sentry_sdk.traces.start_span(
3435
name=f"invoke_agent {agent.name}",
@@ -49,7 +50,14 @@ def invoke_agent_span(
4950

5051
span.set_data(SPANDATA.GEN_AI_OPERATION_NAME, "invoke_agent")
5152

52-
if should_send_default_pii():
53+
record_inputs = False
54+
if has_data_collection_enabled(client_options):
55+
if client_options["data_collection"]["gen_ai"]["inputs"]:
56+
record_inputs = True
57+
elif should_send_default_pii():
58+
record_inputs = True
59+
60+
if record_inputs:
5361
messages = []
5462
if agent.instructions:
5563
message = (
@@ -110,7 +118,13 @@ def update_invoke_agent_span(
110118
if hasattr(context, "usage"):
111119
_set_usage_data(span, context.usage)
112120

113-
if should_send_default_pii():
121+
client = sentry_sdk.get_client()
122+
if has_data_collection_enabled(client.options):
123+
if client.options["data_collection"]["gen_ai"]["outputs"]:
124+
set_data_normalized(
125+
span, SPANDATA.GEN_AI_RESPONSE_TEXT, output, unpack=False
126+
)
127+
elif should_send_default_pii():
114128
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output, unpack=False)
115129

116130
# Add conversation ID from agent

‎tests/integrations/openai_agents/test_openai_agents.py‎

Lines changed: 340 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -698,6 +698,346 @@ async def test_agent_invocation_span_no_pii(
698698
assert ai_client_span["data"]["gen_ai.request.top_p"] == 1.0
699699

700700

701+
@pytest.mark.parametrize("span_streaming", [True, False])
702+
@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
703+
@pytest.mark.parametrize(
704+
"init_kwargs,expect_messages",
705+
[
706+
pytest.param(
707+
{"_experiments": {"data_collection": {"gen_ai": {"inputs": True}}}},
708+
True,
709+
id="gen_ai_inputs_true",
710+
),
711+
pytest.param(
712+
{"_experiments": {"data_collection": {"gen_ai": {"inputs": False}}}},
713+
False,
714+
id="gen_ai_inputs_false",
715+
),
716+
pytest.param(
717+
{"_experiments": {"data_collection": {}}},
718+
True,
719+
id="data_collection_defaults_to_enabled",
720+
),
721+
pytest.param(
722+
{
723+
"send_default_pii": True,
724+
"_experiments": {"data_collection": {"gen_ai": {"inputs": False}}},
725+
},
726+
False,
727+
id="data_collection_wins_over_send_default_pii_true",
728+
),
729+
pytest.param(
730+
{
731+
"send_default_pii": False,
732+
"_experiments": {"data_collection": {"gen_ai": {"inputs": True}}},
733+
},
734+
True,
735+
id="data_collection_wins_over_send_default_pii_false",
736+
),
737+
pytest.param(
738+
{"send_default_pii": True},
739+
True,
740+
id="legacy_send_default_pii_true",
741+
),
742+
pytest.param(
743+
{"send_default_pii": False},
744+
False,
745+
id="legacy_send_default_pii_false",
746+
),
747+
],
748+
)
749+
@pytest.mark.asyncio
750+
async def test_invoke_agent_span_data_collection_inputs(
751+
sentry_init,
752+
capture_events,
753+
capture_items,
754+
test_agent,
755+
nonstreaming_responses_model_response,
756+
get_model_response,
757+
init_kwargs,
758+
expect_messages,
759+
stream_gen_ai_spans,
760+
span_streaming,
761+
):
762+
client = AsyncOpenAI(api_key="test-key")
763+
model = OpenAIResponsesModel(model="gpt-4", openai_client=client)
764+
agent = test_agent.clone(model=model)
765+
766+
response = get_model_response(
767+
nonstreaming_responses_model_response, serialize_pydantic=True
768+
)
769+
770+
system_message = {
771+
"role": "system",
772+
"content": [{"text": "You are a helpful test assistant.", "type": "text"}],
773+
}
774+
user_message = {
775+
"role": "user",
776+
"content": [{"text": "Test input", "type": "text"}],
777+
}
778+
expected_messages = (
779+
[user_message]
780+
if not stream_gen_ai_spans and not span_streaming
781+
else [system_message, user_message]
782+
)
783+
784+
if span_streaming:
785+
with patch.object(
786+
agent.model._client._client,
787+
"send",
788+
return_value=response,
789+
) as _:
790+
sentry_init(
791+
integrations=[OpenAIAgentsIntegration()],
792+
disabled_integrations=[StdlibIntegration],
793+
traces_sample_rate=1.0,
794+
stream_gen_ai_spans=stream_gen_ai_spans,
795+
trace_lifecycle="stream",
796+
**init_kwargs,
797+
)
798+
799+
items = capture_items("span")
800+
801+
result = await agents.Runner.run(
802+
agent, "Test input", run_config=test_run_config
803+
)
804+
805+
assert result is not None
806+
807+
sentry_sdk.flush()
808+
spans = [item.payload for item in items]
809+
invoke_agent_span = next(
810+
span
811+
for span in spans
812+
if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT
813+
)
814+
span_data = invoke_agent_span["attributes"]
815+
elif stream_gen_ai_spans:
816+
with patch.object(
817+
agent.model._client._client,
818+
"send",
819+
return_value=response,
820+
) as _:
821+
sentry_init(
822+
integrations=[OpenAIAgentsIntegration()],
823+
traces_sample_rate=1.0,
824+
stream_gen_ai_spans=stream_gen_ai_spans,
825+
**init_kwargs,
826+
)
827+
828+
items = capture_items("span", "transaction")
829+
830+
result = await agents.Runner.run(
831+
agent, "Test input", run_config=test_run_config
832+
)
833+
834+
assert result is not None
835+
836+
spans = [item.payload for item in items if item.type == "span"]
837+
invoke_agent_span = next(
838+
span
839+
for span in spans
840+
if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT
841+
)
842+
span_data = invoke_agent_span["attributes"]
843+
else:
844+
with patch.object(
845+
agent.model._client._client,
846+
"send",
847+
return_value=response,
848+
) as _:
849+
sentry_init(
850+
integrations=[OpenAIAgentsIntegration()],
851+
traces_sample_rate=1.0,
852+
stream_gen_ai_spans=stream_gen_ai_spans,
853+
**init_kwargs,
854+
)
855+
events = capture_events()
856+
857+
result = await agents.Runner.run(
858+
agent, "Test input", run_config=test_run_config
859+
)
860+
861+
assert result is not None
862+
863+
(transaction,) = events
864+
invoke_agent_span = next(
865+
span
866+
for span in transaction["spans"]
867+
if span["op"] == OP.GEN_AI_INVOKE_AGENT
868+
)
869+
span_data = invoke_agent_span["data"]
870+
871+
if expect_messages:
872+
assert (
873+
json.loads(span_data[SPANDATA.GEN_AI_REQUEST_MESSAGES]) == expected_messages
874+
)
875+
else:
876+
assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span_data
877+
878+
879+
@pytest.mark.parametrize("span_streaming", [True, False])
880+
@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
881+
@pytest.mark.parametrize(
882+
"init_kwargs,expect_response_text",
883+
[
884+
pytest.param(
885+
{"_experiments": {"data_collection": {"gen_ai": {"outputs": True}}}},
886+
True,
887+
id="gen_ai_outputs_true",
888+
),
889+
pytest.param(
890+
{"_experiments": {"data_collection": {"gen_ai": {"outputs": False}}}},
891+
False,
892+
id="gen_ai_outputs_false",
893+
),
894+
pytest.param(
895+
{"_experiments": {"data_collection": {}}},
896+
True,
897+
id="data_collection_defaults_to_enabled",
898+
),
899+
pytest.param(
900+
{
901+
"send_default_pii": True,
902+
"_experiments": {"data_collection": {"gen_ai": {"outputs": False}}},
903+
},
904+
False,
905+
id="data_collection_wins_over_send_default_pii_true",
906+
),
907+
pytest.param(
908+
{
909+
"send_default_pii": False,
910+
"_experiments": {"data_collection": {"gen_ai": {"outputs": True}}},
911+
},
912+
True,
913+
id="data_collection_wins_over_send_default_pii_false",
914+
),
915+
pytest.param(
916+
{"send_default_pii": True},
917+
True,
918+
id="legacy_send_default_pii_true",
919+
),
920+
pytest.param(
921+
{"send_default_pii": False},
922+
False,
923+
id="legacy_send_default_pii_false",
924+
),
925+
],
926+
)
927+
@pytest.mark.asyncio
928+
async def test_invoke_agent_span_data_collection_outputs(
929+
sentry_init,
930+
capture_events,
931+
capture_items,
932+
test_agent,
933+
nonstreaming_responses_model_response,
934+
get_model_response,
935+
init_kwargs,
936+
expect_response_text,
937+
stream_gen_ai_spans,
938+
span_streaming,
939+
):
940+
client = AsyncOpenAI(api_key="test-key")
941+
model = OpenAIResponsesModel(model="gpt-4", openai_client=client)
942+
agent = test_agent.clone(model=model)
943+
944+
response = get_model_response(
945+
nonstreaming_responses_model_response, serialize_pydantic=True
946+
)
947+
948+
if span_streaming:
949+
with patch.object(
950+
agent.model._client._client,
951+
"send",
952+
return_value=response,
953+
) as _:
954+
sentry_init(
955+
integrations=[OpenAIAgentsIntegration()],
956+
disabled_integrations=[StdlibIntegration],
957+
traces_sample_rate=1.0,
958+
stream_gen_ai_spans=stream_gen_ai_spans,
959+
trace_lifecycle="stream",
960+
**init_kwargs,
961+
)
962+
963+
items = capture_items("span")
964+
965+
result = await agents.Runner.run(
966+
agent, "Test input", run_config=test_run_config
967+
)
968+
969+
assert result is not None
970+
971+
sentry_sdk.flush()
972+
spans = [item.payload for item in items]
973+
invoke_agent_span = next(
974+
span
975+
for span in spans
976+
if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT
977+
)
978+
span_data = invoke_agent_span["attributes"]
979+
elif stream_gen_ai_spans:
980+
with patch.object(
981+
agent.model._client._client,
982+
"send",
983+
return_value=response,
984+
) as _:
985+
sentry_init(
986+
integrations=[OpenAIAgentsIntegration()],
987+
traces_sample_rate=1.0,
988+
stream_gen_ai_spans=stream_gen_ai_spans,
989+
**init_kwargs,
990+
)
991+
992+
items = capture_items("span", "transaction")
993+
994+
result = await agents.Runner.run(
995+
agent, "Test input", run_config=test_run_config
996+
)
997+
998+
assert result is not None
999+
1000+
spans = [item.payload for item in items if item.type == "span"]
1001+
invoke_agent_span = next(
1002+
span
1003+
for span in spans
1004+
if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT
1005+
)
1006+
span_data = invoke_agent_span["attributes"]
1007+
else:
1008+
with patch.object(
1009+
agent.model._client._client,
1010+
"send",
1011+
return_value=response,
1012+
) as _:
1013+
sentry_init(
1014+
integrations=[OpenAIAgentsIntegration()],
1015+
traces_sample_rate=1.0,
1016+
stream_gen_ai_spans=stream_gen_ai_spans,
1017+
**init_kwargs,
1018+
)
1019+
events = capture_events()
1020+
1021+
result = await agents.Runner.run(
1022+
agent, "Test input", run_config=test_run_config
1023+
)
1024+
1025+
assert result is not None
1026+
1027+
(transaction,) = events
1028+
invoke_agent_span = next(
1029+
span
1030+
for span in transaction["spans"]
1031+
if span["op"] == OP.GEN_AI_INVOKE_AGENT
1032+
)
1033+
span_data = invoke_agent_span["data"]
1034+
1035+
if expect_response_text:
1036+
assert span_data[SPANDATA.GEN_AI_RESPONSE_TEXT] == "Hello, how can I help you?"
1037+
else:
1038+
assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span_data
1039+
1040+
7011041
@pytest.mark.parametrize("span_streaming", [True, False])
7021042
@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
7031043
@pytest.mark.asyncio

0 commit comments

Comments
 (0)