Skip to content

Commit af05b2c

Browse files
committed
feat(openai): Gate Chat Completions inputs behind data collection
Apply data_collection.gen_ai.inputs to Chat Completions prompt messages, system instructions, and tool definitions, matching the Responses API. Preserve legacy send_default_pii behavior when data collection is not configured. Refs PY-2588
1 parent ac8eae7 commit af05b2c

2 files changed

Lines changed: 188 additions & 21 deletions

File tree

sentry_sdk/integrations/openai.py

Lines changed: 40 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -476,19 +476,12 @@ def _set_completions_api_input_data(
476476
kwargs: "dict[str, Any]",
477477
integration: "OpenAIIntegration",
478478
) -> None:
479-
messages: "Optional[Union[str, Iterable[ChatCompletionMessageParam]]]" = kwargs.get(
480-
"messages"
481-
)
482-
483-
tools = kwargs.get("tools")
484-
if tools is not None and _is_given(tools) and len(tools) > 0:
485-
set_data_normalized(
486-
span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools)
487-
)
488-
489479
set_on_span = (
490480
span.set_attribute if isinstance(span, StreamedSpan) else span.set_data
491481
)
482+
483+
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
484+
492485
model = kwargs.get("model")
493486
if model is not None:
494487
set_on_span(SPANDATA.GEN_AI_REQUEST_MODEL, model)
@@ -517,17 +510,47 @@ def _set_completions_api_input_data(
517510
if reasoning_level is not None and _is_given(reasoning_level):
518511
set_on_span(SPANDATA.GEN_AI_REQUEST_REASONING_LEVEL, reasoning_level)
519512

520-
if (
521-
not should_send_default_pii()
522-
or not integration.include_prompts
523-
or messages is None
524-
):
525-
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
513+
client = sentry_sdk.get_client()
514+
if has_data_collection_enabled(client.options):
515+
if (
516+
integration.include_prompts
517+
and client.options["data_collection"]["gen_ai"]["inputs"]
518+
):
519+
tools = kwargs.get("tools")
520+
if tools is not None and _is_given(tools) and len(tools) > 0:
521+
set_data_normalized(
522+
span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools)
523+
)
524+
else:
525+
# Pre-data collection this was always set, so this needs to be left here for now until
526+
# we deprecate `send_default_pii`. Once we do, this 'else' branch should be removed,
527+
# and the above branch placed below the "if not should_send_default_pii() or not integration.include_prompts"
528+
# line below
529+
tools = kwargs.get("tools")
530+
if tools is not None and _is_given(tools) and len(tools) > 0:
531+
set_data_normalized(
532+
span, SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS, safe_serialize(tools)
533+
)
534+
535+
536+
messages: "Optional[Union[str, Iterable[ChatCompletionMessageParam]]]" = kwargs.get(
537+
"messages"
538+
)
539+
540+
if has_data_collection_enabled(client.options):
541+
# This takes precedence over the global data collection settings
542+
if not integration.include_prompts:
543+
return
544+
if not client.options["data_collection"]["gen_ai"]["inputs"]:
545+
return
546+
elif not should_send_default_pii() or not integration.include_prompts:
547+
return
548+
549+
if messages is None:
526550
return
527551

