From 566ea5cf538e7db36507822b6b486aafc6c8d6c3 Mon Sep 17 00:00:00 2001 From: Erica Pisani Date: Tue, 4 Aug 2026 16:16:54 -0400 Subject: [PATCH 1/2] feat(openai): Gate response output collection on data collection options Gate OpenAI response text and tool calls collection on the new data_collection settings, giving precedence over the legacy send_default_pii flag for backwards compatibility. Applies to Chat Completions API, Responses API, and their streaming variants. Refs PY-2588 --- sentry_sdk/integrations/openai.py | 80 +- tests/integrations/openai/test_openai.py | 7188 ++++++++++++---------- 2 files changed, 4110 insertions(+), 3158 deletions(-) diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index 304a3a0899..8525edc930 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -681,12 +681,24 @@ def _set_common_output_data( integration: "OpenAIIntegration", finish_span: bool = True, ) -> None: + client = sentry_sdk.get_client() if hasattr(response, "model"): set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_MODEL, response.model) # Chat Completions API if hasattr(response, "choices") and response.choices is not None: - if should_send_default_pii() and integration.include_prompts: + if has_data_collection_enabled(client.options): + if client.options["data_collection"]["gen_ai"]["outputs"]: + response_text = [ + choice.message.model_dump() + for choice in response.choices + if choice.message is not None + ] + if len(response_text) > 0: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, response_text + ) + elif should_send_default_pii() and integration.include_prompts: response_text = [ choice.message.model_dump() for choice in response.choices @@ -709,12 +721,40 @@ def _set_common_output_data( # Responses API elif hasattr(response, "output"): - if should_send_default_pii() and integration.include_prompts: - output_messages: "dict[str, list[Any]]" = { - "response": [], - "tool": [], - } - + output_messages: "dict[str, list[Any]]" = { + "response": [], + "tool": [], + } + + if has_data_collection_enabled(client.options): + if client.options["data_collection"]["gen_ai"]["outputs"]: + for output in response.output: + if output.type == "function_call": + output_messages["tool"].append(output.dict()) + elif output.type == "message": + for output_message in output.content: + try: + output_messages["response"].append(output_message.text) + except AttributeError: + # Unknown output message type, just return the json + output_messages["response"].append( + output_message.dict() + ) + + if len(output_messages["tool"]) > 0: + set_data_normalized( + span, + SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS, + output_messages["tool"], + unpack=False, + ) + + if len(output_messages["response"]) > 0: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, output_messages["response"] + ) + + elif should_send_default_pii() and integration.include_prompts: for output in response.output: if output.type == "function_call": output_messages["tool"].append(output.dict()) @@ -961,6 +1001,7 @@ def _wrap_synchronous_completions_chunk_iterator( ttft = None data_buf: "list[list[str]]" = [] # one for each choice streaming_message_total_token_usage = None + client = sentry_sdk.get_client() for x in old_iterator: if isinstance(span, StreamedSpan): @@ -993,7 +1034,10 @@ def _wrap_synchronous_completions_chunk_iterator( all_responses = None if len(data_buf) > 0: all_responses = ["".join(chunk) for chunk in data_buf] - if should_send_default_pii() and integration.include_prompts: + if has_data_collection_enabled(client.options): + if client.options["data_collection"]["gen_ai"]["outputs"]: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) + elif should_send_default_pii() and integration.include_prompts: set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) _calculate_completions_token_usage( @@ -1026,6 +1070,7 @@ async def _wrap_asynchronous_completions_chunk_iterator( ttft = None data_buf: "list[list[str]]" = [] # one for each choice streaming_message_total_token_usage = None + client = sentry_sdk.get_client() async for x in old_iterator: if isinstance(span, StreamedSpan): @@ -1058,7 +1103,10 @@ async def _wrap_asynchronous_completions_chunk_iterator( all_responses = None if len(data_buf) > 0: all_responses = ["".join(chunk) for chunk in data_buf] - if should_send_default_pii() and integration.include_prompts: + if has_data_collection_enabled(client.options): + if client.options["data_collection"]["gen_ai"]["outputs"]: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) + elif should_send_default_pii() and integration.include_prompts: set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) _calculate_completions_token_usage( @@ -1090,6 +1138,7 @@ def _wrap_synchronous_responses_event_iterator( """ ttft = None data_buf: "list[list[str]]" = [] # one for each choice + client = sentry_sdk.get_client() count_tokens_manually = True for x in old_iterator: @@ -1125,7 +1174,10 @@ def _wrap_synchronous_responses_event_iterator( ) if len(data_buf) > 0: all_responses = ["".join(chunk) for chunk in data_buf] - if should_send_default_pii() and integration.include_prompts: + if has_data_collection_enabled(client.options): + if client.options["data_collection"]["gen_ai"]["outputs"]: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) + elif should_send_default_pii() and integration.include_prompts: set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) if count_tokens_manually: @@ -1157,6 +1209,7 @@ async def _wrap_asynchronous_responses_event_iterator( """ ttft: "Optional[float]" = None data_buf: "list[list[str]]" = [] # one for each choice + client = sentry_sdk.get_client() count_tokens_manually = True async for x in old_iterator: @@ -1192,8 +1245,13 @@ async def _wrap_asynchronous_responses_event_iterator( ) if len(data_buf) > 0: all_responses = ["".join(chunk) for chunk in data_buf] - if should_send_default_pii() and integration.include_prompts: + + if has_data_collection_enabled(client.options): + if client.options["data_collection"]["gen_ai"]["outputs"]: + set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) + elif should_send_default_pii() and integration.include_prompts: set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) + if count_tokens_manually: _calculate_responses_token_usage( input=input, diff --git a/tests/integrations/openai/test_openai.py b/tests/integrations/openai/test_openai.py index e88b67e6d1..ebb16f8a02 100644 --- a/tests/integrations/openai/test_openai.py +++ b/tests/integrations/openai/test_openai.py @@ -17,7 +17,11 @@ from openai import AsyncOpenAI, AsyncStream, OpenAI, OpenAIError, Stream from openai.types import CompletionUsage, CreateEmbeddingResponse, Embedding -from openai.types.chat import ChatCompletionChunk, ChatCompletionMessage +from openai.types.chat import ( + ChatCompletion, + ChatCompletionChunk, + ChatCompletionMessage, +) from openai.types.chat.chat_completion import Choice from openai.types.chat.chat_completion_chunk import Choice as DeltaChoice from openai.types.chat.chat_completion_chunk import ChoiceDelta @@ -40,7 +44,9 @@ CustomToolParam, FunctionToolParam, Response, + ResponseFunctionToolCall, ResponseOutputMessage, + ResponseOutputRefusal, ResponseOutputText, ResponseUsage, WebSearchToolParam, @@ -809,36 +815,73 @@ def test_completions_api_data_collection( @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.asyncio @pytest.mark.parametrize( - "send_default_pii, include_prompts", + "data_collection,send_default_pii,expect_output", [ - (True, False), - (False, True), - (False, False), + 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( + None, + False, + False, + id="no-gen-ai-data-collection-falls-back-to-send-default-pii", + ), + pytest.param( + None, + True, + True, + id="no-gen-ai-data-collection-pii-enabled-collects", + ), ], ) -async def test_nonstreaming_chat_completion_async_no_prompts( +def test_completions_api_data_collection_outputs( sentry_init, capture_events, capture_items, + data_collection, send_default_pii, - include_prompts, + expect_output, nonstreaming_chat_completions_model_response, stream_gen_ai_spans, span_streaming, ): - sentry_init( - integrations=[OpenAIIntegration(include_prompts=include_prompts)], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - stream_gen_ai_spans=stream_gen_ai_spans, - trace_lifecycle="stream" if span_streaming else "static", - ) + init_kwargs = { + "integrations": [OpenAIIntegration()], + "disabled_integrations": [StdlibIntegration], + "traces_sample_rate": 1.0, + "send_default_pii": send_default_pii, + "stream_gen_ai_spans": stream_gen_ai_spans, + "trace_lifecycle": "stream" if span_streaming else "static", + } + if data_collection is not None: + init_kwargs["_experiments"] = {"data_collection": data_collection} - client = AsyncOpenAI(api_key="z") - client.chat.completions._post = mock.AsyncMock( + sentry_init(**init_kwargs) + + client = OpenAI(api_key="z") + client.chat.completions._post = mock.Mock( return_value=nonstreaming_chat_completions_model_response( response_id="chat-id", response_model="gpt-3.5-turbo", @@ -856,182 +899,102 @@ async def test_nonstreaming_chat_completion_async_no_prompts( items = capture_items("span") with start_transaction(name="openai tx"): - response = await client.chat.completions.create( + client.chat.completions.create( model="some-model", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "hello"}, - ], - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, + messages=[{"role": "user", "content": "hello"}], ) - response = response.choices[0].message.content - assert response == "the model response" sentry_sdk.flush() - span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.chat" - assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False - - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - - assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in span["attributes"] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] - - assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 - assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 - assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 + (span,) = (item.payload for item in items) + span_data = span["attributes"] else: events = capture_events() with start_transaction(name="openai tx"): - response = await client.chat.completions.create( + client.chat.completions.create( model="some-model", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "hello"}, - ], - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, + messages=[{"role": "user", "content": "hello"}], ) - response = response.choices[0].message.content - - assert response == "the model response" - tx = events[0] - assert tx["type"] == "transaction" - span = tx["spans"][0] - assert span["op"] == "gen_ai.chat" - assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False - assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" - assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 - assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + (transaction,) = events + (span,) = transaction["spans"] + span_data = span["data"] - assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in span["data"] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + assert span_data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-3.5-turbo" - assert span["data"]["gen_ai.usage.output_tokens"] == 10 - assert span["data"]["gen_ai.usage.input_tokens"] == 20 - assert span["data"]["gen_ai.usage.total_tokens"] == 30 + if expect_output: + assert "the model response" in span_data[SPANDATA.GEN_AI_RESPONSE_TEXT] + 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 @pytest.mark.parametrize( - "get_messages,expected_system_instructions", + "data_collection,send_default_pii,expect_output", [ - ( - lambda: [ - { - "role": "system", - "content": "You are a helpful assistant.", - }, - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "hello"}, - ], - [ - { - "type": "text", - "content": "You are a helpful assistant.", - } - ], + pytest.param( + {"gen_ai": {"outputs": True}}, + False, + True, + id="gen-ai-outputs-enabled-overrides-pii-disabled", ), - ( - lambda: [ - { - "role": "system", - "content": [ - {"type": "text", "text": "You are a helpful assistant."}, - {"type": "text", "text": "Be concise and clear."}, - ], - }, - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "hello"}, - ], - [ - { - "type": "text", - "content": "You are a helpful assistant.", - }, - { - "type": "text", - "content": "Be concise and clear.", - }, - ], + pytest.param( + {"gen_ai": {"outputs": False}}, + True, + False, + id="gen-ai-outputs-disabled-overrides-pii-enabled", ), - ( - lambda: iter( - [ - { - "role": "system", - "content": [ - {"type": "text", "text": "You are a helpful assistant."}, - {"type": "text", "text": "Be concise and clear."}, - ], - }, - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "hello"}, - ] - ), - [ - { - "type": "text", - "content": "You are a helpful assistant.", - }, - { - "type": "text", - "content": "Be concise and clear.", - }, - ], + 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( + None, + False, + False, + id="no-gen-ai-data-collection-falls-back-to-send-default-pii", + ), + pytest.param( + None, + True, + True, + id="no-gen-ai-data-collection-pii-enabled-collects", ), ], ) -async def test_nonstreaming_chat_completion_async( +async def test_completions_api_data_collection_outputs_async( sentry_init, capture_events, capture_items, - get_messages, - expected_system_instructions, + data_collection, + send_default_pii, + expect_output, nonstreaming_chat_completions_model_response, stream_gen_ai_spans, span_streaming, ): - sentry_init( - integrations=[OpenAIIntegration(include_prompts=True)], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, - trace_lifecycle="stream" if span_streaming else "static", - ) + init_kwargs = { + "integrations": [OpenAIIntegration()], + "disabled_integrations": [StdlibIntegration], + "traces_sample_rate": 1.0, + "send_default_pii": send_default_pii, + "stream_gen_ai_spans": stream_gen_ai_spans, + "trace_lifecycle": "stream" if span_streaming else "static", + } + if data_collection is not None: + init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**init_kwargs) client = AsyncOpenAI(api_key="z") client.chat.completions._post = AsyncMock( @@ -1052,130 +1015,137 @@ async def test_nonstreaming_chat_completion_async( items = capture_items("span") with start_transaction(name="openai tx"): - response = await client.chat.completions.create( + await client.chat.completions.create( model="some-model", - messages=get_messages(), - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, + messages=[{"role": "user", "content": "hello"}], ) - response = response.choices[0].message.content - assert response == "the model response" sentry_sdk.flush() - span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.chat" - assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False - - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - - assert ( - json.loads(span["attributes"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) - == expected_system_instructions - ) - - assert "hello" in span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert ( - "Message demonstrating the absence of truncation." - in span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - assert "the model response" in span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - - assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 - assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 - assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 + (span,) = (item.payload for item in items) + span_data = span["attributes"] else: events = capture_events() with start_transaction(name="openai tx"): - response = await client.chat.completions.create( + await client.chat.completions.create( model="some-model", - messages=get_messages(), - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, + messages=[{"role": "user", "content": "hello"}], ) - response = response.choices[0].message.content - assert response == "the model response" - tx = events[0] - assert tx["type"] == "transaction" - span = tx["spans"][0] - assert span["op"] == "gen_ai.chat" - assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + (transaction,) = events + (span,) = transaction["spans"] + span_data = span["data"] - assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" - assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 - assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + assert span_data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "gpt-3.5-turbo" - assert ( - json.loads(span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) - == expected_system_instructions + if expect_output: + assert "the model response" in span_data[SPANDATA.GEN_AI_RESPONSE_TEXT] + 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]) +def test_completions_api_data_collection_outputs_empty_choices( + sentry_init, + capture_events, + capture_items, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + _experiments={"data_collection": {"gen_ai": {"outputs": True}}}, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = OpenAI(api_key="z") + client.chat.completions._post = mock.Mock( + return_value=ChatCompletion( + id="chat-id", + choices=[], + created=10000000, + model="gpt-3.5-turbo", + object="chat.completion", + usage=CompletionUsage( + prompt_tokens=20, + completion_tokens=10, + total_tokens=30, + ), ) + ) - assert "hello" in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert "the model response" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") - assert span["data"]["gen_ai.usage.output_tokens"] == 10 - assert span["data"]["gen_ai.usage.input_tokens"] == 20 - assert span["data"]["gen_ai.usage.total_tokens"] == 30 + with start_transaction(name="openai tx"): + client.chat.completions.create( + model="some-model", + messages=[{"role": "user", "content": "hello"}], + ) + sentry_sdk.flush() + (span,) = (item.payload for item in items) + span_data = span["attributes"] + else: + events = capture_events() -def tiktoken_encoding_if_installed(): - try: - import tiktoken # type: ignore # noqa # pylint: disable=unused-import + with start_transaction(name="openai tx"): + client.chat.completions.create( + model="some-model", + messages=[{"role": "user", "content": "hello"}], + ) - return "cl100k_base" - except ImportError: - return None + (transaction,) = events + (span,) = transaction["spans"] + span_data = span["data"] + + # No choices means no output data, even with outputs collection enabled + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span_data -# noinspection PyTypeChecker @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.parametrize( - "send_default_pii, include_prompts", + "data_collection,expect_output", [ - (True, False), - (False, True), - (False, False), + pytest.param( + {"gen_ai": {"outputs": True}}, + True, + id="gen-ai-outputs-enabled", + ), + pytest.param( + {"gen_ai": {"outputs": False}}, + False, + id="gen-ai-outputs-disabled", + ), + pytest.param( + {}, + True, + id="gen-ai-omitted-defaults-to-enabled", + ), ], ) -def test_streaming_chat_completion_no_prompts( +def test_streaming_chat_completion_data_collection_outputs( sentry_init, capture_events, capture_items, - send_default_pii, - include_prompts, + data_collection, + expect_output, get_model_response, server_side_event_chunks, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[ - OpenAIIntegration( - include_prompts=include_prompts, - tiktoken_encoding_name=tiktoken_encoding_if_installed(), - ) - ], + integrations=[OpenAIIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=send_default_pii, + send_default_pii=False, + _experiments={"data_collection": data_collection}, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) @@ -1189,33 +1159,7 @@ def test_streaming_chat_completion_no_prompts( choices=[ DeltaChoice( index=0, - delta=ChoiceDelta(content="hel"), - finish_reason=None, - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=1, - delta=ChoiceDelta(content="lo "), - finish_reason=None, - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=2, - delta=ChoiceDelta(content="world"), + delta=ChoiceDelta(content="hello"), finish_reason="stop", ) ], @@ -1238,49 +1182,17 @@ def test_streaming_chat_completion_no_prompts( ), start_transaction(name="openai tx"): response_stream = client.chat.completions.create( model="some-model", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "hello"}, - ], + messages=[{"role": "user", "content": "hello"}], stream=True, - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, ) response_string = "".join( map(lambda x: x.choices[0].delta.content, response_stream) ) - assert response_string == "hello world" + assert response_string == "hello" sentry_sdk.flush() - span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.chat" - assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - - assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" - - assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in span["attributes"] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] - - try: - import tiktoken # type: ignore # noqa # pylint: disable=unused-import - - assert span["attributes"]["gen_ai.usage.output_tokens"] == 2 - assert span["attributes"]["gen_ai.usage.input_tokens"] == 7 - assert span["attributes"]["gen_ai.usage.total_tokens"] == 9 - except ImportError: - pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly + (span,) = (item.payload for item in items) + span_data = span["attributes"] else: events = capture_events() @@ -1291,114 +1203,90 @@ def test_streaming_chat_completion_no_prompts( ), start_transaction(name="openai tx"): response_stream = client.chat.completions.create( model="some-model", - messages=[ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "hello"}, - ], + messages=[{"role": "user", "content": "hello"}], stream=True, - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, ) response_string = "".join( map(lambda x: x.choices[0].delta.content, response_stream) ) - assert response_string == "hello world" - tx = events[0] - assert tx["type"] == "transaction" - span = tx["spans"][0] - assert span["op"] == "gen_ai.chat" - assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - - assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" - assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 - assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - - assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" + assert response_string == "hello" + (transaction,) = events + (span,) = transaction["spans"] + span_data = span["data"] - assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in span["data"] - assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] - assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] - - try: - import tiktoken # type: ignore # noqa # pylint: disable=unused-import - - assert span["data"]["gen_ai.usage.output_tokens"] == 2 - assert span["data"]["gen_ai.usage.input_tokens"] == 7 - assert span["data"]["gen_ai.usage.total_tokens"] == 9 - except ImportError: - pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly + if expect_output: + assert "hello" in span_data[SPANDATA.GEN_AI_RESPONSE_TEXT] + 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.skipif( - OPENAI_VERSION <= (1, 1, 0), - reason="OpenAI versions <=1.1.0 do not support the stream_options parameter.", +@pytest.mark.asyncio +@pytest.mark.parametrize( + "data_collection,expect_output", + [ + pytest.param( + {"gen_ai": {"outputs": True}}, + True, + id="gen-ai-outputs-enabled", + ), + pytest.param( + {"gen_ai": {"outputs": False}}, + False, + id="gen-ai-outputs-disabled", + ), + pytest.param( + {}, + True, + id="gen-ai-omitted-defaults-to-enabled", + ), + ], ) -def test_streaming_chat_completion_with_usage_in_stream( +async def test_streaming_chat_completion_data_collection_outputs_async( sentry_init, capture_events, capture_items, + data_collection, + expect_output, get_model_response, + async_iterator, server_side_event_chunks, stream_gen_ai_spans, span_streaming, ): - """When stream_options=include_usage is set, token usage comes from the final chunk's usage field.""" sentry_init( - integrations=[OpenAIIntegration(include_prompts=False)], + integrations=[OpenAIIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=False, + _experiments={"data_collection": data_collection}, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) - client = OpenAI(api_key="z") + client = AsyncOpenAI(api_key="z") returned_stream = get_model_response( - server_side_event_chunks( - [ - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, - delta=ChoiceDelta(content="hel"), - finish_reason=None, - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, - delta=ChoiceDelta(content="lo"), - finish_reason="stop", - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - usage=CompletionUsage( - prompt_tokens=20, - completion_tokens=10, - total_tokens=30, + async_iterator( + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="hello"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", ), - ), - ], - include_event_type=False, + ], + include_event_type=False, + ) ) ) @@ -1410,21 +1298,19 @@ def test_streaming_chat_completion_with_usage_in_stream( "send", return_value=returned_stream, ), start_transaction(name="openai tx"): - response_stream = client.chat.completions.create( + response_stream = await client.chat.completions.create( model="some-model", messages=[{"role": "user", "content": "hello"}], stream=True, - stream_options={"include_usage": True}, ) - for _ in response_stream: - pass + response_string = "" + async for x in response_stream: + response_string += x.choices[0].delta.content + assert response_string == "hello" sentry_sdk.flush() - span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.chat" - assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 - assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 - assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 + (span,) = (item.payload for item in items) + span_data = span["attributes"] else: events = capture_events() @@ -1433,341 +1319,157 @@ def test_streaming_chat_completion_with_usage_in_stream( "send", return_value=returned_stream, ), start_transaction(name="openai tx"): - response_stream = client.chat.completions.create( + response_stream = await client.chat.completions.create( model="some-model", messages=[{"role": "user", "content": "hello"}], stream=True, - stream_options={"include_usage": True}, ) - for _ in response_stream: - pass + response_string = "" + async for x in response_stream: + response_string += x.choices[0].delta.content - tx = events[0] - assert tx["type"] == "transaction" - span = tx["spans"][0] - assert span["op"] == "gen_ai.chat" - assert span["data"]["gen_ai.usage.input_tokens"] == 20 - assert span["data"]["gen_ai.usage.output_tokens"] == 10 - assert span["data"]["gen_ai.usage.total_tokens"] == 30 + assert response_string == "hello" + (transaction,) = events + (span,) = transaction["spans"] + span_data = span["data"] + + if expect_output: + assert "hello" in span_data[SPANDATA.GEN_AI_RESPONSE_TEXT] + 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.skipif( - OPENAI_VERSION <= (1, 1, 0), - reason="OpenAI versions <=1.1.0 do not support the stream_options parameter.", +@pytest.mark.asyncio +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [ + (True, False), + (False, True), + (False, False), + ], ) -def test_streaming_chat_completion_empty_content_preserves_token_usage( +async def test_nonstreaming_chat_completion_async_no_prompts( sentry_init, capture_events, capture_items, - get_model_response, - server_side_event_chunks, + send_default_pii, + include_prompts, + nonstreaming_chat_completions_model_response, stream_gen_ai_spans, span_streaming, ): - """Token usage from the stream is recorded even when no content is produced (e.g. content filter).""" sentry_init( - integrations=[OpenAIIntegration(include_prompts=False)], + integrations=[OpenAIIntegration(include_prompts=include_prompts)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=False, + send_default_pii=send_default_pii, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) - client = OpenAI(api_key="z") - returned_stream = get_model_response( - server_side_event_chunks( - [ - ChatCompletionChunk( - id="1", - choices=[], - created=100000, - model="model-id", - object="chat.completion.chunk", - usage=CompletionUsage( - prompt_tokens=20, - completion_tokens=0, - total_tokens=20, - ), - ), - ], - include_event_type=False, + client = AsyncOpenAI(api_key="z") + client.chat.completions._post = mock.AsyncMock( + return_value=nonstreaming_chat_completions_model_response( + response_id="chat-id", + response_model="gpt-3.5-turbo", + message_content="the model response", + created=10000000, + usage=CompletionUsage( + prompt_tokens=20, + completion_tokens=10, + total_tokens=30, + ), ) ) if span_streaming or stream_gen_ai_spans: items = capture_items("span") - with mock.patch.object( - client.chat._client._client, - "send", - return_value=returned_stream, - ), start_transaction(name="openai tx"): - response_stream = client.chat.completions.create( + with start_transaction(name="openai tx"): + response = await client.chat.completions.create( model="some-model", - messages=[{"role": "user", "content": "hello"}], - stream=True, - stream_options={"include_usage": True}, + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "hello"}, + ], + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, ) - for _ in response_stream: - pass + response = response.choices[0].message.content + assert response == "the model response" sentry_sdk.flush() span = next(item.payload for item in items) assert span["attributes"]["sentry.op"] == "gen_ai.chat" + assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + + assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in span["attributes"] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + + assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 - assert "gen_ai.usage.output_tokens" not in span["attributes"] - assert span["attributes"]["gen_ai.usage.total_tokens"] == 20 + assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 else: events = capture_events() - with mock.patch.object( - client.chat._client._client, - "send", - return_value=returned_stream, - ), start_transaction(name="openai tx"): - response_stream = client.chat.completions.create( + with start_transaction(name="openai tx"): + response = await client.chat.completions.create( model="some-model", - messages=[{"role": "user", "content": "hello"}], - stream=True, - stream_options={"include_usage": True}, + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "hello"}, + ], + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, ) - for _ in response_stream: - pass + response = response.choices[0].message.content + assert response == "the model response" tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] assert span["op"] == "gen_ai.chat" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False + + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 + assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 + assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 + assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + + assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in span["data"] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + + assert span["data"]["gen_ai.usage.output_tokens"] == 10 assert span["data"]["gen_ai.usage.input_tokens"] == 20 - assert "gen_ai.usage.output_tokens" not in span["data"] - assert span["data"]["gen_ai.usage.total_tokens"] == 20 + assert span["data"]["gen_ai.usage.total_tokens"] == 30 @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.skipif( - OPENAI_VERSION <= (1, 1, 0), - reason="OpenAI versions <=1.1.0 do not support the stream_options parameter.", -) -@pytest.mark.asyncio -async def test_streaming_chat_completion_empty_content_preserves_token_usage_async( - sentry_init, - capture_events, - capture_items, - get_model_response, - async_iterator, - server_side_event_chunks, - stream_gen_ai_spans, - span_streaming, -): - """Token usage from the stream is recorded even when no content is produced - async variant.""" - sentry_init( - integrations=[OpenAIIntegration(include_prompts=False)], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - send_default_pii=False, - stream_gen_ai_spans=stream_gen_ai_spans, - trace_lifecycle="stream" if span_streaming else "static", - ) - - client = AsyncOpenAI(api_key="z") - returned_stream = get_model_response( - async_iterator( - server_side_event_chunks( - [ - ChatCompletionChunk( - id="1", - choices=[], - created=100000, - model="model-id", - object="chat.completion.chunk", - usage=CompletionUsage( - prompt_tokens=20, - completion_tokens=0, - total_tokens=20, - ), - ), - ], - include_event_type=False, - ) - ) - ) - - if span_streaming or stream_gen_ai_spans: - items = capture_items("span") - - with mock.patch.object( - client.chat._client._client, - "send", - return_value=returned_stream, - ), start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( - model="some-model", - messages=[{"role": "user", "content": "hello"}], - stream=True, - stream_options={"include_usage": True}, - ) - async for _ in response_stream: - pass - - sentry_sdk.flush() - span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.chat" - assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 - assert "gen_ai.usage.output_tokens" not in span["attributes"] - assert span["attributes"]["gen_ai.usage.total_tokens"] == 20 - else: - events = capture_events() - - with mock.patch.object( - client.chat._client._client, - "send", - return_value=returned_stream, - ), start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( - model="some-model", - messages=[{"role": "user", "content": "hello"}], - stream=True, - stream_options={"include_usage": True}, - ) - async for _ in response_stream: - pass - - tx = events[0] - assert tx["type"] == "transaction" - span = tx["spans"][0] - assert span["op"] == "gen_ai.chat" - assert span["data"]["gen_ai.usage.input_tokens"] == 20 - assert "gen_ai.usage.output_tokens" not in span["data"] - assert span["data"]["gen_ai.usage.total_tokens"] == 20 - - -@pytest.mark.parametrize("span_streaming", [True, False]) -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.skipif( - OPENAI_VERSION <= (1, 1, 0), - reason="OpenAI versions <=1.1.0 do not support the stream_options parameter.", -) @pytest.mark.asyncio -async def test_streaming_chat_completion_async_with_usage_in_stream( - sentry_init, - capture_events, - capture_items, - get_model_response, - async_iterator, - server_side_event_chunks, - stream_gen_ai_spans, - span_streaming, -): - """When stream_options=include_usage is set, token usage comes from the final chunk's usage field (async).""" - sentry_init( - integrations=[OpenAIIntegration(include_prompts=False)], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - send_default_pii=False, - stream_gen_ai_spans=stream_gen_ai_spans, - trace_lifecycle="stream" if span_streaming else "static", - ) - - client = AsyncOpenAI(api_key="z") - returned_stream = get_model_response( - async_iterator( - server_side_event_chunks( - [ - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, - delta=ChoiceDelta(content="hel"), - finish_reason=None, - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, - delta=ChoiceDelta(content="lo"), - finish_reason="stop", - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - usage=CompletionUsage( - prompt_tokens=20, - completion_tokens=10, - total_tokens=30, - ), - ), - ], - include_event_type=False, - ) - ) - ) - - if span_streaming or stream_gen_ai_spans: - items = capture_items("span") - - with mock.patch.object( - client.chat._client._client, - "send", - return_value=returned_stream, - ), start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( - model="some-model", - messages=[{"role": "user", "content": "hello"}], - stream=True, - stream_options={"include_usage": True}, - ) - async for _ in response_stream: - pass - - sentry_sdk.flush() - span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.chat" - assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 - assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 - assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 - else: - events = capture_events() - - with mock.patch.object( - client.chat._client._client, - "send", - return_value=returned_stream, - ), start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( - model="some-model", - messages=[{"role": "user", "content": "hello"}], - stream=True, - stream_options={"include_usage": True}, - ) - async for _ in response_stream: - pass - - tx = events[0] - assert tx["type"] == "transaction" - span = tx["spans"][0] - assert span["op"] == "gen_ai.chat" - assert span["data"]["gen_ai.usage.input_tokens"] == 20 - assert span["data"]["gen_ai.usage.output_tokens"] == 10 - assert span["data"]["gen_ai.usage.total_tokens"] == 30 - - -# noinspection PyTypeChecker -@pytest.mark.parametrize("span_streaming", [True, False]) -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.parametrize( - "get_messages,expected_system_instructions,expected_output_tokens,expected_input_tokens", + "get_messages,expected_system_instructions", [ ( lambda: [ @@ -1787,8 +1489,6 @@ async def test_streaming_chat_completion_async_with_usage_in_stream( "content": "You are a helpful assistant.", } ], - 2, - 15, ), ( lambda: [ @@ -1815,8 +1515,6 @@ async def test_streaming_chat_completion_async_with_usage_in_stream( "content": "Be concise and clear.", }, ], - 2, - 20, ), ( lambda: iter( @@ -1845,31 +1543,21 @@ async def test_streaming_chat_completion_async_with_usage_in_stream( "content": "Be concise and clear.", }, ], - 2, - 20, ), ], ) -def test_streaming_chat_completion( +async def test_nonstreaming_chat_completion_async( sentry_init, capture_events, capture_items, get_messages, expected_system_instructions, - expected_output_tokens, - expected_input_tokens, - get_model_response, - server_side_event_chunks, + nonstreaming_chat_completions_model_response, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[ - OpenAIIntegration( - include_prompts=True, - tiktoken_encoding_name=tiktoken_encoding_if_installed(), - ) - ], + integrations=[OpenAIIntegration(include_prompts=True)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, @@ -1877,81 +1565,42 @@ def test_streaming_chat_completion( trace_lifecycle="stream" if span_streaming else "static", ) - client = OpenAI(api_key="z") - returned_stream = get_model_response( - server_side_event_chunks( - [ - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, - delta=ChoiceDelta(content="hel"), - finish_reason=None, - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=1, - delta=ChoiceDelta(content="lo "), - finish_reason=None, - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=2, - delta=ChoiceDelta(content="world"), - finish_reason="stop", - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ], - include_event_type=False, + client = AsyncOpenAI(api_key="z") + client.chat.completions._post = AsyncMock( + return_value=nonstreaming_chat_completions_model_response( + response_id="chat-id", + response_model="gpt-3.5-turbo", + message_content="the model response", + created=10000000, + usage=CompletionUsage( + prompt_tokens=20, + completion_tokens=10, + total_tokens=30, + ), ) ) if span_streaming or stream_gen_ai_spans: items = capture_items("span") - with mock.patch.object( - client.chat._client._client, - "send", - return_value=returned_stream, - ), start_transaction(name="openai tx"): - response_stream = client.chat.completions.create( + with start_transaction(name="openai tx"): + response = await client.chat.completions.create( model="some-model", messages=get_messages(), - stream=True, max_tokens=100, presence_penalty=0.1, frequency_penalty=0.2, temperature=0.7, top_p=0.9, ) - response_string = "".join( - map(lambda x: x.choices[0].delta.content, response_stream) - ) - assert response_string == "hello world" + response = response.choices[0].message.content + + assert response == "the model response" sentry_sdk.flush() span = next(item.payload for item in items) assert span["attributes"]["sentry.op"] == "gen_ai.chat" assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 @@ -1965,60 +1614,38 @@ def test_streaming_chat_completion( == expected_system_instructions ) - assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" - + assert "hello" in span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] assert ( "Message demonstrating the absence of truncation." in span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] ) - assert "hello" in span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert "hello world" in span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - - try: - import tiktoken # type: ignore # noqa # pylint: disable=unused-import - - assert ( - span["attributes"]["gen_ai.usage.output_tokens"] - == expected_output_tokens - ) - assert ( - span["attributes"]["gen_ai.usage.input_tokens"] == expected_input_tokens - ) - assert ( - span["attributes"]["gen_ai.usage.total_tokens"] - == expected_output_tokens + expected_input_tokens - ) + assert "the model response" in span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - except ImportError: - pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly + assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 + assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 + assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 else: events = capture_events() - with mock.patch.object( - client.chat._client._client, - "send", - return_value=returned_stream, - ), start_transaction(name="openai tx"): - response_stream = client.chat.completions.create( + with start_transaction(name="openai tx"): + response = await client.chat.completions.create( model="some-model", messages=get_messages(), - stream=True, max_tokens=100, presence_penalty=0.1, frequency_penalty=0.2, temperature=0.7, top_p=0.9, ) - response_string = "".join( - map(lambda x: x.choices[0].delta.content, response_stream) - ) - assert response_string == "hello world" + response = response.choices[0].message.content + + assert response == "the model response" tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] assert span["op"] == "gen_ai.chat" assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is False assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 @@ -2032,29 +1659,26 @@ def test_streaming_chat_completion( == expected_system_instructions ) - assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" - assert "hello" in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert "hello world" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] + assert "the model response" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] - try: - import tiktoken # type: ignore # noqa # pylint: disable=unused-import + assert span["data"]["gen_ai.usage.output_tokens"] == 10 + assert span["data"]["gen_ai.usage.input_tokens"] == 20 + assert span["data"]["gen_ai.usage.total_tokens"] == 30 - assert span["data"]["gen_ai.usage.output_tokens"] == expected_output_tokens - assert span["data"]["gen_ai.usage.input_tokens"] == expected_input_tokens - assert ( - span["data"]["gen_ai.usage.total_tokens"] - == expected_output_tokens + expected_input_tokens - ) - except ImportError: - pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly +def tiktoken_encoding_if_installed(): + try: + import tiktoken # type: ignore # noqa # pylint: disable=unused-import + + return "cl100k_base" + except ImportError: + return None # noinspection PyTypeChecker @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.asyncio @pytest.mark.parametrize( "send_default_pii, include_prompts", [ @@ -2063,14 +1687,13 @@ def test_streaming_chat_completion( (False, False), ], ) -async def test_streaming_chat_completion_async_no_prompts( +def test_streaming_chat_completion_no_prompts( sentry_init, capture_events, capture_items, send_default_pii, include_prompts, get_model_response, - async_iterator, server_side_event_chunks, stream_gen_ai_spans, span_streaming, @@ -2089,53 +1712,51 @@ async def test_streaming_chat_completion_async_no_prompts( trace_lifecycle="stream" if span_streaming else "static", ) - client = AsyncOpenAI(api_key="z") + client = OpenAI(api_key="z") returned_stream = get_model_response( - async_iterator( - server_side_event_chunks( - [ - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, - delta=ChoiceDelta(content="hel"), - finish_reason=None, - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=1, - delta=ChoiceDelta(content="lo "), - finish_reason=None, - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=2, - delta=ChoiceDelta(content="world"), - finish_reason="stop", - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ], - include_event_type=False, - ) + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="hel"), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=1, + delta=ChoiceDelta(content="lo "), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=2, + delta=ChoiceDelta(content="world"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ], + include_event_type=False, ) ) @@ -2147,7 +1768,7 @@ async def test_streaming_chat_completion_async_no_prompts( "send", return_value=returned_stream, ), start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( + response_stream = client.chat.completions.create( model="some-model", messages=[ {"role": "system", "content": "You are a helpful assistant."}, @@ -2160,10 +1781,9 @@ async def test_streaming_chat_completion_async_no_prompts( temperature=0.7, top_p=0.9, ) - - response_string = "" - async for x in response_stream: - response_string += x.choices[0].delta.content + response_string = "".join( + map(lambda x: x.choices[0].delta.content, response_stream) + ) assert response_string == "hello world" sentry_sdk.flush() @@ -2191,7 +1811,6 @@ async def test_streaming_chat_completion_async_no_prompts( assert span["attributes"]["gen_ai.usage.output_tokens"] == 2 assert span["attributes"]["gen_ai.usage.input_tokens"] == 7 assert span["attributes"]["gen_ai.usage.total_tokens"] == 9 - except ImportError: pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly else: @@ -2202,7 +1821,7 @@ async def test_streaming_chat_completion_async_no_prompts( "send", return_value=returned_stream, ), start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( + response_stream = client.chat.completions.create( model="some-model", messages=[ {"role": "system", "content": "You are a helpful assistant."}, @@ -2215,10 +1834,9 @@ async def test_streaming_chat_completion_async_no_prompts( temperature=0.7, top_p=0.9, ) - - response_string = "" - async for x in response_stream: - response_string += x.choices[0].delta.content + response_string = "".join( + map(lambda x: x.choices[0].delta.content, response_stream) + ) assert response_string == "hello world" tx = events[0] @@ -2247,175 +1865,72 @@ async def test_streaming_chat_completion_async_no_prompts( assert span["data"]["gen_ai.usage.output_tokens"] == 2 assert span["data"]["gen_ai.usage.input_tokens"] == 7 assert span["data"]["gen_ai.usage.total_tokens"] == 9 - except ImportError: pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly -# noinspection PyTypeChecker @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "get_messages,expected_system_instructions,expected_output_tokens,expected_input_tokens", - [ - ( - lambda: [ - { - "role": "system", - "content": "You are a helpful assistant.", - }, - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "hello"}, - ], - [ - { - "type": "text", - "content": "You are a helpful assistant.", - } - ], - 2, - 15, - ), - ( - lambda: [ - { - "role": "system", - "content": [ - {"type": "text", "text": "You are a helpful assistant."}, - {"type": "text", "text": "Be concise and clear."}, - ], - }, - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "hello"}, - ], - [ - { - "type": "text", - "content": "You are a helpful assistant.", - }, - { - "type": "text", - "content": "Be concise and clear.", - }, - ], - 2, - 20, - ), - ( - lambda: iter( - [ - { - "role": "system", - "content": [ - {"type": "text", "text": "You are a helpful assistant."}, - {"type": "text", "text": "Be concise and clear."}, - ], - }, - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "hello"}, - ] - ), - [ - { - "type": "text", - "content": "You are a helpful assistant.", - }, - { - "type": "text", - "content": "Be concise and clear.", - }, - ], - 2, - 20, - ), - ], +@pytest.mark.skipif( + OPENAI_VERSION <= (1, 1, 0), + reason="OpenAI versions <=1.1.0 do not support the stream_options parameter.", ) -async def test_streaming_chat_completion_async( +def test_streaming_chat_completion_with_usage_in_stream( sentry_init, capture_events, capture_items, - get_messages, - expected_system_instructions, - expected_output_tokens, - expected_input_tokens, get_model_response, - async_iterator, server_side_event_chunks, stream_gen_ai_spans, span_streaming, ): + """When stream_options=include_usage is set, token usage comes from the final chunk's usage field.""" sentry_init( - integrations=[ - OpenAIIntegration( - include_prompts=True, - tiktoken_encoding_name=tiktoken_encoding_if_installed(), - ) - ], + integrations=[OpenAIIntegration(include_prompts=False)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + send_default_pii=False, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) - client = AsyncOpenAI(api_key="z") - + client = OpenAI(api_key="z") returned_stream = get_model_response( - async_iterator( - server_side_event_chunks( - [ - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, - delta=ChoiceDelta(content="hel"), - finish_reason=None, - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=1, - delta=ChoiceDelta(content="lo "), - finish_reason=None, - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=2, - delta=ChoiceDelta(content="world"), - finish_reason="stop", - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="hel"), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="lo"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + usage=CompletionUsage( + prompt_tokens=20, + completion_tokens=10, + total_tokens=30, ), - ], - include_event_type=False, - ) + ), + ], + include_event_type=False, ) ) @@ -2427,66 +1942,21 @@ async def test_streaming_chat_completion_async( "send", return_value=returned_stream, ), start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( + response_stream = client.chat.completions.create( model="some-model", - messages=get_messages(), + messages=[{"role": "user", "content": "hello"}], stream=True, - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, + stream_options={"include_usage": True}, ) + for _ in response_stream: + pass - response_string = "" - async for x in response_stream: - response_string += x.choices[0].delta.content - - assert response_string == "hello world" sentry_sdk.flush() span = next(item.payload for item in items) assert span["attributes"]["sentry.op"] == "gen_ai.chat" - assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - - assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" - - assert ( - json.loads(span["attributes"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) - == expected_system_instructions - ) - - assert ( - "Message demonstrating the absence of truncation." - in span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - ) - assert "hello" in span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert "hello world" in span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - - try: - import tiktoken # type: ignore # noqa # pylint: disable=unused-import - - assert ( - span["attributes"]["gen_ai.usage.output_tokens"] - == expected_output_tokens - ) - assert ( - span["attributes"]["gen_ai.usage.input_tokens"] == expected_input_tokens - ) - assert ( - span["attributes"]["gen_ai.usage.total_tokens"] - == expected_output_tokens + expected_input_tokens - ) - - except ImportError: - pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly + assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 + assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 + assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 else: events = capture_events() @@ -2495,418 +1965,443 @@ async def test_streaming_chat_completion_async( "send", return_value=returned_stream, ), start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( + response_stream = client.chat.completions.create( model="some-model", - messages=get_messages(), + messages=[{"role": "user", "content": "hello"}], stream=True, - max_tokens=100, - presence_penalty=0.1, - frequency_penalty=0.2, - temperature=0.7, - top_p=0.9, + stream_options={"include_usage": True}, ) + for _ in response_stream: + pass - response_string = "" - async for x in response_stream: - response_string += x.choices[0].delta.content - - assert response_string == "hello world" tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] assert span["op"] == "gen_ai.chat" - assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - - assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" - assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 - assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 - assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 - assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - - assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" - - assert ( - json.loads(span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) - == expected_system_instructions - ) - - assert "hello" in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] - assert "hello world" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] - - try: - import tiktoken # type: ignore # noqa # pylint: disable=unused-import - - assert span["data"]["gen_ai.usage.output_tokens"] == expected_output_tokens - assert span["data"]["gen_ai.usage.input_tokens"] == expected_input_tokens - assert ( - span["data"]["gen_ai.usage.total_tokens"] - == expected_output_tokens + expected_input_tokens - ) - - except ImportError: - pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly + assert span["data"]["gen_ai.usage.input_tokens"] == 20 + assert span["data"]["gen_ai.usage.output_tokens"] == 10 + assert span["data"]["gen_ai.usage.total_tokens"] == 30 @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_bad_chat_completion( +@pytest.mark.skipif( + OPENAI_VERSION <= (1, 1, 0), + reason="OpenAI versions <=1.1.0 do not support the stream_options parameter.", +) +def test_streaming_chat_completion_empty_content_preserves_token_usage( sentry_init, capture_events, capture_items, + get_model_response, + server_side_event_chunks, stream_gen_ai_spans, span_streaming, ): + """Token usage from the stream is recorded even when no content is produced (e.g. content filter).""" sentry_init( - integrations=[OpenAIIntegration()], + integrations=[OpenAIIntegration(include_prompts=False)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + send_default_pii=False, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) - if span_streaming: - items = capture_items("event", "span") - - client = OpenAI(api_key="z") - client.chat.completions._post = mock.Mock( - side_effect=OpenAIError("API rate limit reached") + client = OpenAI(api_key="z") + returned_stream = get_model_response( + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[], + created=100000, + model="model-id", + object="chat.completion.chunk", + usage=CompletionUsage( + prompt_tokens=20, + completion_tokens=0, + total_tokens=20, + ), + ), + ], + include_event_type=False, ) - with pytest.raises(OpenAIError): - client.chat.completions.create( - model="some-model", - messages=[{"role": "system", "content": "hello"}], - ) + ) - (event,) = (item.payload for item in items if item.type == "event") - sentry_sdk.flush() - (span,) = (item.payload for item in items if item.type == "span") - assert event["level"] == "error" - assert span["status"] == "error" - elif stream_gen_ai_spans: - items = capture_items("event", "transaction") + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") - client = OpenAI(api_key="z") - client.chat.completions._post = mock.Mock( - side_effect=OpenAIError("API rate limit reached") - ) - with pytest.raises(OpenAIError): - client.chat.completions.create( + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = client.chat.completions.create( model="some-model", - messages=[{"role": "system", "content": "hello"}], + messages=[{"role": "user", "content": "hello"}], + stream=True, + stream_options={"include_usage": True}, ) + for _ in response_stream: + pass - (event,) = (item.payload for item in items if item.type == "event") - (transaction,) = (item.payload for item in items if item.type == "transaction") - assert event["level"] == "error" - assert transaction["contexts"]["trace"]["status"] == "internal_error" + sentry_sdk.flush() + span = next(item.payload for item in items) + assert span["attributes"]["sentry.op"] == "gen_ai.chat" + assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 + assert "gen_ai.usage.output_tokens" not in span["attributes"] + assert span["attributes"]["gen_ai.usage.total_tokens"] == 20 else: events = capture_events() - client = OpenAI(api_key="z") - client.chat.completions._post = mock.Mock( - side_effect=OpenAIError("API rate limit reached") - ) - with pytest.raises(OpenAIError): - client.chat.completions.create( + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = client.chat.completions.create( model="some-model", - messages=[{"role": "system", "content": "hello"}], + messages=[{"role": "user", "content": "hello"}], + stream=True, + stream_options={"include_usage": True}, ) + for _ in response_stream: + pass - (event, transaction) = events - assert event["level"] == "error" - assert transaction["contexts"]["trace"]["status"] == "internal_error" + tx = events[0] + assert tx["type"] == "transaction" + span = tx["spans"][0] + assert span["op"] == "gen_ai.chat" + assert span["data"]["gen_ai.usage.input_tokens"] == 20 + assert "gen_ai.usage.output_tokens" not in span["data"] + assert span["data"]["gen_ai.usage.total_tokens"] == 20 @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_span_status_error( +@pytest.mark.skipif( + OPENAI_VERSION <= (1, 1, 0), + reason="OpenAI versions <=1.1.0 do not support the stream_options parameter.", +) +@pytest.mark.asyncio +async def test_streaming_chat_completion_empty_content_preserves_token_usage_async( sentry_init, capture_events, capture_items, + get_model_response, + async_iterator, + server_side_event_chunks, stream_gen_ai_spans, span_streaming, ): + """Token usage from the stream is recorded even when no content is produced - async variant.""" sentry_init( - integrations=[OpenAIIntegration()], + integrations=[OpenAIIntegration(include_prompts=False)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + send_default_pii=False, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) + client = AsyncOpenAI(api_key="z") + returned_stream = get_model_response( + async_iterator( + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[], + created=100000, + model="model-id", + object="chat.completion.chunk", + usage=CompletionUsage( + prompt_tokens=20, + completion_tokens=0, + total_tokens=20, + ), + ), + ], + include_event_type=False, + ) + ) + ) + if span_streaming or stream_gen_ai_spans: - items = capture_items("event", "span") + items = capture_items("span") - with start_transaction(name="test"): - client = OpenAI(api_key="z") - client.chat.completions._post = mock.Mock( - side_effect=OpenAIError("API rate limit reached") + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", + messages=[{"role": "user", "content": "hello"}], + stream=True, + stream_options={"include_usage": True}, ) - with pytest.raises(OpenAIError): - client.chat.completions.create( - model="some-model", - messages=[{"role": "system", "content": "hello"}], - ) - - (error,) = (item.payload for item in items if item.type == "event") - assert error["level"] == "error" + async for _ in response_stream: + pass sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[0]["status"] == "error" + span = next(item.payload for item in items) + assert span["attributes"]["sentry.op"] == "gen_ai.chat" + assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 + assert "gen_ai.usage.output_tokens" not in span["attributes"] + assert span["attributes"]["gen_ai.usage.total_tokens"] == 20 else: events = capture_events() - with start_transaction(name="test"): - client = OpenAI(api_key="z") - client.chat.completions._post = mock.Mock( - side_effect=OpenAIError("API rate limit reached") + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", + messages=[{"role": "user", "content": "hello"}], + stream=True, + stream_options={"include_usage": True}, ) - with pytest.raises(OpenAIError): - client.chat.completions.create( - model="some-model", - messages=[{"role": "system", "content": "hello"}], - ) + async for _ in response_stream: + pass - (error, transaction) = events - assert error["level"] == "error" - assert transaction["spans"][0]["status"] == "internal_error" - assert transaction["spans"][0]["tags"]["status"] == "internal_error" + tx = events[0] + assert tx["type"] == "transaction" + span = tx["spans"][0] + assert span["op"] == "gen_ai.chat" + assert span["data"]["gen_ai.usage.input_tokens"] == 20 + assert "gen_ai.usage.output_tokens" not in span["data"] + assert span["data"]["gen_ai.usage.total_tokens"] == 20 @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.skipif( + OPENAI_VERSION <= (1, 1, 0), + reason="OpenAI versions <=1.1.0 do not support the stream_options parameter.", +) @pytest.mark.asyncio -async def test_bad_chat_completion_async( +async def test_streaming_chat_completion_async_with_usage_in_stream( sentry_init, capture_events, capture_items, + get_model_response, + async_iterator, + server_side_event_chunks, stream_gen_ai_spans, span_streaming, ): + """When stream_options=include_usage is set, token usage comes from the final chunk's usage field (async).""" sentry_init( - integrations=[OpenAIIntegration()], + integrations=[OpenAIIntegration(include_prompts=False)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + send_default_pii=False, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) client = AsyncOpenAI(api_key="z") - client.chat.completions._post = AsyncMock( - side_effect=OpenAIError("API rate limit reached") - ) - - if span_streaming: - items = capture_items("event", "span") - - with pytest.raises(OpenAIError): - await client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] - ) - - (event,) = (item.payload for item in items if item.type == "event") - sentry_sdk.flush() - (span,) = (item.payload for item in items if item.type == "span") - assert event["level"] == "error" - assert span["status"] == "error" - elif stream_gen_ai_spans: - items = capture_items("event", "transaction") - - with pytest.raises(OpenAIError): - await client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] - ) - - (event,) = (item.payload for item in items if item.type == "event") - (transaction,) = (item.payload for item in items if item.type == "transaction") - assert event["level"] == "error" - assert transaction["contexts"]["trace"]["status"] == "internal_error" - else: - events = capture_events() - - with pytest.raises(OpenAIError): - await client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] + returned_stream = get_model_response( + async_iterator( + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="hel"), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="lo"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + usage=CompletionUsage( + prompt_tokens=20, + completion_tokens=10, + total_tokens=30, + ), + ), + ], + include_event_type=False, ) - - (event, transaction) = events - assert event["level"] == "error" - assert transaction["contexts"]["trace"]["status"] == "internal_error" - - -@pytest.mark.parametrize("span_streaming", [True, False]) -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.parametrize( - "send_default_pii, include_prompts", - [ - (True, False), - (False, True), - (False, False), - ], -) -def test_embeddings_create_no_pii( - sentry_init, - capture_events, - capture_items, - send_default_pii, - include_prompts, - stream_gen_ai_spans, - span_streaming, -): - sentry_init( - integrations=[OpenAIIntegration(include_prompts=include_prompts)], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - stream_gen_ai_spans=stream_gen_ai_spans, - trace_lifecycle="stream" if span_streaming else "static", - ) - - client = OpenAI(api_key="z") - - returned_embedding = CreateEmbeddingResponse( - data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], - model="some-model", - object="list", - usage=EmbeddingTokenUsage( - prompt_tokens=20, - total_tokens=30, - ), + ) ) - client.embeddings._post = mock.Mock(return_value=returned_embedding) - if span_streaming or stream_gen_ai_spans: items = capture_items("span") - with start_transaction(name="openai tx"): - response = client.embeddings.create( - input="hello", model="text-embedding-3-large" + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", + messages=[{"role": "user", "content": "hello"}], + stream=True, + stream_options={"include_usage": True}, ) - - assert len(response.data[0].embedding) == 3 + async for _ in response_stream: + pass sentry_sdk.flush() span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" - assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert ( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] - == "text-embedding-3-large" - ) - - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["attributes"] - + assert span["attributes"]["sentry.op"] == "gen_ai.chat" assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 + assert span["attributes"]["gen_ai.usage.output_tokens"] == 10 assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 else: events = capture_events() - with start_transaction(name="openai tx"): - response = client.embeddings.create( - input="hello", model="text-embedding-3-large" + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", + messages=[{"role": "user", "content": "hello"}], + stream=True, + stream_options={"include_usage": True}, ) - - assert len(response.data[0].embedding) == 3 + async for _ in response_stream: + pass tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] - assert span["op"] == "gen_ai.embeddings" - assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" - - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["data"] - + assert span["op"] == "gen_ai.chat" assert span["data"]["gen_ai.usage.input_tokens"] == 20 + assert span["data"]["gen_ai.usage.output_tokens"] == 10 assert span["data"]["gen_ai.usage.total_tokens"] == 30 +# noinspection PyTypeChecker @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.parametrize( - "get_input,expected_embeddings_input", + "get_messages,expected_system_instructions,expected_output_tokens,expected_input_tokens", [ ( - lambda: "hello", - ["hello"], - ), - ( - lambda: ["First text", "Second text", "Third text"], - [ - "First text", - "Second text", - "Third text", + lambda: [ + { + "role": "system", + "content": "You are a helpful assistant.", + }, + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "hello"}, ], - ), - ( - lambda: iter(["First text", "Second text", "Third text"]), [ - "First text", - "Second text", - "Third text", + { + "type": "text", + "content": "You are a helpful assistant.", + } ], + 2, + 15, ), ( - lambda: [5, 8, 13, 21, 34], + lambda: [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Be concise and clear."}, + ], + }, + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "hello"}, + ], [ - 5, - 8, - 13, - 21, - 34, + { + "type": "text", + "content": "You are a helpful assistant.", + }, + { + "type": "text", + "content": "Be concise and clear.", + }, ], + 2, + 20, ), ( lambda: iter( - [5, 8, 13, 21, 34], + [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Be concise and clear."}, + ], + }, + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "hello"}, + ] ), [ - 5, - 8, - 13, - 21, - 34, - ], - ), - ( - lambda: [ - [5, 8, 13, 21, 34], - [8, 13, 21, 34, 55], - ], - [ - [5, 8, 13, 21, 34], - [8, 13, 21, 34, 55], - ], - ), - ( - lambda: iter( - [ - [5, 8, 13, 21, 34], - [8, 13, 21, 34, 55], - ] - ), - [ - [5, 8, 13, 21, 34], - [8, 13, 21, 34, 55], + { + "type": "text", + "content": "You are a helpful assistant.", + }, + { + "type": "text", + "content": "Be concise and clear.", + }, ], + 2, + 20, ), ], ) -def test_embeddings_create( +def test_streaming_chat_completion( sentry_init, capture_events, capture_items, - get_input, - expected_embeddings_input, + get_messages, + expected_system_instructions, + expected_output_tokens, + expected_input_tokens, + get_model_response, + server_side_event_chunks, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[OpenAIIntegration(include_prompts=True)], + integrations=[ + OpenAIIntegration( + include_prompts=True, + tiktoken_encoding_name=tiktoken_encoding_if_installed(), + ) + ], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, @@ -2915,293 +2410,210 @@ def test_embeddings_create( ) client = OpenAI(api_key="z") - - returned_embedding = CreateEmbeddingResponse( - data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], - model="some-model", - object="list", - usage=EmbeddingTokenUsage( - prompt_tokens=20, - total_tokens=30, - ), + returned_stream = get_model_response( + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="hel"), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=1, + delta=ChoiceDelta(content="lo "), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=2, + delta=ChoiceDelta(content="world"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ], + include_event_type=False, + ) ) - client.embeddings._post = mock.Mock(return_value=returned_embedding) - if span_streaming or stream_gen_ai_spans: items = capture_items("span") - with start_transaction(name="openai tx"): - response = client.embeddings.create( - input=get_input(), model="text-embedding-3-large" + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = client.chat.completions.create( + model="some-model", + messages=get_messages(), + stream=True, + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, ) - - assert len(response.data[0].embedding) == 3 - + response_string = "".join( + map(lambda x: x.choices[0].delta.content, response_stream) + ) + assert response_string == "hello world" sentry_sdk.flush() span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" + assert span["attributes"]["sentry.op"] == "gen_ai.chat" assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 + assert ( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] - == "text-embedding-3-large" + json.loads(span["attributes"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) + == expected_system_instructions ) + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" + assert ( - json.loads(span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) - == expected_embeddings_input + "Message demonstrating the absence of truncation." + in span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] ) + assert "hello" in span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert "hello world" in span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 - assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 - else: - events = capture_events() + try: + import tiktoken # type: ignore # noqa # pylint: disable=unused-import - with start_transaction(name="openai tx"): - response = client.embeddings.create( - input=get_input(), model="text-embedding-3-large" + assert ( + span["attributes"]["gen_ai.usage.output_tokens"] + == expected_output_tokens + ) + assert ( + span["attributes"]["gen_ai.usage.input_tokens"] == expected_input_tokens + ) + assert ( + span["attributes"]["gen_ai.usage.total_tokens"] + == expected_output_tokens + expected_input_tokens ) - assert len(response.data[0].embedding) == 3 + except ImportError: + pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly + else: + events = capture_events() + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = client.chat.completions.create( + model="some-model", + messages=get_messages(), + stream=True, + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, + ) + response_string = "".join( + map(lambda x: x.choices[0].delta.content, response_stream) + ) + assert response_string == "hello world" tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] - assert span["op"] == "gen_ai.embeddings" + assert span["op"] == "gen_ai.chat" assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True + + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 + assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 + assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 + assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 assert ( - json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) - == expected_embeddings_input + json.loads(span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) + == expected_system_instructions ) - assert span["data"]["gen_ai.usage.input_tokens"] == 20 - assert span["data"]["gen_ai.usage.total_tokens"] == 30 - - -def _collect_embeddings_span_data( - capture_events, capture_items, span_streaming, stream_gen_ai_spans, create -): - if span_streaming or stream_gen_ai_spans: - items = capture_items("span") - - with start_transaction(name="openai tx"): - response = create() - - assert len(response.data[0].embedding) == 3 - - sentry_sdk.flush() - span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" - return span["attributes"] + assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" - events = capture_events() + assert "hello" in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert "hello world" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] - with start_transaction(name="openai tx"): - response = create() + try: + import tiktoken # type: ignore # noqa # pylint: disable=unused-import - assert len(response.data[0].embedding) == 3 + assert span["data"]["gen_ai.usage.output_tokens"] == expected_output_tokens + assert span["data"]["gen_ai.usage.input_tokens"] == expected_input_tokens + assert ( + span["data"]["gen_ai.usage.total_tokens"] + == expected_output_tokens + expected_input_tokens + ) - tx = events[0] - assert tx["type"] == "transaction" - span = tx["spans"][0] - assert span["op"] == "gen_ai.embeddings" - return span["data"] + except ImportError: + pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly +# noinspection PyTypeChecker @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio @pytest.mark.parametrize( - "data_collection,send_default_pii,include_prompts,expect_input", - [ - pytest.param( - {"gen_ai": {"inputs": True}}, - False, - True, - True, - id="inputs-enabled-overrides-pii-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": False}}, - True, - True, - False, - id="inputs-disabled-overrides-pii-enabled", - ), - pytest.param( - {}, - False, - True, - True, - id="gen-ai-omitted-defaults-to-enabled", - ), - pytest.param( - {"gen_ai": {"inputs": False}}, - False, - True, - False, - id="inputs-disabled-and-pii-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": True}}, - True, - False, - False, - id="include-prompts-disabled-overrides-inputs-enabled", - ), - pytest.param( - None, - False, - True, - False, - id="no-experiment-falls-back-to-pii", - ), - ], -) -def test_embeddings_create_data_collection( - sentry_init, - capture_events, - capture_items, - data_collection, - send_default_pii, - include_prompts, - expect_input, - stream_gen_ai_spans, - span_streaming, -): - init_kwargs = { - "integrations": [OpenAIIntegration(include_prompts=include_prompts)], - "disabled_integrations": [StdlibIntegration], - "traces_sample_rate": 1.0, - "send_default_pii": send_default_pii, - "stream_gen_ai_spans": stream_gen_ai_spans, - "trace_lifecycle": "stream" if span_streaming else "static", - } - - sentry_init_kwargs = dict(init_kwargs) - if data_collection is not None: - sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} - - sentry_init(**sentry_init_kwargs) - - client = OpenAI(api_key="z") - - returned_embedding = CreateEmbeddingResponse( - data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], - model="some-model", - object="list", - usage=EmbeddingTokenUsage( - prompt_tokens=20, - total_tokens=30, - ), - ) - - client.embeddings._post = mock.Mock(return_value=returned_embedding) - - span_data = _collect_embeddings_span_data( - capture_events, - capture_items, - span_streaming, - stream_gen_ai_spans, - lambda: client.embeddings.create(input="hello", model="text-embedding-3-large"), - ) - - assert span_data[SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" - - if expect_input: - assert json.loads(span_data[SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == ["hello"] - else: - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span_data - - assert span_data["gen_ai.usage.input_tokens"] == 20 - assert span_data["gen_ai.usage.total_tokens"] == 30 - - -@pytest.mark.parametrize("span_streaming", [True, False]) -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.parametrize( - "get_input", - [ - lambda: "hello", - lambda: ["First text", "Second text"], - lambda: iter(["First text", "Second text"]), - lambda: [5, 8, 13, 21, 34], - lambda: [[5, 8, 13], [8, 13, 21]], - lambda: {"text": "hello"}, - ], -) -def test_embeddings_create_data_collection_inputs_disabled_input_shapes( - sentry_init, - capture_events, - capture_items, - get_input, - stream_gen_ai_spans, - span_streaming, -): - sentry_init( - integrations=[OpenAIIntegration(include_prompts=True)], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - send_default_pii=True, - stream_gen_ai_spans=stream_gen_ai_spans, - trace_lifecycle="stream" if span_streaming else "static", - _experiments={"data_collection": {"gen_ai": {"inputs": False}}}, - ) - - client = OpenAI(api_key="z") - - returned_embedding = CreateEmbeddingResponse( - data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], - model="some-model", - object="list", - usage=EmbeddingTokenUsage( - prompt_tokens=20, - total_tokens=30, - ), - ) - - client.embeddings._post = mock.Mock(return_value=returned_embedding) - - span_data = _collect_embeddings_span_data( - capture_events, - capture_items, - span_streaming, - stream_gen_ai_spans, - lambda: client.embeddings.create( - input=get_input(), model="text-embedding-3-large" - ), - ) - - assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span_data - - -@pytest.mark.parametrize("span_streaming", [True, False]) -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "send_default_pii, include_prompts", + "send_default_pii, include_prompts", [ (True, False), (False, True), (False, False), ], ) -async def test_embeddings_create_async_no_pii( +async def test_streaming_chat_completion_async_no_prompts( sentry_init, capture_events, capture_items, send_default_pii, include_prompts, + get_model_response, + async_iterator, + server_side_event_chunks, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[OpenAIIntegration(include_prompts=include_prompts)], + integrations=[ + OpenAIIntegration( + include_prompts=include_prompts, + tiktoken_encoding_name=tiktoken_encoding_if_installed(), + ) + ], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=send_default_pii, @@ -3210,148 +2622,277 @@ async def test_embeddings_create_async_no_pii( ) client = AsyncOpenAI(api_key="z") - - returned_embedding = CreateEmbeddingResponse( - data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], - model="some-model", - object="list", - usage=EmbeddingTokenUsage( - prompt_tokens=20, - total_tokens=30, - ), + returned_stream = get_model_response( + async_iterator( + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="hel"), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=1, + delta=ChoiceDelta(content="lo "), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=2, + delta=ChoiceDelta(content="world"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ], + include_event_type=False, + ) + ) ) - client.embeddings._post = AsyncMock(return_value=returned_embedding) - if span_streaming or stream_gen_ai_spans: items = capture_items("span") - with start_transaction(name="openai tx"): - response = await client.embeddings.create( - input="hello", model="text-embedding-3-large" + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "hello"}, + ], + stream=True, + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, ) - assert len(response.data[0].embedding) == 3 + response_string = "" + async for x in response_stream: + response_string += x.choices[0].delta.content + assert response_string == "hello world" sentry_sdk.flush() span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" + assert span["attributes"]["sentry.op"] == "gen_ai.chat" assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert ( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] - == "text-embedding-3-large" - ) + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["attributes"] + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 - assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" + + assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in span["attributes"] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["attributes"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["attributes"] + + try: + import tiktoken # type: ignore # noqa # pylint: disable=unused-import + + assert span["attributes"]["gen_ai.usage.output_tokens"] == 2 + assert span["attributes"]["gen_ai.usage.input_tokens"] == 7 + assert span["attributes"]["gen_ai.usage.total_tokens"] == 9 + + except ImportError: + pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly else: events = capture_events() - with start_transaction(name="openai tx"): - response = await client.embeddings.create( - input="hello", model="text-embedding-3-large" + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", + messages=[ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "hello"}, + ], + stream=True, + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, ) - assert len(response.data[0].embedding) == 3 + response_string = "" + async for x in response_stream: + response_string += x.choices[0].delta.content + assert response_string == "hello world" tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] - assert span["op"] == "gen_ai.embeddings" + assert span["op"] == "gen_ai.chat" assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["data"] + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 + assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 + assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 + assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - assert span["data"]["gen_ai.usage.input_tokens"] == 20 - assert span["data"]["gen_ai.usage.total_tokens"] == 30 + assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" + + assert SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS not in span["data"] + assert SPANDATA.GEN_AI_REQUEST_MESSAGES not in span["data"] + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span["data"] + + try: + import tiktoken # type: ignore # noqa # pylint: disable=unused-import + + assert span["data"]["gen_ai.usage.output_tokens"] == 2 + assert span["data"]["gen_ai.usage.input_tokens"] == 7 + assert span["data"]["gen_ai.usage.total_tokens"] == 9 + except ImportError: + pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly + +# noinspection PyTypeChecker @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio @pytest.mark.parametrize( - "get_input,expected_embeddings_input", + "get_messages,expected_system_instructions,expected_output_tokens,expected_input_tokens", [ ( - lambda: "hello", - ["hello"], - ), - ( - lambda: ["First text", "Second text", "Third text"], - [ - "First text", - "Second text", - "Third text", - ], - ), - ( - lambda: iter(["First text", "Second text", "Third text"]), - [ - "First text", - "Second text", - "Third text", - ], - ), - ( - lambda: [5, 8, 13, 21, 34], - [ - 5, - 8, - 13, - 21, - 34, + lambda: [ + { + "role": "system", + "content": "You are a helpful assistant.", + }, + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "hello"}, ], - ), - ( - lambda: iter( - [5, 8, 13, 21, 34], - ), [ - 5, - 8, - 13, - 21, - 34, + { + "type": "text", + "content": "You are a helpful assistant.", + } ], + 2, + 15, ), ( lambda: [ - [5, 8, 13, 21, 34], - [8, 13, 21, 34, 55], + { + "role": "system", + "content": [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Be concise and clear."}, + ], + }, + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "hello"}, ], [ - [5, 8, 13, 21, 34], - [8, 13, 21, 34, 55], + { + "type": "text", + "content": "You are a helpful assistant.", + }, + { + "type": "text", + "content": "Be concise and clear.", + }, ], + 2, + 20, ), ( lambda: iter( [ - [5, 8, 13, 21, 34], - [8, 13, 21, 34, 55], + { + "role": "system", + "content": [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Be concise and clear."}, + ], + }, + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "hello"}, ] ), [ - [5, 8, 13, 21, 34], - [8, 13, 21, 34, 55], + { + "type": "text", + "content": "You are a helpful assistant.", + }, + { + "type": "text", + "content": "Be concise and clear.", + }, ], + 2, + 20, ), ], ) -async def test_embeddings_create_async( +async def test_streaming_chat_completion_async( sentry_init, capture_events, capture_items, - get_input, - expected_embeddings_input, + get_messages, + expected_system_instructions, + expected_output_tokens, + expected_input_tokens, + get_model_response, + async_iterator, + server_side_event_chunks, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[OpenAIIntegration(include_prompts=True)], + integrations=[ + OpenAIIntegration( + include_prompts=True, + tiktoken_encoding_name=tiktoken_encoding_if_installed(), + ) + ], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, send_default_pii=True, @@ -3361,295 +2902,213 @@ async def test_embeddings_create_async( client = AsyncOpenAI(api_key="z") - returned_embedding = CreateEmbeddingResponse( - data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], - model="some-model", - object="list", - usage=EmbeddingTokenUsage( - prompt_tokens=20, - total_tokens=30, - ), - ) - - client.embeddings._post = AsyncMock(return_value=returned_embedding) - - if span_streaming or stream_gen_ai_spans: - items = capture_items("span") - - with start_transaction(name="openai tx"): - response = await client.embeddings.create( - input=get_input(), model="text-embedding-3-large" - ) - - assert len(response.data[0].embedding) == 3 - - sentry_sdk.flush() - span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" - assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert ( - span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] - == "text-embedding-3-large" - ) - - assert ( - json.loads(span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) - == expected_embeddings_input - ) - - assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 - assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 - else: - events = capture_events() - - with start_transaction(name="openai tx"): - response = await client.embeddings.create( - input=get_input(), model="text-embedding-3-large" - ) - - assert len(response.data[0].embedding) == 3 - - tx = events[0] - assert tx["type"] == "transaction" - span = tx["spans"][0] - assert span["op"] == "gen_ai.embeddings" - assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" - - assert ( - json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) - == expected_embeddings_input + returned_stream = get_model_response( + async_iterator( + server_side_event_chunks( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, + delta=ChoiceDelta(content="hel"), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=1, + delta=ChoiceDelta(content="lo "), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=2, + delta=ChoiceDelta(content="world"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ], + include_event_type=False, + ) ) + ) - assert span["data"]["gen_ai.usage.input_tokens"] == 20 - assert span["data"]["gen_ai.usage.total_tokens"] == 30 + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", + messages=get_messages(), + stream=True, + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, + ) -@pytest.mark.parametrize("span_streaming", [True, False]) -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "data_collection,send_default_pii,include_prompts,expect_input", - [ - pytest.param( - {"gen_ai": {"inputs": True}}, - False, - True, - True, - id="inputs-enabled-overrides-pii-disabled", - ), - pytest.param( - {"gen_ai": {"inputs": False}}, - True, - True, - False, - id="inputs-disabled-overrides-pii-enabled", - ), - pytest.param( - {}, - False, - True, - True, - id="gen-ai-omitted-defaults-to-enabled", - ), - pytest.param( - {"gen_ai": {"inputs": True}}, - True, - False, - False, - id="include-prompts-disabled-overrides-inputs-enabled", - ), - pytest.param( - None, - False, - True, - False, - id="no-experiment-falls-back-to-pii", - ), - ], -) -async def test_embeddings_create_async_data_collection( - sentry_init, - capture_events, - capture_items, - data_collection, - send_default_pii, - include_prompts, - expect_input, - stream_gen_ai_spans, - span_streaming, -): - init_kwargs = { - "integrations": [OpenAIIntegration(include_prompts=include_prompts)], - "disabled_integrations": [StdlibIntegration], - "traces_sample_rate": 1.0, - "send_default_pii": send_default_pii, - "stream_gen_ai_spans": stream_gen_ai_spans, - "trace_lifecycle": "stream" if span_streaming else "static", - } + response_string = "" + async for x in response_stream: + response_string += x.choices[0].delta.content - sentry_init_kwargs = dict(init_kwargs) - if data_collection is not None: - sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} + assert response_string == "hello world" + sentry_sdk.flush() + span = next(item.payload for item in items) + assert span["attributes"]["sentry.op"] == "gen_ai.chat" + assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - sentry_init(**sentry_init_kwargs) + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 + assert span["attributes"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - client = AsyncOpenAI(api_key="z") + assert span["attributes"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" - returned_embedding = CreateEmbeddingResponse( - data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], - model="some-model", - object="list", - usage=EmbeddingTokenUsage( - prompt_tokens=20, - total_tokens=30, - ), - ) + assert ( + json.loads(span["attributes"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) + == expected_system_instructions + ) - client.embeddings._post = AsyncMock(return_value=returned_embedding) + assert ( + "Message demonstrating the absence of truncation." + in span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + ) + assert "hello" in span["attributes"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert "hello world" in span["attributes"][SPANDATA.GEN_AI_RESPONSE_TEXT] - if span_streaming or stream_gen_ai_spans: - items = capture_items("span") + try: + import tiktoken # type: ignore # noqa # pylint: disable=unused-import - with start_transaction(name="openai tx"): - response = await client.embeddings.create( - input="hello", model="text-embedding-3-large" + assert ( + span["attributes"]["gen_ai.usage.output_tokens"] + == expected_output_tokens + ) + assert ( + span["attributes"]["gen_ai.usage.input_tokens"] == expected_input_tokens + ) + assert ( + span["attributes"]["gen_ai.usage.total_tokens"] + == expected_output_tokens + expected_input_tokens ) - assert len(response.data[0].embedding) == 3 - - sentry_sdk.flush() - span = next(item.payload for item in items) - assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" - span_data = span["attributes"] + except ImportError: + pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly else: events = capture_events() - with start_transaction(name="openai tx"): - response = await client.embeddings.create( - input="hello", model="text-embedding-3-large" + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", + messages=get_messages(), + stream=True, + max_tokens=100, + presence_penalty=0.1, + frequency_penalty=0.2, + temperature=0.7, + top_p=0.9, ) - assert len(response.data[0].embedding) == 3 + response_string = "" + async for x in response_stream: + response_string += x.choices[0].delta.content + assert response_string == "hello world" tx = events[0] assert tx["type"] == "transaction" span = tx["spans"][0] - assert span["op"] == "gen_ai.embeddings" - span_data = span["data"] + assert span["op"] == "gen_ai.chat" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert span["data"][SPANDATA.GEN_AI_RESPONSE_STREAMING] is True - assert span_data[SPANDATA.GEN_AI_SYSTEM] == "openai" - assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" - assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "some-model" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MAX_TOKENS] == 100 + assert span["data"][SPANDATA.GEN_AI_REQUEST_PRESENCE_PENALTY] == 0.1 + assert span["data"][SPANDATA.GEN_AI_REQUEST_FREQUENCY_PENALTY] == 0.2 + assert span["data"][SPANDATA.GEN_AI_REQUEST_TEMPERATURE] == 0.7 + assert span["data"][SPANDATA.GEN_AI_REQUEST_TOP_P] == 0.9 - if expect_input: - assert json.loads(span_data[SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == ["hello"] - else: - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span_data + assert span["data"][SPANDATA.GEN_AI_RESPONSE_MODEL] == "model-id" - assert span_data["gen_ai.usage.input_tokens"] == 20 - assert span_data["gen_ai.usage.total_tokens"] == 30 + assert ( + json.loads(span["data"][SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS]) + == expected_system_instructions + ) + assert "hello" in span["data"][SPANDATA.GEN_AI_REQUEST_MESSAGES] + assert "hello world" in span["data"][SPANDATA.GEN_AI_RESPONSE_TEXT] -@pytest.mark.parametrize("span_streaming", [True, False]) -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.parametrize( - "send_default_pii, include_prompts", - [(True, True), (True, False), (False, True), (False, False)], -) -def test_embeddings_create_raises_error( - sentry_init, - capture_events, - capture_items, - send_default_pii, - include_prompts, - stream_gen_ai_spans, - span_streaming, -): - sentry_init( - integrations=[OpenAIIntegration(include_prompts=include_prompts)], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - send_default_pii=send_default_pii, - stream_gen_ai_spans=stream_gen_ai_spans, - trace_lifecycle="stream" if span_streaming else "static", - ) - - client = OpenAI(api_key="z") - - client.embeddings._post = mock.Mock( - side_effect=OpenAIError("API rate limit reached") - ) - - if span_streaming: - items = capture_items("event", "span") - - with pytest.raises(OpenAIError): - client.embeddings.create(input="hello", model="text-embedding-3-large") - - (event,) = (item.payload for item in items if item.type == "event") - sentry_sdk.flush() - (span,) = (item.payload for item in items if item.type == "span") - assert event["level"] == "error" - assert span["status"] == "error" - elif stream_gen_ai_spans: - items = capture_items("event", "transaction") - - with pytest.raises(OpenAIError): - client.embeddings.create(input="hello", model="text-embedding-3-large") - - (event,) = (item.payload for item in items if item.type == "event") - (transaction,) = (item.payload for item in items if item.type == "transaction") - assert event["level"] == "error" - assert transaction["contexts"]["trace"]["status"] == "internal_error" - else: - events = capture_events() + try: + import tiktoken # type: ignore # noqa # pylint: disable=unused-import - with pytest.raises(OpenAIError): - client.embeddings.create(input="hello", model="text-embedding-3-large") + assert span["data"]["gen_ai.usage.output_tokens"] == expected_output_tokens + assert span["data"]["gen_ai.usage.input_tokens"] == expected_input_tokens + assert ( + span["data"]["gen_ai.usage.total_tokens"] + == expected_output_tokens + expected_input_tokens + ) - (event, transaction) = events - assert event["level"] == "error" - assert transaction["contexts"]["trace"]["status"] == "internal_error" + except ImportError: + pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "send_default_pii, include_prompts", - [(True, True), (True, False), (False, True), (False, False)], -) -async def test_embeddings_create_raises_error_async( +def test_bad_chat_completion( sentry_init, capture_events, capture_items, - send_default_pii, - include_prompts, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[OpenAIIntegration(include_prompts=include_prompts)], + integrations=[OpenAIIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=send_default_pii, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) - client = AsyncOpenAI(api_key="z") - - client.embeddings._post = AsyncMock( - side_effect=OpenAIError("API rate limit reached") - ) - if span_streaming: items = capture_items("event", "span") + client = OpenAI(api_key="z") + client.chat.completions._post = mock.Mock( + side_effect=OpenAIError("API rate limit reached") + ) with pytest.raises(OpenAIError): - await client.embeddings.create( - input="hello", model="text-embedding-3-large" + client.chat.completions.create( + model="some-model", + messages=[{"role": "system", "content": "hello"}], ) (event,) = (item.payload for item in items if item.type == "event") @@ -3660,9 +3119,14 @@ async def test_embeddings_create_raises_error_async( elif stream_gen_ai_spans: items = capture_items("event", "transaction") + client = OpenAI(api_key="z") + client.chat.completions._post = mock.Mock( + side_effect=OpenAIError("API rate limit reached") + ) with pytest.raises(OpenAIError): - await client.embeddings.create( - input="hello", model="text-embedding-3-large" + client.chat.completions.create( + model="some-model", + messages=[{"role": "system", "content": "hello"}], ) (event,) = (item.payload for item in items if item.type == "event") @@ -3672,9 +3136,14 @@ async def test_embeddings_create_raises_error_async( else: events = capture_events() + client = OpenAI(api_key="z") + client.chat.completions._post = mock.Mock( + side_effect=OpenAIError("API rate limit reached") + ) with pytest.raises(OpenAIError): - await client.embeddings.create( - input="hello", model="text-embedding-3-large" + client.chat.completions.create( + model="some-model", + messages=[{"role": "system", "content": "hello"}], ) (event, transaction) = events @@ -3684,84 +3153,68 @@ async def test_embeddings_create_raises_error_async( @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_span_origin_nonstreaming_chat( +def test_span_status_error( sentry_init, capture_events, capture_items, - nonstreaming_chat_completions_model_response, stream_gen_ai_spans, span_streaming, ): sentry_init( integrations=[OpenAIIntegration()], + disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) - client = OpenAI(api_key="z") - client.chat.completions._post = mock.Mock( - return_value=nonstreaming_chat_completions_model_response( - response_id="chat-id", - response_model="gpt-3.5-turbo", - message_content="the model response", - created=10000000, - usage=CompletionUsage( - prompt_tokens=20, - completion_tokens=10, - total_tokens=30, - ), - ) - ) - - if span_streaming: - items = capture_items("transaction", "span") - - with sentry_sdk.traces.start_span(name="openai tx"): - client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] - ) - - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[1]["attributes"]["sentry.origin"] == "manual" - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" - elif stream_gen_ai_spans: - items = capture_items("transaction", "span") + if span_streaming or stream_gen_ai_spans: + items = capture_items("event", "span") - with start_transaction(name="openai tx"): - client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] + with start_transaction(name="test"): + client = OpenAI(api_key="z") + client.chat.completions._post = mock.Mock( + side_effect=OpenAIError("API rate limit reached") ) + with pytest.raises(OpenAIError): + client.chat.completions.create( + model="some-model", + messages=[{"role": "system", "content": "hello"}], + ) - (event,) = (item.payload for item in items if item.type == "transaction") - assert event["contexts"]["trace"]["origin"] == "manual" + (error,) = (item.payload for item in items if item.type == "event") + assert error["level"] == "error" sentry_sdk.flush() spans = [item.payload for item in items if item.type == "span"] - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + assert spans[0]["status"] == "error" else: events = capture_events() - with start_transaction(name="openai tx"): - client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] + with start_transaction(name="test"): + client = OpenAI(api_key="z") + client.chat.completions._post = mock.Mock( + side_effect=OpenAIError("API rate limit reached") ) + with pytest.raises(OpenAIError): + client.chat.completions.create( + model="some-model", + messages=[{"role": "system", "content": "hello"}], + ) - (event,) = events - - assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["origin"] == "auto.ai.openai" + (error, transaction) = events + assert error["level"] == "error" + assert transaction["spans"][0]["status"] == "internal_error" + assert transaction["spans"][0]["tags"]["status"] == "internal_error" @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.asyncio -async def test_span_origin_nonstreaming_chat_async( +async def test_bad_chat_completion_async( sentry_init, capture_events, capture_items, - nonstreaming_chat_completions_model_response, stream_gen_ai_spans, span_streaming, ): @@ -3775,294 +3228,394 @@ async def test_span_origin_nonstreaming_chat_async( client = AsyncOpenAI(api_key="z") client.chat.completions._post = AsyncMock( - return_value=nonstreaming_chat_completions_model_response( - response_id="chat-id", - response_model="gpt-3.5-turbo", - message_content="the model response", - created=10000000, - usage=CompletionUsage( - prompt_tokens=20, - completion_tokens=10, - total_tokens=30, - ), - ) + side_effect=OpenAIError("API rate limit reached") ) if span_streaming: - items = capture_items("transaction", "span") + items = capture_items("event", "span") - with sentry_sdk.traces.start_span(name="openai tx"): + with pytest.raises(OpenAIError): await client.chat.completions.create( model="some-model", messages=[{"role": "system", "content": "hello"}] ) + (event,) = (item.payload for item in items if item.type == "event") sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[1]["attributes"]["sentry.origin"] == "manual" - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + (span,) = (item.payload for item in items if item.type == "span") + assert event["level"] == "error" + assert span["status"] == "error" elif stream_gen_ai_spans: - items = capture_items("transaction", "span") + items = capture_items("event", "transaction") - with start_transaction(name="openai tx"): + with pytest.raises(OpenAIError): await client.chat.completions.create( model="some-model", messages=[{"role": "system", "content": "hello"}] ) - (event,) = (item.payload for item in items if item.type == "transaction") - assert event["contexts"]["trace"]["origin"] == "manual" - - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + (event,) = (item.payload for item in items if item.type == "event") + (transaction,) = (item.payload for item in items if item.type == "transaction") + assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" else: events = capture_events() - with start_transaction(name="openai tx"): + with pytest.raises(OpenAIError): await client.chat.completions.create( model="some-model", messages=[{"role": "system", "content": "hello"}] ) - (event,) = events - - assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["origin"] == "auto.ai.openai" + (event, transaction) = events + assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_span_origin_streaming_chat( +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [ + (True, False), + (False, True), + (False, False), + ], +) +def test_embeddings_create_no_pii( sentry_init, capture_events, capture_items, + send_default_pii, + include_prompts, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[OpenAIIntegration()], + integrations=[OpenAIIntegration(include_prompts=include_prompts)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + send_default_pii=send_default_pii, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) client = OpenAI(api_key="z") - returned_stream = Stream(cast_to=None, response=None, client=client) - returned_stream._iterator = [ - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, delta=ChoiceDelta(content="hel"), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=1, delta=ChoiceDelta(content="lo "), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=2, delta=ChoiceDelta(content="world"), finish_reason="stop" - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ] - if span_streaming: - items = capture_items("transaction", "span") - - client.chat.completions._post = mock.Mock(return_value=returned_stream) - with sentry_sdk.traces.start_span(name="openai tx"): - response_stream = client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] - ) + returned_embedding = CreateEmbeddingResponse( + data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], + model="some-model", + object="list", + usage=EmbeddingTokenUsage( + prompt_tokens=20, + total_tokens=30, + ), + ) - "".join(map(lambda x: x.choices[0].delta.content, response_stream)) + client.embeddings._post = mock.Mock(return_value=returned_embedding) - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[1]["attributes"]["sentry.origin"] == "manual" - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" - elif stream_gen_ai_spans: - items = capture_items("transaction", "span") + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") - client.chat.completions._post = mock.Mock(return_value=returned_stream) with start_transaction(name="openai tx"): - response_stream = client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] + response = client.embeddings.create( + input="hello", model="text-embedding-3-large" ) - "".join(map(lambda x: x.choices[0].delta.content, response_stream)) - - (event,) = (item.payload for item in items if item.type == "transaction") - assert event["contexts"]["trace"]["origin"] == "manual" + assert len(response.data[0].embedding) == 3 sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + span = next(item.payload for item in items) + assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" + assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert ( + span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] + == "text-embedding-3-large" + ) + + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["attributes"] + + assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 + assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 else: events = capture_events() - client.chat.completions._post = mock.Mock(return_value=returned_stream) with start_transaction(name="openai tx"): - response_stream = client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] + response = client.embeddings.create( + input="hello", model="text-embedding-3-large" ) - "".join(map(lambda x: x.choices[0].delta.content, response_stream)) + assert len(response.data[0].embedding) == 3 - (event,) = events + tx = events[0] + assert tx["type"] == "transaction" + span = tx["spans"][0] + assert span["op"] == "gen_ai.embeddings" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" - assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["origin"] == "auto.ai.openai" + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["data"] + + assert span["data"]["gen_ai.usage.input_tokens"] == 20 + assert span["data"]["gen_ai.usage.total_tokens"] == 30 @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.asyncio -async def test_span_origin_streaming_chat_async( +@pytest.mark.parametrize( + "get_input,expected_embeddings_input", + [ + ( + lambda: "hello", + ["hello"], + ), + ( + lambda: ["First text", "Second text", "Third text"], + [ + "First text", + "Second text", + "Third text", + ], + ), + ( + lambda: iter(["First text", "Second text", "Third text"]), + [ + "First text", + "Second text", + "Third text", + ], + ), + ( + lambda: [5, 8, 13, 21, 34], + [ + 5, + 8, + 13, + 21, + 34, + ], + ), + ( + lambda: iter( + [5, 8, 13, 21, 34], + ), + [ + 5, + 8, + 13, + 21, + 34, + ], + ), + ( + lambda: [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ], + [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ], + ), + ( + lambda: iter( + [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ] + ), + [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ], + ), + ], +) +def test_embeddings_create( sentry_init, capture_events, capture_items, - async_iterator, + get_input, + expected_embeddings_input, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[OpenAIIntegration()], + integrations=[OpenAIIntegration(include_prompts=True)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) - client = AsyncOpenAI(api_key="z") - returned_stream = AsyncStream(cast_to=None, response=None, client=client) - returned_stream._iterator = async_iterator( - [ - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=0, delta=ChoiceDelta(content="hel"), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=1, delta=ChoiceDelta(content="lo "), finish_reason=None - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ChatCompletionChunk( - id="1", - choices=[ - DeltaChoice( - index=2, - delta=ChoiceDelta(content="world"), - finish_reason="stop", - ) - ], - created=100000, - model="model-id", - object="chat.completion.chunk", - ), - ] + client = OpenAI(api_key="z") + + returned_embedding = CreateEmbeddingResponse( + data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], + model="some-model", + object="list", + usage=EmbeddingTokenUsage( + prompt_tokens=20, + total_tokens=30, + ), ) - client.chat.completions._post = AsyncMock(return_value=returned_stream) + client.embeddings._post = mock.Mock(return_value=returned_embedding) - if span_streaming: - items = capture_items("transaction", "span") + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") - with sentry_sdk.traces.start_span(name="openai tx"): - response_stream = await client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] + with start_transaction(name="openai tx"): + response = client.embeddings.create( + input=get_input(), model="text-embedding-3-large" ) - async for _ in response_stream: - pass - # "".join(map(lambda x: x.choices[0].delta.content, response_stream)) + assert len(response.data[0].embedding) == 3 sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[1]["attributes"]["sentry.origin"] == "manual" - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" - elif stream_gen_ai_spans: - items = capture_items("transaction", "span") + span = next(item.payload for item in items) + assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" + assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert ( + span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] + == "text-embedding-3-large" + ) - with start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] - ) - async for _ in response_stream: - pass - - # "".join(map(lambda x: x.choices[0].delta.content, response_stream)) - - (event,) = (item.payload for item in items if item.type == "transaction") - assert event["contexts"]["trace"]["origin"] == "manual" + assert ( + json.loads(span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) + == expected_embeddings_input + ) - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 + assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 else: events = capture_events() with start_transaction(name="openai tx"): - response_stream = await client.chat.completions.create( - model="some-model", messages=[{"role": "system", "content": "hello"}] + response = client.embeddings.create( + input=get_input(), model="text-embedding-3-large" ) - async for _ in response_stream: - pass - # "".join(map(lambda x: x.choices[0].delta.content, response_stream)) + assert len(response.data[0].embedding) == 3 - (event,) = events + tx = events[0] + assert tx["type"] == "transaction" + span = tx["spans"][0] + assert span["op"] == "gen_ai.embeddings" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" - assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["origin"] == "auto.ai.openai" + assert ( + json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) + == expected_embeddings_input + ) + + assert span["data"]["gen_ai.usage.input_tokens"] == 20 + assert span["data"]["gen_ai.usage.total_tokens"] == 30 + + +def _collect_embeddings_span_data( + capture_events, capture_items, span_streaming, stream_gen_ai_spans, create +): + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + + with start_transaction(name="openai tx"): + response = create() + + assert len(response.data[0].embedding) == 3 + + sentry_sdk.flush() + span = next(item.payload for item in items) + assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" + return span["attributes"] + + events = capture_events() + + with start_transaction(name="openai tx"): + response = create() + + assert len(response.data[0].embedding) == 3 + + tx = events[0] + assert tx["type"] == "transaction" + span = tx["spans"][0] + assert span["op"] == "gen_ai.embeddings" + return span["data"] @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -def test_span_origin_embeddings( +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,expect_input", + [ + pytest.param( + {"gen_ai": {"inputs": True}}, + False, + True, + True, + id="inputs-enabled-overrides-pii-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False}}, + True, + True, + False, + id="inputs-disabled-overrides-pii-enabled", + ), + pytest.param( + {}, + False, + True, + True, + id="gen-ai-omitted-defaults-to-enabled", + ), + pytest.param( + {"gen_ai": {"inputs": False}}, + False, + True, + False, + id="inputs-disabled-and-pii-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": True}}, + True, + False, + False, + id="include-prompts-disabled-overrides-inputs-enabled", + ), + pytest.param( + None, + False, + True, + False, + id="no-experiment-falls-back-to-pii", + ), + ], +) +def test_embeddings_create_data_collection( sentry_init, capture_events, capture_items, + data_collection, + send_default_pii, + include_prompts, + expect_input, stream_gen_ai_spans, span_streaming, ): - sentry_init( - integrations=[OpenAIIntegration()], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - stream_gen_ai_spans=stream_gen_ai_spans, - trace_lifecycle="stream" if span_streaming else "static", - ) + init_kwargs = { + "integrations": [OpenAIIntegration(include_prompts=include_prompts)], + "disabled_integrations": [StdlibIntegration], + "traces_sample_rate": 1.0, + "send_default_pii": send_default_pii, + "stream_gen_ai_spans": stream_gen_ai_spans, + "trace_lifecycle": "stream" if span_streaming else "static", + } + + sentry_init_kwargs = dict(init_kwargs) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**sentry_init_kwargs) client = OpenAI(api_key="z") @@ -4078,59 +3631,59 @@ def test_span_origin_embeddings( client.embeddings._post = mock.Mock(return_value=returned_embedding) - if span_streaming: - items = capture_items("transaction", "span") - - with sentry_sdk.traces.start_span(name="openai tx"): - client.embeddings.create(input="hello", model="text-embedding-3-large") - - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[1]["attributes"]["sentry.origin"] == "manual" - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" - elif stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with start_transaction(name="openai tx"): - client.embeddings.create(input="hello", model="text-embedding-3-large") + span_data = _collect_embeddings_span_data( + capture_events, + capture_items, + span_streaming, + stream_gen_ai_spans, + lambda: client.embeddings.create(input="hello", model="text-embedding-3-large"), + ) - (event,) = [item.payload for item in items if item.type == "transaction"] - assert event["contexts"]["trace"]["origin"] == "manual" + assert span_data[SPANDATA.GEN_AI_SYSTEM] == "openai" + assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + if expect_input: + assert json.loads(span_data[SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == ["hello"] else: - events = capture_events() - - with start_transaction(name="openai tx"): - client.embeddings.create(input="hello", model="text-embedding-3-large") - - (event,) = events + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span_data - assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["origin"] == "auto.ai.openai" + assert span_data["gen_ai.usage.input_tokens"] == 20 + assert span_data["gen_ai.usage.total_tokens"] == 30 @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.asyncio -async def test_span_origin_embeddings_async( +@pytest.mark.parametrize( + "get_input", + [ + lambda: "hello", + lambda: ["First text", "Second text"], + lambda: iter(["First text", "Second text"]), + lambda: [5, 8, 13, 21, 34], + lambda: [[5, 8, 13], [8, 13, 21]], + lambda: {"text": "hello"}, + ], +) +def test_embeddings_create_data_collection_inputs_disabled_input_shapes( sentry_init, capture_events, capture_items, + get_input, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[OpenAIIntegration()], + integrations=[OpenAIIntegration(include_prompts=True)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", + _experiments={"data_collection": {"gen_ai": {"inputs": False}}}, ) - client = AsyncOpenAI(api_key="z") + client = OpenAI(api_key="z") returned_embedding = CreateEmbeddingResponse( data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], @@ -4142,424 +3695,1897 @@ async def test_span_origin_embeddings_async( ), ) - client.embeddings._post = AsyncMock(return_value=returned_embedding) - - if span_streaming: - items = capture_items("transaction", "span") - - with sentry_sdk.traces.start_span(name="openai tx"): - await client.embeddings.create( - input="hello", model="text-embedding-3-large" - ) - - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[1]["attributes"]["sentry.origin"] == "manual" - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" - elif stream_gen_ai_spans: - items = capture_items("transaction", "span") - - with start_transaction(name="openai tx"): - await client.embeddings.create( - input="hello", model="text-embedding-3-large" - ) + client.embeddings._post = mock.Mock(return_value=returned_embedding) - (event,) = [item.payload for item in items if item.type == "transaction"] - assert event["contexts"]["trace"]["origin"] == "manual" + span_data = _collect_embeddings_span_data( + capture_events, + capture_items, + span_streaming, + stream_gen_ai_spans, + lambda: client.embeddings.create( + input=get_input(), model="text-embedding-3-large" + ), + ) - sentry_sdk.flush() - spans = [item.payload for item in items if item.type == "span"] - assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" - else: - events = capture_events() + assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span_data - with start_transaction(name="openai tx"): - await client.embeddings.create( - input="hello", model="text-embedding-3-large" - ) - (event,) = events - - assert event["contexts"]["trace"]["origin"] == "manual" - assert event["spans"][0]["origin"] == "auto.ai.openai" - - -def test_completions_token_usage_from_response(): - """Token counts are extracted from response.usage using Completions API field names.""" - span = mock.MagicMock() - - def count_tokens(msg): - return len(str(msg)) - - response = mock.MagicMock() - response.usage = mock.MagicMock() - response.usage.completion_tokens = 10 - response.usage.prompt_tokens = 20 - response.usage.total_tokens = 30 - messages = [] - streaming_message_responses = [] - - with mock.patch( - "sentry_sdk.integrations.openai.record_token_usage" - ) as mock_record_token_usage: - _calculate_completions_token_usage( - messages=messages, - response=response, - span=span, - streaming_message_responses=streaming_message_responses, - streaming_message_total_token_usage=None, - count_tokens=count_tokens, - ) - mock_record_token_usage.assert_called_once_with( - span, - input_tokens=20, - input_tokens_cached=None, - output_tokens=10, - output_tokens_reasoning=None, - total_tokens=30, - ) - - -def test_completions_token_usage_with_detailed_fields(): - """Cached and reasoning token counts are extracted from prompt_tokens_details and completion_tokens_details.""" - span = mock.MagicMock() - - def count_tokens(msg): - return len(str(msg)) +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [ + (True, False), + (False, True), + (False, False), + ], +) +async def test_embeddings_create_async_no_pii( + sentry_init, + capture_events, + capture_items, + send_default_pii, + include_prompts, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration(include_prompts=include_prompts)], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) - response = mock.MagicMock() - response.usage = mock.MagicMock() - response.usage.prompt_tokens = 20 - response.usage.prompt_tokens_details = mock.MagicMock() - response.usage.prompt_tokens_details.cached_tokens = 5 - response.usage.completion_tokens = 10 - response.usage.completion_tokens_details = mock.MagicMock() - response.usage.completion_tokens_details.reasoning_tokens = 8 - response.usage.total_tokens = 30 + client = AsyncOpenAI(api_key="z") - with mock.patch( - "sentry_sdk.integrations.openai.record_token_usage" - ) as mock_record_token_usage: - _calculate_completions_token_usage( - messages=[], - response=response, - span=span, - streaming_message_responses=[], - streaming_message_total_token_usage=None, - count_tokens=count_tokens, - ) - mock_record_token_usage.assert_called_once_with( - span, - input_tokens=20, - input_tokens_cached=5, - output_tokens=10, - output_tokens_reasoning=8, + returned_embedding = CreateEmbeddingResponse( + data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], + model="some-model", + object="list", + usage=EmbeddingTokenUsage( + prompt_tokens=20, total_tokens=30, - ) + ), + ) + client.embeddings._post = AsyncMock(return_value=returned_embedding) -def test_completions_token_usage_manual_input_counting(): - """When prompt_tokens is missing, input tokens are counted manually from messages.""" - span = mock.MagicMock() + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") - def count_tokens(msg): - return len(str(msg)) + with start_transaction(name="openai tx"): + response = await client.embeddings.create( + input="hello", model="text-embedding-3-large" + ) - response = mock.MagicMock() - response.usage = mock.MagicMock() - response.usage.completion_tokens = 10 - response.usage.total_tokens = 10 - messages = [ - {"content": "one"}, - {"content": "two"}, - {"content": "three"}, - ] - streaming_message_responses = [] + assert len(response.data[0].embedding) == 3 - with mock.patch( - "sentry_sdk.integrations.openai.record_token_usage" - ) as mock_record_token_usage: - _calculate_completions_token_usage( - messages=messages, - response=response, - span=span, - streaming_message_responses=streaming_message_responses, - streaming_message_total_token_usage=None, - count_tokens=count_tokens, - ) - mock_record_token_usage.assert_called_once_with( - span, - input_tokens=11, - input_tokens_cached=None, - output_tokens=10, - output_tokens_reasoning=None, - total_tokens=10, + sentry_sdk.flush() + span = next(item.payload for item in items) + assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" + assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert ( + span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] + == "text-embedding-3-large" ) + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["attributes"] -def test_completions_token_usage_manual_output_counting_streaming(): - """When completion_tokens is missing, output tokens are counted from streaming responses.""" - span = mock.MagicMock() + assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 + assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 + else: + events = capture_events() - def count_tokens(msg): - return len(str(msg)) + with start_transaction(name="openai tx"): + response = await client.embeddings.create( + input="hello", model="text-embedding-3-large" + ) - response = mock.MagicMock() - response.usage = mock.MagicMock() - response.usage.prompt_tokens = 20 - response.usage.total_tokens = 20 - messages = [] - streaming_message_responses = [ - "one", - "two", - "three", - ] + assert len(response.data[0].embedding) == 3 - with mock.patch( - "sentry_sdk.integrations.openai.record_token_usage" - ) as mock_record_token_usage: - _calculate_completions_token_usage( - messages=messages, - response=response, - span=span, - streaming_message_responses=streaming_message_responses, - streaming_message_total_token_usage=None, - count_tokens=count_tokens, - ) - mock_record_token_usage.assert_called_once_with( - span, - input_tokens=20, - input_tokens_cached=None, - output_tokens=11, - output_tokens_reasoning=None, - total_tokens=20, - ) + tx = events[0] + assert tx["type"] == "transaction" + span = tx["spans"][0] + assert span["op"] == "gen_ai.embeddings" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["data"] -def test_completions_token_usage_manual_output_counting_choices(): - """When completion_tokens is missing, output tokens are counted from response.choices.""" - span = mock.MagicMock() + assert span["data"]["gen_ai.usage.input_tokens"] == 20 + assert span["data"]["gen_ai.usage.total_tokens"] == 30 - def count_tokens(msg): - return len(str(msg)) - response = mock.MagicMock() - response.usage = mock.MagicMock() - response.usage.prompt_tokens = 20 - response.usage.total_tokens = 20 - response.choices = [ - Choice( - index=0, - finish_reason="stop", - message=ChatCompletionMessage(role="assistant", content="one"), - ), - Choice( - index=1, - finish_reason="stop", - message=ChatCompletionMessage(role="assistant", content="two"), +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "get_input,expected_embeddings_input", + [ + ( + lambda: "hello", + ["hello"], ), - Choice( - index=2, - finish_reason="stop", - message=ChatCompletionMessage(role="assistant", content="three"), + ( + lambda: ["First text", "Second text", "Third text"], + [ + "First text", + "Second text", + "Third text", + ], ), - ] - messages = [] - streaming_message_responses = None - - with mock.patch( - "sentry_sdk.integrations.openai.record_token_usage" - ) as mock_record_token_usage: - _calculate_completions_token_usage( - messages=messages, - response=response, - span=span, - streaming_message_responses=streaming_message_responses, - streaming_message_total_token_usage=None, - count_tokens=count_tokens, - ) - mock_record_token_usage.assert_called_once_with( - span, - input_tokens=20, - input_tokens_cached=None, - output_tokens=11, - output_tokens_reasoning=None, - total_tokens=20, - ) + ( + lambda: iter(["First text", "Second text", "Third text"]), + [ + "First text", + "Second text", + "Third text", + ], + ), + ( + lambda: [5, 8, 13, 21, 34], + [ + 5, + 8, + 13, + 21, + 34, + ], + ), + ( + lambda: iter( + [5, 8, 13, 21, 34], + ), + [ + 5, + 8, + 13, + 21, + 34, + ], + ), + ( + lambda: [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ], + [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ], + ), + ( + lambda: iter( + [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ] + ), + [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ], + ), + ], +) +async def test_embeddings_create_async( + sentry_init, + capture_events, + capture_items, + get_input, + expected_embeddings_input, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration(include_prompts=True)], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + send_default_pii=True, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + client = AsyncOpenAI(api_key="z") -def test_completions_token_usage_no_usage_data(): - """When response has no usage data and no streaming responses, all tokens are None.""" - span = mock.MagicMock() + returned_embedding = CreateEmbeddingResponse( + data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], + model="some-model", + object="list", + usage=EmbeddingTokenUsage( + prompt_tokens=20, + total_tokens=30, + ), + ) - def count_tokens(msg): - return len(str(msg)) + client.embeddings._post = AsyncMock(return_value=returned_embedding) - response = mock.MagicMock() - messages = [] - streaming_message_responses = None + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") - with mock.patch( - "sentry_sdk.integrations.openai.record_token_usage" - ) as mock_record_token_usage: - _calculate_completions_token_usage( - messages=messages, - response=response, - span=span, - streaming_message_responses=streaming_message_responses, - streaming_message_total_token_usage=None, - count_tokens=count_tokens, + with start_transaction(name="openai tx"): + response = await client.embeddings.create( + input=get_input(), model="text-embedding-3-large" + ) + + assert len(response.data[0].embedding) == 3 + + sentry_sdk.flush() + span = next(item.payload for item in items) + assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" + assert span["attributes"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert ( + span["attributes"][SPANDATA.GEN_AI_REQUEST_MODEL] + == "text-embedding-3-large" ) - mock_record_token_usage.assert_called_once_with( - span, - input_tokens=None, - input_tokens_cached=None, - output_tokens=None, - output_tokens_reasoning=None, - total_tokens=None, + + assert ( + json.loads(span["attributes"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) + == expected_embeddings_input ) + assert span["attributes"]["gen_ai.usage.input_tokens"] == 20 + assert span["attributes"]["gen_ai.usage.total_tokens"] == 30 + else: + events = capture_events() -@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") -def test_responses_token_usage_from_response(): - """Token counts including cached and reasoning tokens are extracted from Responses API.""" - span = mock.MagicMock() + with start_transaction(name="openai tx"): + response = await client.embeddings.create( + input=get_input(), model="text-embedding-3-large" + ) - def count_tokens(msg): - return len(str(msg)) + assert len(response.data[0].embedding) == 3 - response = mock.MagicMock() - response.usage = mock.MagicMock() - response.usage.input_tokens = 20 - response.usage.input_tokens_details = mock.MagicMock() - response.usage.input_tokens_details.cached_tokens = 5 - response.usage.output_tokens = 10 - response.usage.output_tokens_details = mock.MagicMock() - response.usage.output_tokens_details.reasoning_tokens = 8 - response.usage.total_tokens = 30 - input = [] + tx = events[0] + assert tx["type"] == "transaction" + span = tx["spans"][0] + assert span["op"] == "gen_ai.embeddings" + assert span["data"][SPANDATA.GEN_AI_SYSTEM] == "openai" + assert span["data"][SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" - with mock.patch( - "sentry_sdk.integrations.openai.record_token_usage" - ) as mock_record_token_usage: - _calculate_responses_token_usage(input, response, span, None, count_tokens) - mock_record_token_usage.assert_called_once_with( - span, - input_tokens=20, - input_tokens_cached=5, - output_tokens=10, - output_tokens_reasoning=8, - total_tokens=30, + assert ( + json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) + == expected_embeddings_input ) + assert span["data"]["gen_ai.usage.input_tokens"] == 20 + assert span["data"]["gen_ai.usage.total_tokens"] == 30 -@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") -def test_responses_token_usage_no_usage_data(): - """When Responses API response has no usage data, all tokens are None.""" - span = mock.MagicMock() - - def count_tokens(msg): - return len(str(msg)) - - response = mock.MagicMock() - response.usage = None - input = [] - streaming_message_responses = None - with mock.patch( - "sentry_sdk.integrations.openai.record_token_usage" - ) as mock_record_token_usage: - _calculate_responses_token_usage( - input, response, span, streaming_message_responses, count_tokens - ) - mock_record_token_usage.assert_called_once_with( - span, - input_tokens=None, - input_tokens_cached=None, - output_tokens=None, - output_tokens_reasoning=None, - total_tokens=None, - ) +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "data_collection,send_default_pii,include_prompts,expect_input", + [ + pytest.param( + {"gen_ai": {"inputs": True}}, + False, + True, + True, + id="inputs-enabled-overrides-pii-disabled", + ), + pytest.param( + {"gen_ai": {"inputs": False}}, + True, + True, + False, + id="inputs-disabled-overrides-pii-enabled", + ), + pytest.param( + {}, + False, + True, + True, + id="gen-ai-omitted-defaults-to-enabled", + ), + pytest.param( + {"gen_ai": {"inputs": True}}, + True, + False, + False, + id="include-prompts-disabled-overrides-inputs-enabled", + ), + pytest.param( + None, + False, + True, + False, + id="no-experiment-falls-back-to-pii", + ), + ], +) +async def test_embeddings_create_async_data_collection( + sentry_init, + capture_events, + capture_items, + data_collection, + send_default_pii, + include_prompts, + expect_input, + stream_gen_ai_spans, + span_streaming, +): + init_kwargs = { + "integrations": [OpenAIIntegration(include_prompts=include_prompts)], + "disabled_integrations": [StdlibIntegration], + "traces_sample_rate": 1.0, + "send_default_pii": send_default_pii, + "stream_gen_ai_spans": stream_gen_ai_spans, + "trace_lifecycle": "stream" if span_streaming else "static", + } + + sentry_init_kwargs = dict(init_kwargs) + if data_collection is not None: + sentry_init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**sentry_init_kwargs) + + client = AsyncOpenAI(api_key="z") + + returned_embedding = CreateEmbeddingResponse( + data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], + model="some-model", + object="list", + usage=EmbeddingTokenUsage( + prompt_tokens=20, + total_tokens=30, + ), + ) + + client.embeddings._post = AsyncMock(return_value=returned_embedding) + + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + + with start_transaction(name="openai tx"): + response = await client.embeddings.create( + input="hello", model="text-embedding-3-large" + ) + + assert len(response.data[0].embedding) == 3 + + sentry_sdk.flush() + span = next(item.payload for item in items) + assert span["attributes"]["sentry.op"] == "gen_ai.embeddings" + span_data = span["attributes"] + else: + events = capture_events() + + with start_transaction(name="openai tx"): + response = await client.embeddings.create( + input="hello", model="text-embedding-3-large" + ) + + assert len(response.data[0].embedding) == 3 + + tx = events[0] + assert tx["type"] == "transaction" + span = tx["spans"][0] + assert span["op"] == "gen_ai.embeddings" + span_data = span["data"] + + assert span_data[SPANDATA.GEN_AI_SYSTEM] == "openai" + assert span_data[SPANDATA.GEN_AI_OPERATION_NAME] == "embeddings" + assert span_data[SPANDATA.GEN_AI_REQUEST_MODEL] == "text-embedding-3-large" + + if expect_input: + assert json.loads(span_data[SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == ["hello"] + else: + assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span_data + + assert span_data["gen_ai.usage.input_tokens"] == 20 + assert span_data["gen_ai.usage.total_tokens"] == 30 + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (True, False), (False, True), (False, False)], +) +def test_embeddings_create_raises_error( + sentry_init, + capture_events, + capture_items, + send_default_pii, + include_prompts, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration(include_prompts=include_prompts)], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = OpenAI(api_key="z") + + client.embeddings._post = mock.Mock( + side_effect=OpenAIError("API rate limit reached") + ) + + if span_streaming: + items = capture_items("event", "span") + + with pytest.raises(OpenAIError): + client.embeddings.create(input="hello", model="text-embedding-3-large") + + (event,) = (item.payload for item in items if item.type == "event") + sentry_sdk.flush() + (span,) = (item.payload for item in items if item.type == "span") + assert event["level"] == "error" + assert span["status"] == "error" + elif stream_gen_ai_spans: + items = capture_items("event", "transaction") + + with pytest.raises(OpenAIError): + client.embeddings.create(input="hello", model="text-embedding-3-large") + + (event,) = (item.payload for item in items if item.type == "event") + (transaction,) = (item.payload for item in items if item.type == "transaction") + assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" + else: + events = capture_events() + + with pytest.raises(OpenAIError): + client.embeddings.create(input="hello", model="text-embedding-3-large") + + (event, transaction) = events + assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +@pytest.mark.parametrize( + "send_default_pii, include_prompts", + [(True, True), (True, False), (False, True), (False, False)], +) +async def test_embeddings_create_raises_error_async( + sentry_init, + capture_events, + capture_items, + send_default_pii, + include_prompts, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration(include_prompts=include_prompts)], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + send_default_pii=send_default_pii, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = AsyncOpenAI(api_key="z") + + client.embeddings._post = AsyncMock( + side_effect=OpenAIError("API rate limit reached") + ) + + if span_streaming: + items = capture_items("event", "span") + + with pytest.raises(OpenAIError): + await client.embeddings.create( + input="hello", model="text-embedding-3-large" + ) + + (event,) = (item.payload for item in items if item.type == "event") + sentry_sdk.flush() + (span,) = (item.payload for item in items if item.type == "span") + assert event["level"] == "error" + assert span["status"] == "error" + elif stream_gen_ai_spans: + items = capture_items("event", "transaction") + + with pytest.raises(OpenAIError): + await client.embeddings.create( + input="hello", model="text-embedding-3-large" + ) + + (event,) = (item.payload for item in items if item.type == "event") + (transaction,) = (item.payload for item in items if item.type == "transaction") + assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" + else: + events = capture_events() + + with pytest.raises(OpenAIError): + await client.embeddings.create( + input="hello", model="text-embedding-3-large" + ) + + (event, transaction) = events + assert event["level"] == "error" + assert transaction["contexts"]["trace"]["status"] == "internal_error" + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +def test_span_origin_nonstreaming_chat( + sentry_init, + capture_events, + capture_items, + nonstreaming_chat_completions_model_response, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration()], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = OpenAI(api_key="z") + client.chat.completions._post = mock.Mock( + return_value=nonstreaming_chat_completions_model_response( + response_id="chat-id", + response_model="gpt-3.5-turbo", + message_content="the model response", + created=10000000, + usage=CompletionUsage( + prompt_tokens=20, + completion_tokens=10, + total_tokens=30, + ), + ) + ) + + if span_streaming: + items = capture_items("transaction", "span") + + with sentry_sdk.traces.start_span(name="openai tx"): + client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[1]["attributes"]["sentry.origin"] == "manual" + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + elif stream_gen_ai_spans: + items = capture_items("transaction", "span") + + with start_transaction(name="openai tx"): + client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + + (event,) = (item.payload for item in items if item.type == "transaction") + assert event["contexts"]["trace"]["origin"] == "manual" + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + else: + events = capture_events() + + with start_transaction(name="openai tx"): + client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + + (event,) = events + + assert event["contexts"]["trace"]["origin"] == "manual" + assert event["spans"][0]["origin"] == "auto.ai.openai" + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +async def test_span_origin_nonstreaming_chat_async( + sentry_init, + capture_events, + capture_items, + nonstreaming_chat_completions_model_response, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = AsyncOpenAI(api_key="z") + client.chat.completions._post = AsyncMock( + return_value=nonstreaming_chat_completions_model_response( + response_id="chat-id", + response_model="gpt-3.5-turbo", + message_content="the model response", + created=10000000, + usage=CompletionUsage( + prompt_tokens=20, + completion_tokens=10, + total_tokens=30, + ), + ) + ) + + if span_streaming: + items = capture_items("transaction", "span") + + with sentry_sdk.traces.start_span(name="openai tx"): + await client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[1]["attributes"]["sentry.origin"] == "manual" + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + elif stream_gen_ai_spans: + items = capture_items("transaction", "span") + + with start_transaction(name="openai tx"): + await client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + + (event,) = (item.payload for item in items if item.type == "transaction") + assert event["contexts"]["trace"]["origin"] == "manual" + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + else: + events = capture_events() + + with start_transaction(name="openai tx"): + await client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + + (event,) = events + + assert event["contexts"]["trace"]["origin"] == "manual" + assert event["spans"][0]["origin"] == "auto.ai.openai" + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +def test_span_origin_streaming_chat( + sentry_init, + capture_events, + capture_items, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = OpenAI(api_key="z") + returned_stream = Stream(cast_to=None, response=None, client=client) + returned_stream._iterator = [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, delta=ChoiceDelta(content="hel"), finish_reason=None + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=1, delta=ChoiceDelta(content="lo "), finish_reason=None + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=2, delta=ChoiceDelta(content="world"), finish_reason="stop" + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ] + + if span_streaming: + items = capture_items("transaction", "span") + + client.chat.completions._post = mock.Mock(return_value=returned_stream) + with sentry_sdk.traces.start_span(name="openai tx"): + response_stream = client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + + "".join(map(lambda x: x.choices[0].delta.content, response_stream)) + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[1]["attributes"]["sentry.origin"] == "manual" + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + elif stream_gen_ai_spans: + items = capture_items("transaction", "span") + + client.chat.completions._post = mock.Mock(return_value=returned_stream) + with start_transaction(name="openai tx"): + response_stream = client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + + "".join(map(lambda x: x.choices[0].delta.content, response_stream)) + + (event,) = (item.payload for item in items if item.type == "transaction") + assert event["contexts"]["trace"]["origin"] == "manual" + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + else: + events = capture_events() + + client.chat.completions._post = mock.Mock(return_value=returned_stream) + with start_transaction(name="openai tx"): + response_stream = client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + + "".join(map(lambda x: x.choices[0].delta.content, response_stream)) + + (event,) = events + + assert event["contexts"]["trace"]["origin"] == "manual" + assert event["spans"][0]["origin"] == "auto.ai.openai" + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +async def test_span_origin_streaming_chat_async( + sentry_init, + capture_events, + capture_items, + async_iterator, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = AsyncOpenAI(api_key="z") + returned_stream = AsyncStream(cast_to=None, response=None, client=client) + returned_stream._iterator = async_iterator( + [ + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=0, delta=ChoiceDelta(content="hel"), finish_reason=None + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=1, delta=ChoiceDelta(content="lo "), finish_reason=None + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ChatCompletionChunk( + id="1", + choices=[ + DeltaChoice( + index=2, + delta=ChoiceDelta(content="world"), + finish_reason="stop", + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ] + ) + + client.chat.completions._post = AsyncMock(return_value=returned_stream) + + if span_streaming: + items = capture_items("transaction", "span") + + with sentry_sdk.traces.start_span(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + async for _ in response_stream: + pass + + # "".join(map(lambda x: x.choices[0].delta.content, response_stream)) + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[1]["attributes"]["sentry.origin"] == "manual" + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + elif stream_gen_ai_spans: + items = capture_items("transaction", "span") + + with start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + async for _ in response_stream: + pass + + # "".join(map(lambda x: x.choices[0].delta.content, response_stream)) + + (event,) = (item.payload for item in items if item.type == "transaction") + assert event["contexts"]["trace"]["origin"] == "manual" + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + else: + events = capture_events() + + with start_transaction(name="openai tx"): + response_stream = await client.chat.completions.create( + model="some-model", messages=[{"role": "system", "content": "hello"}] + ) + async for _ in response_stream: + pass + + # "".join(map(lambda x: x.choices[0].delta.content, response_stream)) + + (event,) = events + + assert event["contexts"]["trace"]["origin"] == "manual" + assert event["spans"][0]["origin"] == "auto.ai.openai" + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +def test_span_origin_embeddings( + sentry_init, + capture_events, + capture_items, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = OpenAI(api_key="z") + + returned_embedding = CreateEmbeddingResponse( + data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], + model="some-model", + object="list", + usage=EmbeddingTokenUsage( + prompt_tokens=20, + total_tokens=30, + ), + ) + + client.embeddings._post = mock.Mock(return_value=returned_embedding) + + if span_streaming: + items = capture_items("transaction", "span") + + with sentry_sdk.traces.start_span(name="openai tx"): + client.embeddings.create(input="hello", model="text-embedding-3-large") + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[1]["attributes"]["sentry.origin"] == "manual" + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + elif stream_gen_ai_spans: + items = capture_items("transaction", "span") + + with start_transaction(name="openai tx"): + client.embeddings.create(input="hello", model="text-embedding-3-large") + + (event,) = [item.payload for item in items if item.type == "transaction"] + assert event["contexts"]["trace"]["origin"] == "manual" + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + else: + events = capture_events() + + with start_transaction(name="openai tx"): + client.embeddings.create(input="hello", model="text-embedding-3-large") + + (event,) = events + + assert event["contexts"]["trace"]["origin"] == "manual" + assert event["spans"][0]["origin"] == "auto.ai.openai" + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.asyncio +async def test_span_origin_embeddings_async( + sentry_init, + capture_events, + capture_items, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = AsyncOpenAI(api_key="z") + + returned_embedding = CreateEmbeddingResponse( + data=[Embedding(object="embedding", index=0, embedding=[1.0, 2.0, 3.0])], + model="some-model", + object="list", + usage=EmbeddingTokenUsage( + prompt_tokens=20, + total_tokens=30, + ), + ) + + client.embeddings._post = AsyncMock(return_value=returned_embedding) + + if span_streaming: + items = capture_items("transaction", "span") + + with sentry_sdk.traces.start_span(name="openai tx"): + await client.embeddings.create( + input="hello", model="text-embedding-3-large" + ) + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[1]["attributes"]["sentry.origin"] == "manual" + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + elif stream_gen_ai_spans: + items = capture_items("transaction", "span") + + with start_transaction(name="openai tx"): + await client.embeddings.create( + input="hello", model="text-embedding-3-large" + ) + + (event,) = [item.payload for item in items if item.type == "transaction"] + assert event["contexts"]["trace"]["origin"] == "manual" + + sentry_sdk.flush() + spans = [item.payload for item in items if item.type == "span"] + assert spans[0]["attributes"]["sentry.origin"] == "auto.ai.openai" + else: + events = capture_events() + + with start_transaction(name="openai tx"): + await client.embeddings.create( + input="hello", model="text-embedding-3-large" + ) + + (event,) = events + + assert event["contexts"]["trace"]["origin"] == "manual" + assert event["spans"][0]["origin"] == "auto.ai.openai" + + +def test_completions_token_usage_from_response(): + """Token counts are extracted from response.usage using Completions API field names.""" + span = mock.MagicMock() + + def count_tokens(msg): + return len(str(msg)) + + response = mock.MagicMock() + response.usage = mock.MagicMock() + response.usage.completion_tokens = 10 + response.usage.prompt_tokens = 20 + response.usage.total_tokens = 30 + messages = [] + streaming_message_responses = [] + + with mock.patch( + "sentry_sdk.integrations.openai.record_token_usage" + ) as mock_record_token_usage: + _calculate_completions_token_usage( + messages=messages, + response=response, + span=span, + streaming_message_responses=streaming_message_responses, + streaming_message_total_token_usage=None, + count_tokens=count_tokens, + ) + mock_record_token_usage.assert_called_once_with( + span, + input_tokens=20, + input_tokens_cached=None, + output_tokens=10, + output_tokens_reasoning=None, + total_tokens=30, + ) + + +def test_completions_token_usage_with_detailed_fields(): + """Cached and reasoning token counts are extracted from prompt_tokens_details and completion_tokens_details.""" + span = mock.MagicMock() + + def count_tokens(msg): + return len(str(msg)) + + response = mock.MagicMock() + response.usage = mock.MagicMock() + response.usage.prompt_tokens = 20 + response.usage.prompt_tokens_details = mock.MagicMock() + response.usage.prompt_tokens_details.cached_tokens = 5 + response.usage.completion_tokens = 10 + response.usage.completion_tokens_details = mock.MagicMock() + response.usage.completion_tokens_details.reasoning_tokens = 8 + response.usage.total_tokens = 30 + + with mock.patch( + "sentry_sdk.integrations.openai.record_token_usage" + ) as mock_record_token_usage: + _calculate_completions_token_usage( + messages=[], + response=response, + span=span, + streaming_message_responses=[], + streaming_message_total_token_usage=None, + count_tokens=count_tokens, + ) + mock_record_token_usage.assert_called_once_with( + span, + input_tokens=20, + input_tokens_cached=5, + output_tokens=10, + output_tokens_reasoning=8, + total_tokens=30, + ) + + +def test_completions_token_usage_manual_input_counting(): + """When prompt_tokens is missing, input tokens are counted manually from messages.""" + span = mock.MagicMock() + + def count_tokens(msg): + return len(str(msg)) + + response = mock.MagicMock() + response.usage = mock.MagicMock() + response.usage.completion_tokens = 10 + response.usage.total_tokens = 10 + messages = [ + {"content": "one"}, + {"content": "two"}, + {"content": "three"}, + ] + streaming_message_responses = [] + + with mock.patch( + "sentry_sdk.integrations.openai.record_token_usage" + ) as mock_record_token_usage: + _calculate_completions_token_usage( + messages=messages, + response=response, + span=span, + streaming_message_responses=streaming_message_responses, + streaming_message_total_token_usage=None, + count_tokens=count_tokens, + ) + mock_record_token_usage.assert_called_once_with( + span, + input_tokens=11, + input_tokens_cached=None, + output_tokens=10, + output_tokens_reasoning=None, + total_tokens=10, + ) + + +def test_completions_token_usage_manual_output_counting_streaming(): + """When completion_tokens is missing, output tokens are counted from streaming responses.""" + span = mock.MagicMock() + + def count_tokens(msg): + return len(str(msg)) + + response = mock.MagicMock() + response.usage = mock.MagicMock() + response.usage.prompt_tokens = 20 + response.usage.total_tokens = 20 + messages = [] + streaming_message_responses = [ + "one", + "two", + "three", + ] + + with mock.patch( + "sentry_sdk.integrations.openai.record_token_usage" + ) as mock_record_token_usage: + _calculate_completions_token_usage( + messages=messages, + response=response, + span=span, + streaming_message_responses=streaming_message_responses, + streaming_message_total_token_usage=None, + count_tokens=count_tokens, + ) + mock_record_token_usage.assert_called_once_with( + span, + input_tokens=20, + input_tokens_cached=None, + output_tokens=11, + output_tokens_reasoning=None, + total_tokens=20, + ) + + +def test_completions_token_usage_manual_output_counting_choices(): + """When completion_tokens is missing, output tokens are counted from response.choices.""" + span = mock.MagicMock() + + def count_tokens(msg): + return len(str(msg)) + + response = mock.MagicMock() + response.usage = mock.MagicMock() + response.usage.prompt_tokens = 20 + response.usage.total_tokens = 20 + response.choices = [ + Choice( + index=0, + finish_reason="stop", + message=ChatCompletionMessage(role="assistant", content="one"), + ), + Choice( + index=1, + finish_reason="stop", + message=ChatCompletionMessage(role="assistant", content="two"), + ), + Choice( + index=2, + finish_reason="stop", + message=ChatCompletionMessage(role="assistant", content="three"), + ), + ] + messages = [] + streaming_message_responses = None + + with mock.patch( + "sentry_sdk.integrations.openai.record_token_usage" + ) as mock_record_token_usage: + _calculate_completions_token_usage( + messages=messages, + response=response, + span=span, + streaming_message_responses=streaming_message_responses, + streaming_message_total_token_usage=None, + count_tokens=count_tokens, + ) + mock_record_token_usage.assert_called_once_with( + span, + input_tokens=20, + input_tokens_cached=None, + output_tokens=11, + output_tokens_reasoning=None, + total_tokens=20, + ) + + +def test_completions_token_usage_no_usage_data(): + """When response has no usage data and no streaming responses, all tokens are None.""" + span = mock.MagicMock() + + def count_tokens(msg): + return len(str(msg)) + + response = mock.MagicMock() + messages = [] + streaming_message_responses = None + + with mock.patch( + "sentry_sdk.integrations.openai.record_token_usage" + ) as mock_record_token_usage: + _calculate_completions_token_usage( + messages=messages, + response=response, + span=span, + streaming_message_responses=streaming_message_responses, + streaming_message_total_token_usage=None, + count_tokens=count_tokens, + ) + mock_record_token_usage.assert_called_once_with( + span, + input_tokens=None, + input_tokens_cached=None, + output_tokens=None, + output_tokens_reasoning=None, + total_tokens=None, + ) + + +@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") +def test_responses_token_usage_from_response(): + """Token counts including cached and reasoning tokens are extracted from Responses API.""" + span = mock.MagicMock() + + def count_tokens(msg): + return len(str(msg)) + + response = mock.MagicMock() + response.usage = mock.MagicMock() + response.usage.input_tokens = 20 + response.usage.input_tokens_details = mock.MagicMock() + response.usage.input_tokens_details.cached_tokens = 5 + response.usage.output_tokens = 10 + response.usage.output_tokens_details = mock.MagicMock() + response.usage.output_tokens_details.reasoning_tokens = 8 + response.usage.total_tokens = 30 + input = [] + + with mock.patch( + "sentry_sdk.integrations.openai.record_token_usage" + ) as mock_record_token_usage: + _calculate_responses_token_usage(input, response, span, None, count_tokens) + mock_record_token_usage.assert_called_once_with( + span, + input_tokens=20, + input_tokens_cached=5, + output_tokens=10, + output_tokens_reasoning=8, + total_tokens=30, + ) + + +@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") +def test_responses_token_usage_no_usage_data(): + """When Responses API response has no usage data, all tokens are None.""" + span = mock.MagicMock() + + def count_tokens(msg): + return len(str(msg)) + + response = mock.MagicMock() + response.usage = None + input = [] + streaming_message_responses = None + + with mock.patch( + "sentry_sdk.integrations.openai.record_token_usage" + ) as mock_record_token_usage: + _calculate_responses_token_usage( + input, response, span, streaming_message_responses, count_tokens + ) + mock_record_token_usage.assert_called_once_with( + span, + input_tokens=None, + input_tokens_cached=None, + output_tokens=None, + output_tokens_reasoning=None, + total_tokens=None, + ) + + +@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") +def test_responses_token_usage_manual_output_counting_response_output(): + """When output_tokens is missing, output tokens are counted from response.output.""" + span = mock.MagicMock() + + def count_tokens(msg): + return len(str(msg)) + + response = mock.MagicMock() + response.usage = mock.MagicMock() + response.usage.input_tokens = 20 + response.usage.total_tokens = 20 + response.output = [ + ResponseOutputMessage( + id="msg-1", + content=[ + ResponseOutputText( + annotations=[], + text="one", + type="output_text", + ), + ], + role="assistant", + status="completed", + type="message", + ), + ResponseOutputMessage( + id="msg-2", + content=[ + ResponseOutputText( + annotations=[], + text="two", + type="output_text", + ), + ResponseOutputText( + annotations=[], + text="three", + type="output_text", + ), + ], + role="assistant", + status="completed", + type="message", + ), + ] + input = [] + streaming_message_responses = None + + with mock.patch( + "sentry_sdk.integrations.openai.record_token_usage" + ) as mock_record_token_usage: + _calculate_responses_token_usage( + input, response, span, streaming_message_responses, count_tokens + ) + mock_record_token_usage.assert_called_once_with( + span, + input_tokens=20, + input_tokens_cached=None, + output_tokens=11, + output_tokens_reasoning=None, + total_tokens=20, + ) + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") +def test_ai_client_span_responses_api_no_pii( + sentry_init, + capture_events, + capture_items, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = OpenAI(api_key="z") + client.responses._post = mock.Mock(return_value=EXAMPLE_RESPONSE) + + if span_streaming: + items = capture_items("span") + + with sentry_sdk.traces.start_span(name="openai tx"): + client.responses.create( + model="gpt-4o", + instructions="You are a coding assistant that talks like a pirate.", + input="How do I check if a Python object is an instance of a class?", + max_output_tokens=100, + temperature=0.7, + top_p=0.9, + reasoning={"effort": "high"}, + ) + + sentry_sdk.flush() + spans = [item.payload for item in items] + + assert len(spans) == 2 + expected_attributes = { + "gen_ai.operation.name": "responses", + "gen_ai.request.max_tokens": 100, + "gen_ai.request.temperature": 0.7, + "gen_ai.request.top_p": 0.9, + "gen_ai.request.reasoning.level": "high", + "gen_ai.request.model": "gpt-4o", + "gen_ai.response.model": "response-model-id", + "gen_ai.response.streaming": False, + "gen_ai.system": "openai", + "gen_ai.usage.input_tokens": 20, + "gen_ai.usage.input_tokens.cached": 5, + "gen_ai.usage.output_tokens": 10, + "gen_ai.usage.output_tokens.reasoning": 8, + "gen_ai.usage.total_tokens": 30, + "sentry.op": "gen_ai.responses", + "sentry.origin": "auto.ai.openai", + "sentry.segment.name": "openai tx", + } + + for attr, value in expected_attributes.items(): + assert spans[0]["attributes"][attr] == value + + assert "gen_ai.system_instructions" not in spans[0]["attributes"] + assert "gen_ai.request.messages" not in spans[0]["attributes"] + assert "gen_ai.response.text" not in spans[0]["attributes"] + + elif stream_gen_ai_spans: + items = capture_items("span") + + with start_transaction(name="openai tx"): + client.responses.create( + model="gpt-4o", + instructions="You are a coding assistant that talks like a pirate.", + input="How do I check if a Python object is an instance of a class?", + max_output_tokens=100, + temperature=0.7, + top_p=0.9, + reasoning={"effort": "high"}, + tools=EXAMPLE_TOOLS, + ) + + spans = [item.payload for item in items] + + assert len(spans) == 1 + expected_attributes = { + "gen_ai.operation.name": "responses", + "gen_ai.request.max_tokens": 100, + "gen_ai.request.temperature": 0.7, + "gen_ai.request.top_p": 0.9, + "gen_ai.request.reasoning.level": "high", + "gen_ai.request.model": "gpt-4o", + "gen_ai.response.model": "response-model-id", + "gen_ai.response.streaming": False, + "gen_ai.system": "openai", + "gen_ai.usage.input_tokens": 20, + "gen_ai.usage.input_tokens.cached": 5, + "gen_ai.usage.output_tokens": 10, + "gen_ai.usage.output_tokens.reasoning": 8, + "gen_ai.usage.total_tokens": 30, + "sentry.op": "gen_ai.responses", + "sentry.origin": "auto.ai.openai", + "sentry.segment.name": "openai tx", + } + + for attr, value in expected_attributes.items(): + assert spans[0]["attributes"][attr] == value + + assert "gen_ai.system_instructions" not in spans[0]["attributes"] + assert "gen_ai.request.messages" not in spans[0]["attributes"] + assert "gen_ai.response.text" not in spans[0]["attributes"] + else: + events = capture_events() + + with start_transaction(name="openai tx"): + client.responses.create( + model="gpt-4o", + instructions="You are a coding assistant that talks like a pirate.", + input="How do I check if a Python object is an instance of a class?", + max_output_tokens=100, + temperature=0.7, + top_p=0.9, + reasoning={"effort": "high"}, + ) + + (transaction,) = events + spans = transaction["spans"] + + assert len(spans) == 1 + assert spans[0]["op"] == "gen_ai.responses" + assert spans[0]["origin"] == "auto.ai.openai" + expected_data = { + "gen_ai.operation.name": "responses", + "gen_ai.request.max_tokens": 100, + "gen_ai.request.temperature": 0.7, + "gen_ai.request.top_p": 0.9, + "gen_ai.request.reasoning.level": "high", + "gen_ai.request.model": "gpt-4o", + "gen_ai.response.model": "response-model-id", + "gen_ai.response.streaming": False, + "gen_ai.system": "openai", + "gen_ai.usage.input_tokens": 20, + "gen_ai.usage.input_tokens.cached": 5, + "gen_ai.usage.output_tokens": 10, + "gen_ai.usage.output_tokens.reasoning": 8, + "gen_ai.usage.total_tokens": 30, + } + + for key, value in expected_data.items(): + assert spans[0]["data"][key] == value + + assert "gen_ai.system_instructions" not in spans[0]["data"] + assert "gen_ai.request.messages" not in spans[0]["data"] + assert "gen_ai.response.text" not in spans[0]["data"] + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") +def test_ai_client_span_responses_tool_definitions( + sentry_init, + capture_events, + capture_items, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = OpenAI(api_key="z") + client.responses._post = mock.Mock(return_value=EXAMPLE_RESPONSE) + + if span_streaming: + items = capture_items("span") + + with sentry_sdk.traces.start_span(name="openai tx"): + client.responses.create( + model="gpt-4o", + input="How do I check if a Python object is an instance of a class?", + tools=[ + FunctionToolParam( + type="function", + name="name", + description="description", + parameters={ + "type": "object", + "properties": { + "city": {"type": "string"}, + "state": {"type": "string"}, + }, + "required": ["city", "state"], + "additionalProperties": False, + }, + strict=True, + ), + CustomToolParam( + type="custom", name="name", description="description" + ), + WebSearchToolParam(type="web_search"), + ], + ) + + sentry_sdk.flush() + spans = [item.payload for item in items] + assert json.loads(spans[0]["attributes"][SPANDATA.GEN_AI_TOOL_DEFINITIONS]) == [ + { + "type": "function", + "name": "name", + "description": "description", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "state": {"type": "string"}, + }, + "required": ["city", "state"], + "additionalProperties": False, + }, + }, + { + "type": "custom", + "name": "name", + "description": "description", + }, + { + "type": "web_search", + }, + ] + elif stream_gen_ai_spans: + items = capture_items("span") + + with start_transaction(name="openai tx"): + client.responses.create( + model="gpt-4o", + input="How do I check if a Python object is an instance of a class?", + tools=[ + FunctionToolParam( + type="function", + name="name", + description="description", + parameters={ + "type": "object", + "properties": { + "city": {"type": "string"}, + "state": {"type": "string"}, + }, + "required": ["city", "state"], + "additionalProperties": False, + }, + strict=True, + ), + CustomToolParam( + type="custom", name="name", description="description" + ), + WebSearchToolParam(type="web_search"), + ], + ) + + spans = [item.payload for item in items] + assert json.loads(spans[0]["attributes"][SPANDATA.GEN_AI_TOOL_DEFINITIONS]) == [ + { + "type": "function", + "name": "name", + "description": "description", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "state": {"type": "string"}, + }, + "required": ["city", "state"], + "additionalProperties": False, + }, + }, + { + "type": "custom", + "name": "name", + "description": "description", + }, + { + "type": "web_search", + }, + ] + else: + events = capture_events() + + with start_transaction(name="openai tx"): + client.responses.create( + model="gpt-4o", + input="How do I check if a Python object is an instance of a class?", + tools=[ + FunctionToolParam( + type="function", + name="name", + description="description", + parameters={ + "type": "object", + "properties": { + "city": {"type": "string"}, + "state": {"type": "string"}, + }, + "required": ["city", "state"], + "additionalProperties": False, + }, + strict=True, + ), + CustomToolParam( + type="custom", name="name", description="description" + ), + WebSearchToolParam(type="web_search"), + ], + ) + (transaction,) = events + spans = transaction["spans"] -@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") -def test_responses_token_usage_manual_output_counting_response_output(): - """When output_tokens is missing, output tokens are counted from response.output.""" - span = mock.MagicMock() + assert json.loads(spans[0]["data"][SPANDATA.GEN_AI_TOOL_DEFINITIONS]) == [ + { + "type": "function", + "name": "name", + "description": "description", + "parameters": { + "type": "object", + "properties": { + "city": {"type": "string"}, + "state": {"type": "string"}, + }, + "required": ["city", "state"], + "additionalProperties": False, + }, + }, + { + "type": "custom", + "name": "name", + "description": "description", + }, + { + "type": "web_search", + }, + ] - def count_tokens(msg): - return len(str(msg)) - response = mock.MagicMock() - response.usage = mock.MagicMock() - response.usage.input_tokens = 20 - response.usage.total_tokens = 20 - response.output = [ - ResponseOutputMessage( - id="msg-1", - content=[ - ResponseOutputText( - annotations=[], - text="one", - type="output_text", - ), +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.parametrize( + "instructions,input,expected_system_instructions,expected_request_messages", + [ + ( + omit, + "How do I check if a Python object is an instance of a class?", + None, + ["How do I check if a Python object is an instance of a class?"], + ), + ( + None, + "How do I check if a Python object is an instance of a class?", + None, + ["How do I check if a Python object is an instance of a class?"], + ), + ( + "You are a coding assistant that talks like a pirate.", + [ + { + "role": "system", + "content": "You are a helpful assistant.", + }, + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "hello"}, + ], + [ + { + "type": "text", + "content": "You are a coding assistant that talks like a pirate.", + }, + {"type": "text", "content": "You are a helpful assistant."}, + ], + [ + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "hello"}, ], - role="assistant", - status="completed", - type="message", ), - ResponseOutputMessage( - id="msg-2", - content=[ - ResponseOutputText( - annotations=[], - text="two", - type="output_text", - ), - ResponseOutputText( - annotations=[], - text="three", - type="output_text", - ), + ( + "You are a coding assistant that talks like a pirate.", + [ + { + "type": "message", + "role": "system", + "content": "You are a helpful assistant.", + }, + { + "type": "message", + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"type": "message", "role": "user", "content": "hello"}, + ], + [ + { + "type": "text", + "content": "You are a coding assistant that talks like a pirate.", + }, + {"type": "text", "content": "You are a helpful assistant."}, + ], + [ + { + "type": "message", + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"type": "message", "role": "user", "content": "hello"}, ], - role="assistant", - status="completed", - type="message", ), - ] - input = [] - streaming_message_responses = None - - with mock.patch( - "sentry_sdk.integrations.openai.record_token_usage" - ) as mock_record_token_usage: - _calculate_responses_token_usage( - input, response, span, streaming_message_responses, count_tokens - ) - mock_record_token_usage.assert_called_once_with( - span, - input_tokens=20, - input_tokens_cached=None, - output_tokens=11, - output_tokens_reasoning=None, - total_tokens=20, - ) - - -@pytest.mark.parametrize("span_streaming", [True, False]) -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) + ( + "You are a coding assistant that talks like a pirate.", + [ + { + "role": "system", + "content": [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Be concise and clear."}, + ], + }, + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "hello"}, + ], + [ + { + "type": "text", + "content": "You are a coding assistant that talks like a pirate.", + }, + {"type": "text", "content": "You are a helpful assistant."}, + {"type": "text", "content": "Be concise and clear."}, + ], + [ + { + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"role": "user", "content": "hello"}, + ], + ), + ( + "You are a coding assistant that talks like a pirate.", + [ + { + "type": "message", + "role": "system", + "content": [ + {"type": "text", "text": "You are a helpful assistant."}, + {"type": "text", "text": "Be concise and clear."}, + ], + }, + { + "type": "message", + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"type": "message", "role": "user", "content": "hello"}, + ], + [ + { + "type": "text", + "content": "You are a coding assistant that talks like a pirate.", + }, + {"type": "text", "content": "You are a helpful assistant."}, + {"type": "text", "content": "Be concise and clear."}, + ], + [ + { + "type": "message", + "role": "user", + "content": "Message demonstrating the absence of truncation.", + }, + {"type": "message", "role": "user", "content": "hello"}, + ], + ), + ], +) @pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") -def test_ai_client_span_responses_api_no_pii( +def test_ai_client_span_responses_api( sentry_init, capture_events, capture_items, + instructions, + input, + expected_system_instructions, + expected_request_messages, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[OpenAIIntegration()], + integrations=[OpenAIIntegration(include_prompts=True)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + send_default_pii=True, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) @@ -4573,8 +5599,8 @@ def test_ai_client_span_responses_api_no_pii( with sentry_sdk.traces.start_span(name="openai tx"): client.responses.create( model="gpt-4o", - instructions="You are a coding assistant that talks like a pirate.", - input="How do I check if a Python object is an instance of a class?", + instructions=instructions, + input=input, max_output_tokens=100, temperature=0.7, top_p=0.9, @@ -4585,32 +5611,36 @@ def test_ai_client_span_responses_api_no_pii( spans = [item.payload for item in items] assert len(spans) == 2 - expected_attributes = { + + expected_data = { "gen_ai.operation.name": "responses", "gen_ai.request.max_tokens": 100, "gen_ai.request.temperature": 0.7, "gen_ai.request.top_p": 0.9, "gen_ai.request.reasoning.level": "high", - "gen_ai.request.model": "gpt-4o", + "gen_ai.system": "openai", "gen_ai.response.model": "response-model-id", "gen_ai.response.streaming": False, - "gen_ai.system": "openai", "gen_ai.usage.input_tokens": 20, "gen_ai.usage.input_tokens.cached": 5, "gen_ai.usage.output_tokens": 10, "gen_ai.usage.output_tokens.reasoning": 8, "gen_ai.usage.total_tokens": 30, + "gen_ai.request.messages": safe_serialize(expected_request_messages), + "gen_ai.request.model": "gpt-4o", + "gen_ai.response.text": "the model response", "sentry.op": "gen_ai.responses", "sentry.origin": "auto.ai.openai", "sentry.segment.name": "openai tx", } - for attr, value in expected_attributes.items(): - assert spans[0]["attributes"][attr] == value + if expected_system_instructions is not None: + expected_data["gen_ai.system_instructions"] = safe_serialize( + expected_system_instructions + ) - assert "gen_ai.system_instructions" not in spans[0]["attributes"] - assert "gen_ai.request.messages" not in spans[0]["attributes"] - assert "gen_ai.response.text" not in spans[0]["attributes"] + for attr, value in expected_data.items(): + assert spans[0]["attributes"][attr] == value elif stream_gen_ai_spans: items = capture_items("span") @@ -4618,52 +5648,56 @@ def test_ai_client_span_responses_api_no_pii( with start_transaction(name="openai tx"): client.responses.create( model="gpt-4o", - instructions="You are a coding assistant that talks like a pirate.", - input="How do I check if a Python object is an instance of a class?", + instructions=instructions, + input=input, max_output_tokens=100, temperature=0.7, top_p=0.9, reasoning={"effort": "high"}, - tools=EXAMPLE_TOOLS, ) spans = [item.payload for item in items] assert len(spans) == 1 - expected_attributes = { + + expected_data = { "gen_ai.operation.name": "responses", "gen_ai.request.max_tokens": 100, "gen_ai.request.temperature": 0.7, "gen_ai.request.top_p": 0.9, "gen_ai.request.reasoning.level": "high", - "gen_ai.request.model": "gpt-4o", + "gen_ai.system": "openai", "gen_ai.response.model": "response-model-id", "gen_ai.response.streaming": False, - "gen_ai.system": "openai", "gen_ai.usage.input_tokens": 20, "gen_ai.usage.input_tokens.cached": 5, "gen_ai.usage.output_tokens": 10, "gen_ai.usage.output_tokens.reasoning": 8, "gen_ai.usage.total_tokens": 30, + "gen_ai.request.messages": safe_serialize(expected_request_messages), + "gen_ai.request.model": "gpt-4o", + "gen_ai.response.text": "the model response", "sentry.op": "gen_ai.responses", "sentry.origin": "auto.ai.openai", "sentry.segment.name": "openai tx", } - for attr, value in expected_attributes.items(): + if expected_system_instructions is not None: + expected_data["gen_ai.system_instructions"] = safe_serialize( + expected_system_instructions + ) + + for attr, value in expected_data.items(): assert spans[0]["attributes"][attr] == value - assert "gen_ai.system_instructions" not in spans[0]["attributes"] - assert "gen_ai.request.messages" not in spans[0]["attributes"] - assert "gen_ai.response.text" not in spans[0]["attributes"] else: events = capture_events() with start_transaction(name="openai tx"): client.responses.create( model="gpt-4o", - instructions="You are a coding assistant that talks like a pirate.", - input="How do I check if a Python object is an instance of a class?", + instructions=instructions, + input=input, max_output_tokens=100, temperature=0.7, top_p=0.9, @@ -4676,384 +5710,194 @@ def test_ai_client_span_responses_api_no_pii( assert len(spans) == 1 assert spans[0]["op"] == "gen_ai.responses" assert spans[0]["origin"] == "auto.ai.openai" + expected_data = { "gen_ai.operation.name": "responses", "gen_ai.request.max_tokens": 100, "gen_ai.request.temperature": 0.7, "gen_ai.request.top_p": 0.9, "gen_ai.request.reasoning.level": "high", - "gen_ai.request.model": "gpt-4o", + "gen_ai.system": "openai", "gen_ai.response.model": "response-model-id", "gen_ai.response.streaming": False, - "gen_ai.system": "openai", "gen_ai.usage.input_tokens": 20, "gen_ai.usage.input_tokens.cached": 5, "gen_ai.usage.output_tokens": 10, "gen_ai.usage.output_tokens.reasoning": 8, "gen_ai.usage.total_tokens": 30, + "gen_ai.request.messages": safe_serialize(expected_request_messages[-1:]), + "gen_ai.request.model": "gpt-4o", + "gen_ai.response.text": "the model response", } - for key, value in expected_data.items(): - assert spans[0]["data"][key] == value + if expected_system_instructions is not None: + expected_data["gen_ai.system_instructions"] = safe_serialize( + expected_system_instructions + ) - assert "gen_ai.system_instructions" not in spans[0]["data"] - assert "gen_ai.request.messages" not in spans[0]["data"] - assert "gen_ai.response.text" not in spans[0]["data"] + for attr, value in expected_data.items(): + assert spans[0]["data"][attr] == value @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") -def test_ai_client_span_responses_tool_definitions( - sentry_init, - capture_events, - capture_items, - stream_gen_ai_spans, - span_streaming, -): - sentry_init( - integrations=[OpenAIIntegration()], - disabled_integrations=[StdlibIntegration], - traces_sample_rate=1.0, - stream_gen_ai_spans=stream_gen_ai_spans, - trace_lifecycle="stream" if span_streaming else "static", - ) - - client = OpenAI(api_key="z") - client.responses._post = mock.Mock(return_value=EXAMPLE_RESPONSE) - - if span_streaming: - items = capture_items("span") - - with sentry_sdk.traces.start_span(name="openai tx"): - client.responses.create( - model="gpt-4o", - input="How do I check if a Python object is an instance of a class?", - tools=[ - FunctionToolParam( - type="function", - name="name", - description="description", - parameters={ - "type": "object", - "properties": { - "city": {"type": "string"}, - "state": {"type": "string"}, - }, - "required": ["city", "state"], - "additionalProperties": False, - }, - strict=True, - ), - CustomToolParam( - type="custom", name="name", description="description" - ), - WebSearchToolParam(type="web_search"), - ], - ) - - sentry_sdk.flush() - spans = [item.payload for item in items] - assert json.loads(spans[0]["attributes"][SPANDATA.GEN_AI_TOOL_DEFINITIONS]) == [ - { - "type": "function", - "name": "name", - "description": "description", - "parameters": { - "type": "object", - "properties": { - "city": {"type": "string"}, - "state": {"type": "string"}, - }, - "required": ["city", "state"], - "additionalProperties": False, - }, - }, - { - "type": "custom", - "name": "name", - "description": "description", - }, - { - "type": "web_search", - }, - ] - elif stream_gen_ai_spans: - items = capture_items("span") - - with start_transaction(name="openai tx"): - client.responses.create( - model="gpt-4o", - input="How do I check if a Python object is an instance of a class?", - tools=[ - FunctionToolParam( - type="function", - name="name", - description="description", - parameters={ - "type": "object", - "properties": { - "city": {"type": "string"}, - "state": {"type": "string"}, - }, - "required": ["city", "state"], - "additionalProperties": False, - }, - strict=True, - ), - CustomToolParam( - type="custom", name="name", description="description" - ), - WebSearchToolParam(type="web_search"), - ], - ) - - spans = [item.payload for item in items] - assert json.loads(spans[0]["attributes"][SPANDATA.GEN_AI_TOOL_DEFINITIONS]) == [ - { - "type": "function", - "name": "name", - "description": "description", - "parameters": { - "type": "object", - "properties": { - "city": {"type": "string"}, - "state": {"type": "string"}, - }, - "required": ["city", "state"], - "additionalProperties": False, - }, - }, - { - "type": "custom", - "name": "name", - "description": "description", - }, +@pytest.mark.parametrize( + "data_collection,extra_kwargs,expected_present,expected_absent,include_prompts", + [ + pytest.param( + {"gen_ai": {"inputs": True}}, { - "type": "web_search", + "instructions": "You are a coding assistant that talks like a pirate.", + "input": "How do I check if a Python object is an instance of a class?", + "tools": EXAMPLE_TOOLS, }, - ] - else: - events = capture_events() - - with start_transaction(name="openai tx"): - client.responses.create( - model="gpt-4o", - input="How do I check if a Python object is an instance of a class?", - tools=[ - FunctionToolParam( - type="function", - name="name", - description="description", - parameters={ - "type": "object", - "properties": { - "city": {"type": "string"}, - "state": {"type": "string"}, - }, - "required": ["city", "state"], - "additionalProperties": False, - }, - strict=True, - ), - CustomToolParam( - type="custom", name="name", description="description" - ), - WebSearchToolParam(type="web_search"), - ], - ) - - (transaction,) = events - spans = transaction["spans"] - - assert json.loads(spans[0]["data"][SPANDATA.GEN_AI_TOOL_DEFINITIONS]) == [ { - "type": "function", - "name": "name", - "description": "description", - "parameters": { - "type": "object", - "properties": { - "city": {"type": "string"}, - "state": {"type": "string"}, - }, - "required": ["city", "state"], - "additionalProperties": False, - }, + SPANDATA.GEN_AI_REQUEST_MESSAGES: safe_serialize( + ["How do I check if a Python object is an instance of a class?"] + ), + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: safe_serialize( + [ + { + "type": "text", + "content": "You are a coding assistant that talks like a pirate.", + } + ] + ), + SPANDATA.GEN_AI_TOOL_DEFINITIONS: safe_serialize(EXAMPLE_TOOLS), }, + [], + True, + id="inputs-enabled-string-input", + ), + pytest.param( + {"gen_ai": {"inputs": True}}, { - "type": "custom", - "name": "name", - "description": "description", + "instructions": "You are a coding assistant that talks like a pirate.", }, { - "type": "web_search", + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: safe_serialize( + [ + { + "type": "text", + "content": "You are a coding assistant that talks like a pirate.", + } + ] + ), }, - ] - - -@pytest.mark.parametrize("span_streaming", [True, False]) -@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) -@pytest.mark.parametrize( - "instructions,input,expected_system_instructions,expected_request_messages", - [ - ( - omit, - "How do I check if a Python object is an instance of a class?", - None, - ["How do I check if a Python object is an instance of a class?"], - ), - ( - None, - "How do I check if a Python object is an instance of a class?", - None, - ["How do I check if a Python object is an instance of a class?"], - ), - ( - "You are a coding assistant that talks like a pirate.", - [ - { - "role": "system", - "content": "You are a helpful assistant.", - }, - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "hello"}, - ], - [ - { - "type": "text", - "content": "You are a coding assistant that talks like a pirate.", - }, - {"type": "text", "content": "You are a helpful assistant."}, - ], [ - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "hello"}, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_TOOL_DEFINITIONS, ], + True, + id="inputs-enabled-instructions-only", ), - ( - "You are a coding assistant that talks like a pirate.", - [ - { - "type": "message", - "role": "system", - "content": "You are a helpful assistant.", - }, - { - "type": "message", - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"type": "message", "role": "user", "content": "hello"}, - ], - [ - { - "type": "text", - "content": "You are a coding assistant that talks like a pirate.", - }, - {"type": "text", "content": "You are a helpful assistant."}, - ], - [ - { - "type": "message", - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"type": "message", "role": "user", "content": "hello"}, - ], + pytest.param( + {"gen_ai": {"inputs": True}}, + { + "instructions": "You are a coding assistant that talks like a pirate.", + "input": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "hello"}, + ], + }, + { + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: safe_serialize( + [ + { + "type": "text", + "content": "You are a coding assistant that talks like a pirate.", + }, + {"type": "text", "content": "You are a helpful assistant."}, + ] + ), + SPANDATA.GEN_AI_REQUEST_MESSAGES: safe_serialize( + [{"role": "user", "content": "hello"}] + ), + }, + [SPANDATA.GEN_AI_TOOL_DEFINITIONS], + True, + id="inputs-enabled-list-input-with-system-message", ), - ( - "You are a coding assistant that talks like a pirate.", - [ - { - "role": "system", - "content": [ - {"type": "text", "text": "You are a helpful assistant."}, - {"type": "text", "text": "Be concise and clear."}, - ], - }, - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "hello"}, - ], - [ - { - "type": "text", - "content": "You are a coding assistant that talks like a pirate.", - }, - {"type": "text", "content": "You are a helpful assistant."}, - {"type": "text", "content": "Be concise and clear."}, - ], + pytest.param( + {"gen_ai": {"inputs": False}}, + { + "instructions": "You are a coding assistant that talks like a pirate.", + "input": "How do I check if a Python object is an instance of a class?", + "tools": EXAMPLE_TOOLS, + }, + {}, [ - { - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"role": "user", "content": "hello"}, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + SPANDATA.GEN_AI_TOOL_DEFINITIONS, ], + True, + id="inputs-disabled", ), - ( - "You are a coding assistant that talks like a pirate.", - [ - { - "type": "message", - "role": "system", - "content": [ - {"type": "text", "text": "You are a helpful assistant."}, - {"type": "text", "text": "Be concise and clear."}, - ], - }, - { - "type": "message", - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"type": "message", "role": "user", "content": "hello"}, - ], + pytest.param( + {}, + { + "input": "How do I check if a Python object is an instance of a class?", + "tools": EXAMPLE_TOOLS, + }, + { + SPANDATA.GEN_AI_REQUEST_MESSAGES: safe_serialize( + ["How do I check if a Python object is an instance of a class?"] + ), + SPANDATA.GEN_AI_TOOL_DEFINITIONS: safe_serialize(EXAMPLE_TOOLS), + }, + [SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS], + True, + id="gen-ai-omitted-defaults-to-enabled", + ), + pytest.param( + {"gen_ai": {"inputs": True}}, + {}, + {}, [ - { - "type": "text", - "content": "You are a coding assistant that talks like a pirate.", - }, - {"type": "text", "content": "You are a helpful assistant."}, - {"type": "text", "content": "Be concise and clear."}, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + SPANDATA.GEN_AI_TOOL_DEFINITIONS, ], + True, + id="inputs-enabled-no-input-provided", + ), + pytest.param( + {"gen_ai": {"inputs": True}}, + { + "instructions": "You are a coding assistant that talks like a pirate.", + "input": "How do I check if a Python object is an instance of a class?", + "tools": EXAMPLE_TOOLS, + }, + {}, [ - { - "type": "message", - "role": "user", - "content": "Message demonstrating the absence of truncation.", - }, - {"type": "message", "role": "user", "content": "hello"}, + SPANDATA.GEN_AI_REQUEST_MESSAGES, + SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, + SPANDATA.GEN_AI_TOOL_DEFINITIONS, ], + False, + id="include-prompts-disabled-overrides-inputs-enabled", ), ], ) @pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") -def test_ai_client_span_responses_api( +def test_responses_api_data_collection( sentry_init, capture_events, capture_items, - instructions, - input, - expected_system_instructions, - expected_request_messages, + data_collection, + extra_kwargs, + expected_present, + expected_absent, + include_prompts, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[OpenAIIntegration(include_prompts=True)], + integrations=[OpenAIIntegration(include_prompts=include_prompts)], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, - send_default_pii=True, + _experiments={"data_collection": data_collection}, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) @@ -5061,375 +5905,425 @@ def test_ai_client_span_responses_api( client = OpenAI(api_key="z") client.responses._post = mock.Mock(return_value=EXAMPLE_RESPONSE) + create_kwargs = { + "model": "gpt-4o", + "max_output_tokens": 100, + "temperature": 0.7, + "top_p": 0.9, + "reasoning": {"effort": "high"}, + } + create_kwargs.update(extra_kwargs) + if span_streaming: items = capture_items("span") with sentry_sdk.traces.start_span(name="openai tx"): - client.responses.create( - model="gpt-4o", - instructions=instructions, - input=input, - max_output_tokens=100, - temperature=0.7, - top_p=0.9, - reasoning={"effort": "high"}, - ) + client.responses.create(**create_kwargs) sentry_sdk.flush() spans = [item.payload for item in items] assert len(spans) == 2 - - expected_data = { - "gen_ai.operation.name": "responses", - "gen_ai.request.max_tokens": 100, - "gen_ai.request.temperature": 0.7, - "gen_ai.request.top_p": 0.9, - "gen_ai.request.reasoning.level": "high", - "gen_ai.system": "openai", - "gen_ai.response.model": "response-model-id", - "gen_ai.response.streaming": False, - "gen_ai.usage.input_tokens": 20, - "gen_ai.usage.input_tokens.cached": 5, - "gen_ai.usage.output_tokens": 10, - "gen_ai.usage.output_tokens.reasoning": 8, - "gen_ai.usage.total_tokens": 30, - "gen_ai.request.messages": safe_serialize(expected_request_messages), - "gen_ai.request.model": "gpt-4o", - "gen_ai.response.text": "the model response", - "sentry.op": "gen_ai.responses", - "sentry.origin": "auto.ai.openai", - "sentry.segment.name": "openai tx", - } - - if expected_system_instructions is not None: - expected_data["gen_ai.system_instructions"] = safe_serialize( - expected_system_instructions - ) - - for attr, value in expected_data.items(): - assert spans[0]["attributes"][attr] == value - + span_data = spans[0]["attributes"] elif stream_gen_ai_spans: items = capture_items("span") with start_transaction(name="openai tx"): - client.responses.create( - model="gpt-4o", - instructions=instructions, - input=input, - max_output_tokens=100, - temperature=0.7, - top_p=0.9, - reasoning={"effort": "high"}, - ) + client.responses.create(**create_kwargs) spans = [item.payload for item in items] assert len(spans) == 1 - - expected_data = { - "gen_ai.operation.name": "responses", - "gen_ai.request.max_tokens": 100, - "gen_ai.request.temperature": 0.7, - "gen_ai.request.top_p": 0.9, - "gen_ai.request.reasoning.level": "high", - "gen_ai.system": "openai", - "gen_ai.response.model": "response-model-id", - "gen_ai.response.streaming": False, - "gen_ai.usage.input_tokens": 20, - "gen_ai.usage.input_tokens.cached": 5, - "gen_ai.usage.output_tokens": 10, - "gen_ai.usage.output_tokens.reasoning": 8, - "gen_ai.usage.total_tokens": 30, - "gen_ai.request.messages": safe_serialize(expected_request_messages), - "gen_ai.request.model": "gpt-4o", - "gen_ai.response.text": "the model response", - "sentry.op": "gen_ai.responses", - "sentry.origin": "auto.ai.openai", - "sentry.segment.name": "openai tx", - } - - if expected_system_instructions is not None: - expected_data["gen_ai.system_instructions"] = safe_serialize( - expected_system_instructions - ) - - for attr, value in expected_data.items(): - assert spans[0]["attributes"][attr] == value - + span_data = spans[0]["attributes"] else: events = capture_events() with start_transaction(name="openai tx"): - client.responses.create( - model="gpt-4o", - instructions=instructions, - input=input, - max_output_tokens=100, - temperature=0.7, - top_p=0.9, - reasoning={"effort": "high"}, - ) + client.responses.create(**create_kwargs) (transaction,) = events spans = transaction["spans"] assert len(spans) == 1 assert spans[0]["op"] == "gen_ai.responses" - assert spans[0]["origin"] == "auto.ai.openai" + span_data = spans[0]["data"] - expected_data = { - "gen_ai.operation.name": "responses", - "gen_ai.request.max_tokens": 100, - "gen_ai.request.temperature": 0.7, - "gen_ai.request.top_p": 0.9, - "gen_ai.request.reasoning.level": "high", - "gen_ai.system": "openai", - "gen_ai.response.model": "response-model-id", - "gen_ai.response.streaming": False, - "gen_ai.usage.input_tokens": 20, - "gen_ai.usage.input_tokens.cached": 5, - "gen_ai.usage.output_tokens": 10, - "gen_ai.usage.output_tokens.reasoning": 8, - "gen_ai.usage.total_tokens": 30, - "gen_ai.request.messages": safe_serialize(expected_request_messages[-1:]), - "gen_ai.request.model": "gpt-4o", - "gen_ai.response.text": "the model response", - } + # Non-input data is always collected, regardless of data collection config + assert span_data["gen_ai.operation.name"] == "responses" + assert span_data["gen_ai.request.model"] == "gpt-4o" + assert span_data["gen_ai.request.max_tokens"] == 100 + assert span_data["gen_ai.request.temperature"] == 0.7 + assert span_data["gen_ai.request.top_p"] == 0.9 + assert span_data["gen_ai.request.reasoning.level"] == "high" + assert span_data["gen_ai.system"] == "openai" - if expected_system_instructions is not None: - expected_data["gen_ai.system_instructions"] = safe_serialize( - expected_system_instructions - ) + for key, value in expected_present.items(): + assert span_data[key] == value - for attr, value in expected_data.items(): - assert spans[0]["data"][attr] == value + for key in expected_absent: + assert key not in span_data + + +def _make_responses_api_output_message(content): + return ResponseOutputMessage( + id="message-id", + content=content, + role="assistant", + status="completed", + type="message", + ) + + +def _make_responses_api_function_call(): + return ResponseFunctionToolCall( + id="fc-id", + call_id="call_123", + name="get_current_weather", + arguments='{"location": "San Francisco, CA"}', + type="function_call", + ) + + +def _make_responses_api_response(output): + return Response( + id="chat-id", + output=output, + parallel_tool_calls=False, + tool_choice="none", + tools=[], + created_at=10000000, + model="response-model-id", + object="response", + usage=ResponseUsage( + input_tokens=20, + input_tokens_details=InputTokensDetails( + cached_tokens=5, + cache_write_tokens=0, + ), + output_tokens=10, + output_tokens_details=OutputTokensDetails( + reasoning_tokens=8, + ), + total_tokens=30, + ), + ) + + +def _collect_responses_span_data( + capture_events, capture_items, span_streaming, stream_gen_ai_spans, create +): + if span_streaming or stream_gen_ai_spans: + items = capture_items("span") + + with start_transaction(name="openai tx"): + create() + + sentry_sdk.flush() + (span,) = (item.payload for item in items) + return span["attributes"] + + events = capture_events() + + with start_transaction(name="openai tx"): + create() + + (transaction,) = events + (span,) = transaction["spans"] + assert span["op"] == "gen_ai.responses" + return span["data"] @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) @pytest.mark.parametrize( - "data_collection,extra_kwargs,expected_present,expected_absent,include_prompts", + "data_collection,send_default_pii,expect_output", [ pytest.param( - {"gen_ai": {"inputs": True}}, - { - "instructions": "You are a coding assistant that talks like a pirate.", - "input": "How do I check if a Python object is an instance of a class?", - "tools": EXAMPLE_TOOLS, - }, - { - SPANDATA.GEN_AI_REQUEST_MESSAGES: safe_serialize( - ["How do I check if a Python object is an instance of a class?"] - ), - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: safe_serialize( - [ - { - "type": "text", - "content": "You are a coding assistant that talks like a pirate.", - } - ] - ), - SPANDATA.GEN_AI_TOOL_DEFINITIONS: safe_serialize(EXAMPLE_TOOLS), - }, - [], + {"gen_ai": {"outputs": True}}, + False, True, - id="inputs-enabled-string-input", + 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": True}}, - { - "instructions": "You are a coding assistant that talks like a pirate.", - }, - { - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: safe_serialize( - [ - { - "type": "text", - "content": "You are a coding assistant that talks like a pirate.", - } + None, + False, + False, + id="no-gen-ai-data-collection-falls-back-to-send-default-pii", + ), + pytest.param( + None, + True, + True, + id="no-gen-ai-data-collection-pii-enabled-collects", + ), + ], +) +@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") +def test_responses_api_data_collection_outputs( + sentry_init, + capture_events, + capture_items, + data_collection, + send_default_pii, + expect_output, + stream_gen_ai_spans, + span_streaming, +): + init_kwargs = { + "integrations": [OpenAIIntegration()], + "disabled_integrations": [StdlibIntegration], + "traces_sample_rate": 1.0, + "send_default_pii": send_default_pii, + "stream_gen_ai_spans": stream_gen_ai_spans, + "trace_lifecycle": "stream" if span_streaming else "static", + } + if data_collection is not None: + init_kwargs["_experiments"] = {"data_collection": data_collection} + + sentry_init(**init_kwargs) + + client = OpenAI(api_key="z") + client.responses._post = mock.Mock( + return_value=_make_responses_api_response( + output=[ + _make_responses_api_output_message( + content=[ + ResponseOutputText( + annotations=[], + text="the model response", + type="output_text", + ), + ] + ), + _make_responses_api_function_call(), + ] + ) + ) + + span_data = _collect_responses_span_data( + capture_events, + capture_items, + span_streaming, + stream_gen_ai_spans, + lambda: client.responses.create(model="gpt-4o", input="hello"), + ) + + assert span_data[SPANDATA.GEN_AI_RESPONSE_MODEL] == "response-model-id" + + if expect_output: + assert "the model response" in span_data[SPANDATA.GEN_AI_RESPONSE_TEXT] + assert "get_current_weather" in span_data[SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + else: + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span_data + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS not in span_data + + +@pytest.mark.parametrize("span_streaming", [True, False]) +@pytest.mark.parametrize("stream_gen_ai_spans", [True, False]) +@pytest.mark.parametrize( + "get_output,expect_text,expect_tool_calls", + [ + pytest.param( + lambda: [ + _make_responses_api_output_message( + content=[ + ResponseOutputText( + annotations=[], + text="the model response", + type="output_text", + ), ] ), - }, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_TOOL_DEFINITIONS, ], True, - id="inputs-enabled-instructions-only", + False, + id="message-only", ), pytest.param( - {"gen_ai": {"inputs": True}}, - { - "instructions": "You are a coding assistant that talks like a pirate.", - "input": [ - {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "hello"}, - ], - }, - { - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS: safe_serialize( - [ - { - "type": "text", - "content": "You are a coding assistant that talks like a pirate.", - }, - {"type": "text", "content": "You are a helpful assistant."}, - ] - ), - SPANDATA.GEN_AI_REQUEST_MESSAGES: safe_serialize( - [{"role": "user", "content": "hello"}] - ), - }, - [SPANDATA.GEN_AI_TOOL_DEFINITIONS], + lambda: [_make_responses_api_function_call()], + False, True, - id="inputs-enabled-list-input-with-system-message", + id="function-call-only", ), pytest.param( - {"gen_ai": {"inputs": False}}, - { - "instructions": "You are a coding assistant that talks like a pirate.", - "input": "How do I check if a Python object is an instance of a class?", - "tools": EXAMPLE_TOOLS, - }, - {}, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_TOOL_DEFINITIONS, + lambda: [ + _make_responses_api_output_message( + content=[ + ResponseOutputRefusal( + refusal="I cannot help with that.", + type="refusal", + ), + ] + ), ], True, - id="inputs-disabled", + False, + id="non-text-content-falls-back-to-dict", ), + ], +) +@pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") +def test_responses_api_data_collection_outputs_shapes( + sentry_init, + capture_events, + capture_items, + get_output, + expect_text, + expect_tool_calls, + stream_gen_ai_spans, + span_streaming, +): + sentry_init( + integrations=[OpenAIIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + _experiments={"data_collection": {"gen_ai": {"outputs": True}}}, + stream_gen_ai_spans=stream_gen_ai_spans, + trace_lifecycle="stream" if span_streaming else "static", + ) + + client = OpenAI(api_key="z") + client.responses._post = mock.Mock( + return_value=_make_responses_api_response(output=get_output()) + ) + + span_data = _collect_responses_span_data( + capture_events, + capture_items, + span_streaming, + stream_gen_ai_spans, + lambda: client.responses.create(model="gpt-4o", input="hello"), + ) + + if expect_text: + assert SPANDATA.GEN_AI_RESPONSE_TEXT in span_data + else: + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span_data + + if expect_tool_calls: + assert "get_current_weather" in span_data[SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS] + else: + assert SPANDATA.GEN_AI_RESPONSE_TOOL_CALLS 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,expect_output", + [ pytest.param( - {}, - { - "input": "How do I check if a Python object is an instance of a class?", - "tools": EXAMPLE_TOOLS, - }, - { - SPANDATA.GEN_AI_REQUEST_MESSAGES: safe_serialize( - ["How do I check if a Python object is an instance of a class?"] - ), - SPANDATA.GEN_AI_TOOL_DEFINITIONS: safe_serialize(EXAMPLE_TOOLS), - }, - [SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS], + {"gen_ai": {"outputs": True}}, True, - id="gen-ai-omitted-defaults-to-enabled", + id="gen-ai-outputs-enabled", ), pytest.param( - {"gen_ai": {"inputs": True}}, - {}, - {}, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_TOOL_DEFINITIONS, - ], - True, - id="inputs-enabled-no-input-provided", + {"gen_ai": {"outputs": False}}, + False, + id="gen-ai-outputs-disabled", ), pytest.param( - {"gen_ai": {"inputs": True}}, - { - "instructions": "You are a coding assistant that talks like a pirate.", - "input": "How do I check if a Python object is an instance of a class?", - "tools": EXAMPLE_TOOLS, - }, {}, - [ - SPANDATA.GEN_AI_REQUEST_MESSAGES, - SPANDATA.GEN_AI_SYSTEM_INSTRUCTIONS, - SPANDATA.GEN_AI_TOOL_DEFINITIONS, - ], - False, - id="include-prompts-disabled-overrides-inputs-enabled", + True, + id="gen-ai-omitted-defaults-to-enabled", ), ], ) @pytest.mark.skipif(SKIP_RESPONSES_TESTS, reason="Responses API not available") -def test_responses_api_data_collection( +def test_streaming_responses_api_data_collection_outputs( sentry_init, capture_events, capture_items, data_collection, - extra_kwargs, - expected_present, - expected_absent, - include_prompts, + expect_output, + get_model_response, + server_side_event_chunks, stream_gen_ai_spans, span_streaming, ): sentry_init( - integrations=[OpenAIIntegration(include_prompts=include_prompts)], + integrations=[OpenAIIntegration()], disabled_integrations=[StdlibIntegration], traces_sample_rate=1.0, + send_default_pii=False, _experiments={"data_collection": data_collection}, stream_gen_ai_spans=stream_gen_ai_spans, trace_lifecycle="stream" if span_streaming else "static", ) client = OpenAI(api_key="z") - client.responses._post = mock.Mock(return_value=EXAMPLE_RESPONSE) - - create_kwargs = { - "model": "gpt-4o", - "max_output_tokens": 100, - "temperature": 0.7, - "top_p": 0.9, - "reasoning": {"effort": "high"}, - } - create_kwargs.update(extra_kwargs) + returned_stream = get_model_response( + server_side_event_chunks( + EXAMPLE_RESPONSES_STREAM, + ) + ) - if span_streaming: + if span_streaming or stream_gen_ai_spans: items = capture_items("span") - with sentry_sdk.traces.start_span(name="openai tx"): - client.responses.create(**create_kwargs) + with mock.patch.object( + client.responses._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = client.responses.create( + model="some-model", + input="hello", + stream=True, + ) + response_string = "" + for item in response_stream: + if hasattr(item, "delta"): + response_string += item.delta + assert response_string == "hello world" sentry_sdk.flush() - spans = [item.payload for item in items] - - assert len(spans) == 2 - span_data = spans[0]["attributes"] - elif stream_gen_ai_spans: - items = capture_items("span") - - with start_transaction(name="openai tx"): - client.responses.create(**create_kwargs) - - spans = [item.payload for item in items] - - assert len(spans) == 1 - span_data = spans[0]["attributes"] + (span,) = (item.payload for item in items) + span_data = span["attributes"] else: events = capture_events() - with start_transaction(name="openai tx"): - client.responses.create(**create_kwargs) + with mock.patch.object( + client.responses._client._client, + "send", + return_value=returned_stream, + ), start_transaction(name="openai tx"): + response_stream = client.responses.create( + model="some-model", + input="hello", + stream=True, + ) + response_string = "" + for item in response_stream: + if hasattr(item, "delta"): + response_string += item.delta + assert response_string == "hello world" (transaction,) = events - spans = transaction["spans"] - - assert len(spans) == 1 - assert spans[0]["op"] == "gen_ai.responses" - span_data = spans[0]["data"] - - # Non-input data is always collected, regardless of data collection config - assert span_data["gen_ai.operation.name"] == "responses" - assert span_data["gen_ai.request.model"] == "gpt-4o" - assert span_data["gen_ai.request.max_tokens"] == 100 - assert span_data["gen_ai.request.temperature"] == 0.7 - assert span_data["gen_ai.request.top_p"] == 0.9 - assert span_data["gen_ai.request.reasoning.level"] == "high" - assert span_data["gen_ai.system"] == "openai" - - for key, value in expected_present.items(): - assert span_data[key] == value + (span,) = transaction["spans"] + span_data = span["data"] - for key in expected_absent: - assert key not in span_data + if expect_output: + assert "hello world" in span_data[SPANDATA.GEN_AI_RESPONSE_TEXT] + else: + assert SPANDATA.GEN_AI_RESPONSE_TEXT not in span_data @pytest.mark.parametrize("span_streaming", [True, False]) From 770189e7d3101371dc442738b5846f7eb3a238d2 Mon Sep 17 00:00:00 2001 From: Erica Pisani Date: Tue, 4 Aug 2026 16:19:00 -0400 Subject: [PATCH 2/2] lint --- sentry_sdk/integrations/openai.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index 8525edc930..27757f8893 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -1036,7 +1036,9 @@ def _wrap_synchronous_completions_chunk_iterator( all_responses = ["".join(chunk) for chunk in data_buf] if has_data_collection_enabled(client.options): if client.options["data_collection"]["gen_ai"]["outputs"]: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses + ) elif should_send_default_pii() and integration.include_prompts: set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) @@ -1105,7 +1107,9 @@ async def _wrap_asynchronous_completions_chunk_iterator( all_responses = ["".join(chunk) for chunk in data_buf] if has_data_collection_enabled(client.options): if client.options["data_collection"]["gen_ai"]["outputs"]: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses + ) elif should_send_default_pii() and integration.include_prompts: set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) @@ -1176,7 +1180,9 @@ def _wrap_synchronous_responses_event_iterator( all_responses = ["".join(chunk) for chunk in data_buf] if has_data_collection_enabled(client.options): if client.options["data_collection"]["gen_ai"]["outputs"]: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses + ) elif should_send_default_pii() and integration.include_prompts: set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) @@ -1248,7 +1254,9 @@ async def _wrap_asynchronous_responses_event_iterator( if has_data_collection_enabled(client.options): if client.options["data_collection"]["gen_ai"]["outputs"]: - set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses) + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses + ) elif should_send_default_pii() and integration.include_prompts: set_data_normalized(span, SPANDATA.GEN_AI_RESPONSE_TEXT, all_responses)