Skip to content

feat(openai-agents): Add data_collection controls for invoke_agent spans #7124

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
ericapisani wants to merge 1 commit into from
+358 −4
Draft

feat(openai-agents): Add data_collection controls for invoke_agent spans #7124

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 18 additions & 4 deletions sentry_sdk/integrations/openai_agents/spans/invoke_agent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
has_span_streaming_enabled,
should_truncate_gen_ai_input,
)
from sentry_sdk.utils import safe_serialize
from sentry_sdk.utils import has_data_collection_enabled, safe_serialize

from ..consts import SPAN_ORIGIN
from ..utils import _set_agent_data, _set_usage_data
Expand All @@ -28,7 +28,8 @@
def invoke_agent_span(
context: "agents.RunContextWrapper", agent: "agents.Agent", kwargs: "dict[str, Any]"
) -> "Union[sentry_sdk.tracing.Span, StreamedSpan]":
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"invoke_agent {agent.name}",
Expand All @@ -49,7 +50,14 @@ def invoke_agent_span(

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

if should_send_default_pii():
record_inputs = False
if has_data_collection_enabled(client_options):
if client_options["data_collection"]["gen_ai"]["inputs"]:
record_inputs = True
elif should_send_default_pii():
record_inputs = True

if record_inputs:
messages = []
if agent.instructions:
message = (
Expand Down Expand Up @@ -110,7 +118,13 @@ def update_invoke_agent_span(
if hasattr(context, "usage"):
_set_usage_data(span, context.usage)

if should_send_default_pii():
client = sentry_sdk.get_client()
if has_data_collection_enabled(client.options):
if client.options["data_collection"]["gen_ai"]["outputs"]:
set_data_normalized(
span, SPANDATA.GEN_AI_RESPONSE_TEXT, output, unpack=False
)
elif should_send_default_pii():
set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, output, unpack=False)

# Add conversation ID from agent
Expand Down
340 changes: 340 additions & 0 deletions tests/integrations/openai_agents/test_openai_agents.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -698,6 +698,346 @@ async def test_agent_invocation_span_no_pii(
assert ai_client_span["data"]["gen_ai.request.top_p"] == 1.0


@pytest.mark.parametrize("span_streaming", [True, False])
@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
@pytest.mark.parametrize(
"init_kwargs,expect_messages",
[
pytest.param(
{"_experiments": {"data_collection": {"gen_ai": {"inputs": True}}}},
True,
id="gen_ai_inputs_true",
),
pytest.param(
{"_experiments": {"data_collection": {"gen_ai": {"inputs": False}}}},
False,
id="gen_ai_inputs_false",
),
pytest.param(
{"_experiments": {"data_collection": {}}},
True,
id="data_collection_defaults_to_enabled",
),
pytest.param(
{
"send_default_pii": True,
"_experiments": {"data_collection": {"gen_ai": {"inputs": False}}},
},
False,
id="data_collection_wins_over_send_default_pii_true",
),
pytest.param(
{
"send_default_pii": False,
"_experiments": {"data_collection": {"gen_ai": {"inputs": True}}},
},
True,
id="data_collection_wins_over_send_default_pii_false",
),
pytest.param(
{"send_default_pii": True},
True,
id="legacy_send_default_pii_true",
),
pytest.param(
{"send_default_pii": False},
False,
id="legacy_send_default_pii_false",
),
],
)
@pytest.mark.asyncio
async def test_invoke_agent_span_data_collection_inputs(
sentry_init,
capture_events,
capture_items,
test_agent,
nonstreaming_responses_model_response,
get_model_response,
init_kwargs,
expect_messages,
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)

response = get_model_response(
nonstreaming_responses_model_response, serialize_pydantic=True
)

system_message = {
"role": "system",
"content": [{"text": "You are a helpful test assistant.", "type": "text"}],
}
user_message = {
"role": "user",
"content": [{"text": "Test input", "type": "text"}],
}
expected_messages = (
[user_message]
if not stream_gen_ai_spans and not span_streaming
else [system_message, user_message]
)

if span_streaming:
with patch.object(
agent.model._client._client,
"send",
return_value=response,
) as _:
sentry_init(
integrations=[OpenAIAgentsIntegration()],
disabled_integrations=[StdlibIntegration],
traces_sample_rate=1.0,
stream_gen_ai_spans=stream_gen_ai_spans,
trace_lifecycle="stream",
**init_kwargs,
)

items = capture_items("span")

result = await agents.Runner.run(
agent, "Test input", run_config=test_run_config
)

assert result is not None

sentry_sdk.flush()
spans = [item.payload for item in items]
invoke_agent_span = next(
span
for span in spans
if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT
)
span_data = invoke_agent_span["attributes"]
elif stream_gen_ai_spans:
with patch.object(
agent.model._client._client,
"send",
return_value=response,
) as _:
sentry_init(
integrations=[OpenAIAgentsIntegration()],
traces_sample_rate=1.0,
stream_gen_ai_spans=stream_gen_ai_spans,
**init_kwargs,
)

items = capture_items("span", "transaction")

result = await agents.Runner.run(
agent, "Test input", run_config=test_run_config
)

assert result is not None

spans = [item.payload for item in items if item.type == "span"]
invoke_agent_span = next(
span
for span in spans
if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT
)
span_data = invoke_agent_span["attributes"]
else:
with patch.object(
agent.model._client._client,
"send",
return_value=response,
) as _:
sentry_init(
integrations=[OpenAIAgentsIntegration()],
traces_sample_rate=1.0,
stream_gen_ai_spans=stream_gen_ai_spans,
**init_kwargs,
)
events = capture_events()

result = await agents.Runner.run(
agent, "Test input", run_config=test_run_config
)

assert result is not None

(transaction,) = events
invoke_agent_span = next(
span
for span in transaction["spans"]
if span["op"] == OP.GEN_AI_INVOKE_AGENT
)
span_data = invoke_agent_span["data"]

if expect_messages:
assert (
json.loads(span_data[SPANDATA.GEN_AI_REQUEST_MESSAGES]) == expected_messages
)
else:
assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span_data


@pytest.mark.parametrize("span_streaming", [True, False])
@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
@pytest.mark.parametrize(
"init_kwargs,expect_response_text",
[
pytest.param(
{"_experiments": {"data_collection": {"gen_ai": {"outputs": True}}}},
True,
id="gen_ai_outputs_true",
),
pytest.param(
{"_experiments": {"data_collection": {"gen_ai": {"outputs": False}}}},
False,
id="gen_ai_outputs_false",
),
pytest.param(
{"_experiments": {"data_collection": {}}},
True,
id="data_collection_defaults_to_enabled",
),
pytest.param(
{
"send_default_pii": True,
"_experiments": {"data_collection": {"gen_ai": {"outputs": False}}},
},
False,
id="data_collection_wins_over_send_default_pii_true",
),
pytest.param(
{
"send_default_pii": False,
"_experiments": {"data_collection": {"gen_ai": {"outputs": True}}},
},
True,
id="data_collection_wins_over_send_default_pii_false",
),
pytest.param(
{"send_default_pii": True},
True,
id="legacy_send_default_pii_true",
),
pytest.param(
{"send_default_pii": False},
False,
id="legacy_send_default_pii_false",
),
],
)
@pytest.mark.asyncio
async def test_invoke_agent_span_data_collection_outputs(
sentry_init,
capture_events,
capture_items,
test_agent,
nonstreaming_responses_model_response,
get_model_response,
init_kwargs,
expect_response_text,
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)

response = get_model_response(
nonstreaming_responses_model_response, serialize_pydantic=True
)

if span_streaming:
with patch.object(
agent.model._client._client,
"send",
return_value=response,
) as _:
sentry_init(
integrations=[OpenAIAgentsIntegration()],
disabled_integrations=[StdlibIntegration],
traces_sample_rate=1.0,
stream_gen_ai_spans=stream_gen_ai_spans,
trace_lifecycle="stream",
**init_kwargs,
)

items = capture_items("span")

result = await agents.Runner.run(
agent, "Test input", run_config=test_run_config
)

assert result is not None

sentry_sdk.flush()
spans = [item.payload for item in items]
invoke_agent_span = next(
span
for span in spans
if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT
)
span_data = invoke_agent_span["attributes"]
elif stream_gen_ai_spans:
with patch.object(
agent.model._client._client,
"send",
return_value=response,
) as _:
sentry_init(
integrations=[OpenAIAgentsIntegration()],
traces_sample_rate=1.0,
stream_gen_ai_spans=stream_gen_ai_spans,
**init_kwargs,
)

items = capture_items("span", "transaction")

result = await agents.Runner.run(
agent, "Test input", run_config=test_run_config
)

assert result is not None

spans = [item.payload for item in items if item.type == "span"]
invoke_agent_span = next(
span
for span in spans
if span["attributes"]["sentry.op"] == OP.GEN_AI_INVOKE_AGENT
)
span_data = invoke_agent_span["attributes"]
else:
with patch.object(
agent.model._client._client,
"send",
return_value=response,
) as _:
sentry_init(
integrations=[OpenAIAgentsIntegration()],
traces_sample_rate=1.0,
stream_gen_ai_spans=stream_gen_ai_spans,
**init_kwargs,
)
events = capture_events()

result = await agents.Runner.run(
agent, "Test input", run_config=test_run_config
)

assert result is not None

(transaction,) = events
invoke_agent_span = next(
span
for span in transaction["spans"]
if span["op"] == OP.GEN_AI_INVOKE_AGENT
)
span_data = invoke_agent_span["data"]

if expect_response_text:
assert span_data[SPANDATA.GEN_AI_RESPONSE_TEXT] == "Hello, how can I help you?"
else:
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.asyncio
Expand Down