528552
if isinstance(messages, str):
529553
normalized_messages = normalize_message_roles([messages]) # type: ignore
530-
client = sentry_sdk.get_client()
531554
scope = sentry_sdk.get_current_scope()
532555
messages_data = (
533556
truncate_and_annotate_messages(normalized_messages, span, scope)
@@ -538,12 +561,10 @@ def _set_completions_api_input_data(
538561
set_data_normalized(
539562
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
540563
)
541-
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
542564
return
543565

544566
# dict special case following https://github.com/openai/openai-python/blob/3e0c05b84a2056870abf3bd6a5e7849020209cc3/src/openai/_utils/_transform.py#L194-L197
545567
if not isinstance(messages, Iterable) or isinstance(messages, dict):
546-
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
547568
return
548569

549570
messages = list(messages)
@@ -575,8 +596,6 @@ def _set_completions_api_input_data(
575596
span, SPANDATA.GEN_AI_REQUEST_MESSAGES, messages_data, unpack=False
576597
)
577598

578-
set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "chat")
579-
580599

581600
def _set_embeddings_input_data(
582601
span: "Union[Span, StreamedSpan]",

tests/integrations/openai/test_openai.py

Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -464,6 +464,154 @@ def test_nonstreaming_chat_completion(
464464
assert span["data"]["gen_ai.usage.total_tokens"] == 30
465465

466466

467+
@pytest.mark.skipif(
468+
OPENAI_VERSION <= (1, 1, 0),
469+
reason="OpenAI versions <=1.1.0 do not support the tools parameter.",
470+
)
471+
@pytest.mark.parametrize("span_streaming", [True, False])
472+
@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
473+
@pytest.mark.parametrize(
474+
"data_collection,include_prompts,expected_present,expected_absent",
475+
[
476+
pytest.param(
477+
{"gen_ai": {"inputs": True}},
478+
True,
479+
{
480+
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: json.dumps(
481+
[{"type": "text", "content": "You are a helpful assistant."}]
482+
),
483+
SPANDATA.GEN_AI_REQUEST_MESSAGES: safe_serialize(
484+
[{"role": "user", "content": "hello"}]
485+
),
486+
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS: safe_serialize(EXAMPLE_TOOLS),
487+
},
488+
[],
489+
id="inputs-enabled",
490+
),
491+
pytest.param(
492+
{"gen_ai": {"inputs": False}},
493+
True,
494+
{},
495+
[
496+
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
497+
SPANDATA.GEN_AI_REQUEST_MESSAGES,
498+
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
499+
],
500+
id="inputs-disabled",
501+
),
502+
pytest.param(
503+
{},
504+
True,
505+
{
506+
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: json.dumps(
507+
[{"type": "text", "content": "You are a helpful assistant."}]
508+
),
509+
SPANDATA.GEN_AI_REQUEST_MESSAGES: safe_serialize(
510+
[{"role": "user", "content": "hello"}]
511+
),
512+
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS: safe_serialize(EXAMPLE_TOOLS),
513+
},
514+
[],
515+
id="gen-ai-omitted-defaults-to-enabled",
516+
),
517+
pytest.param(
518+
{"gen_ai": {"inputs": True}},
519+
False,
520+
{},
521+
[
522+
SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS,
523+
SPANDATA.GEN_AI_REQUEST_MESSAGES,
524+
SPANDATA.GEN_AI_REQUEST_AVAILABLE_TOOLS,
525+
],
526+
id="include-prompts-disabled-overrides-inputs-enabled",
527+
),
528+
],
529+
)
530+
def test_completions_api_data_collection(
531+
sentry_init,
532+
capture_events,
533+
capture_items,
534+
data_collection,
535+
include_prompts,
536+
expected_present,
537+
expected_absent,
538+
nonstreaming_chat_completions_model_response,
539+
stream_gen_ai_spans,
540+
span_streaming,
541+
):
542+
sentry_init(
543+
integrations=[OpenAIIntegration(include_prompts=include_prompts)],
544+
disabled_integrations=[StdlibIntegration],
545+
traces_sample_rate=1.0,
546+
_experiments={"data_collection": data_collection},
547+
stream_gen_ai_spans=stream_gen_ai_spans,
548+
trace_lifecycle="stream" if span_streaming else "static",
549+
)
550+
551+
client = OpenAI(api_key="z")
552+
client.chat.completions._post = mock.Mock(
553+
return_value=nonstreaming_chat_completions_model_response(
554+
response_id="chat-id",
555+
response_model="gpt-3.5-turbo",
556+
message_content="the model response",
557+
created=10000000,
558+
usage=CompletionUsage(
559+
prompt_tokens=20,
560+
completion_tokens=10,
561+
total_tokens=30,
562+
),
563+
)
564+
)
565+
566+
create_kwargs = {
567+
"model": "some-model",
568+
"messages": [
569+
{"role": "system", "content": "You are a helpful assistant."},
570+
{"role": "user", "content": "hello"},
571+
],
572+
"max_tokens": 100,
573+
"presence_penalty": 0.1,
574+
"frequency_penalty": 0.2,
575+
"temperature": 0.7,
576+
"top_p": 0.9,
577+
"tools": EXAMPLE_TOOLS,
578+
}
579+
580+
if span_streaming or stream_gen_ai_spans:
581+
items = capture_items("span")
582+
583+
with start_transaction(name="openai tx"):
584+
client.chat.completions.create(**create_kwargs)
585+
586+
sentry_sdk.flush()
587+
(span,) = (item.payload for item in items)
588+
span_data = span["attributes"]
589+
else:
590+
events = capture_events()
591+
592+
with start_transaction(name="openai tx"):
593+
client.chat.completions.create(**create_kwargs)
594+
595+
(transaction,) = events
596+
(span,) = transaction["spans"]
597+
span_data = span["data"]
598+
599+
assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "chat"
600+
assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model"
601+
assert span_data[SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100
602+
assert span_data[SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1
603+
assert span_data[SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2
604+
assert span_data[SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7
605+
assert span_data[SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9
606+
assert span_data[SPANDATA.GEN_AI_SYSTEM] == "openai"
607+
608+
for key, value in expected_present.items():
609+
assert span_data[key] == value
610+
611+
for key in expected_absent:
612+
assert key not in span_data
613+
614+
467615
@pytest.mark.parametrize("span_streaming", [True, False])
468616
@pytest.mark.parametrize("stream_gen_ai_spans", [True, False])
469617
@pytest.mark.asyncio

0 commit comments

Comments
 (0)