From 7091e62a0d4ae1c8921782e0c0e3d989b1897a49 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Fri, 20 Feb 2026 13:34:42 +0100 Subject: [PATCH 1/4] fix(openai): Avoid consuming iterables passed to the Embeddings API --- sentry_sdk/integrations/openai.py | 57 ++++-- tests/integrations/openai/test_openai.py | 231 ++++++++++++++++++++++- 2 files changed, 258 insertions(+), 30 deletions(-) diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index 863f146a51..766c4f3261 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -50,7 +50,7 @@ from sentry_sdk.tracing import Span from sentry_sdk._types import TextPart - from openai.types.responses import ResponseInputParam + from openai.types.responses import ResponseInputParam, SequenceNotStr from openai import Omit try: @@ -220,20 +220,6 @@ def _calculate_token_usage( ) -def _get_input_messages( - kwargs: "dict[str, Any]", -) -> "Optional[Union[Iterable[Any], list[str]]]": - # Input messages (the prompt or data sent to the model) - messages = kwargs.get("messages") - if messages is None: - messages = kwargs.get("input") - - if isinstance(messages, str): - messages = [messages] - - return messages - - def _commmon_set_input_data( span: "Span", kwargs: "dict[str, Any]", @@ -413,14 +399,45 @@ def _set_embeddings_input_data( kwargs: "dict[str, Any]", integration: "OpenAIIntegration", ) -> None: - messages = _get_input_messages(kwargs) + messages: "Union[str, SequenceNotStr[str], Iterable[int], Iterable[Iterable[int]]]" = kwargs.get( + "input" + ) if ( - messages is not None - and len(messages) > 0 # type: ignore - and should_send_default_pii() - and integration.include_prompts + not should_send_default_pii() + or not integration.include_prompts + or messages is None ): + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + _commmon_set_input_data(span, kwargs) + + return + + if isinstance(messages, str): + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + _commmon_set_input_data(span, kwargs) + + normalized_messages = normalize_message_roles([messages]) # type: ignore + scope = sentry_sdk.get_current_scope() + messages_data = truncate_and_annotate_embedding_inputs( + normalized_messages, span, scope + ) + if messages_data is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_EMBEDDINGS_INPUT, messages_data, unpack=False + ) + + return + + if not isinstance(messages, Iterable): + set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") + _commmon_set_input_data(span, kwargs) + return + + messages = list(messages) + kwargs["input"] = messages + + if len(messages) > 0: normalized_messages = normalize_message_roles(messages) # type: ignore scope = sentry_sdk.get_current_scope() messages_data = truncate_and_annotate_embedding_inputs( diff --git a/tests/integrations/openai/test_openai.py b/tests/integrations/openai/test_openai.py index 094b659b2c..4ac49ff124 100644 --- a/tests/integrations/openai/test_openai.py +++ b/tests/integrations/openai/test_openai.py @@ -928,11 +928,16 @@ async def test_bad_chat_completion_async(sentry_init, capture_events): assert event["level"] == "error" +@pytest.mark.asyncio @pytest.mark.parametrize( "send_default_pii, include_prompts", - [(True, True), (True, False), (False, True), (False, False)], + [ + (True, False), + (False, True), + (False, False), + ], ) -def test_embeddings_create( +def test_embeddings_create_no_pii( sentry_init, capture_events, send_default_pii, include_prompts ): sentry_init( @@ -966,10 +971,110 @@ def test_embeddings_create( assert tx["type"] == "transaction" span = tx["spans"][0] assert span["op"] == "gen_ai.embeddings" - if send_default_pii and include_prompts: - assert "hello" in span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + + 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.asyncio +@pytest.mark.parametrize( + "input", + [ + pytest.param( + "hello", + id="string", + ), + pytest.param( + ["First text", "Second text", "Third text"], + id="string_sequence", + ), + pytest.param( + iter(["First text", "Second text", "Third text"]), + id="string_iterable", + ), + pytest.param( + [5, 8, 13, 21, 34], + id="tokens", + ), + pytest.param( + iter( + [5, 8, 13, 21, 34], + ), + id="token_iterable", + ), + pytest.param( + [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ], + id="tokens_sequence", + ), + pytest.param( + iter( + [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ] + ), + id="tokens_sequence_iterable", + ), + ], +) +def test_embeddings_create(sentry_init, capture_events, input, request): + sentry_init( + integrations=[OpenAIIntegration(include_prompts=True)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + + 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) + with start_transaction(name="openai tx"): + response = client.embeddings.create(input=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" + + param_id = request.node.callspec.id + if param_id == "string": + assert json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == ["hello"] + elif param_id == "string_sequence" or param_id == "string_iterable": + assert json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == [ + "First text", + "Second text", + "Third text", + ] + elif param_id == "tokens" or param_id == "token_iterable": + assert json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == [ + 5, + 8, + 13, + 21, + 34, + ] else: - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["data"] + assert json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ] assert span["data"]["gen_ai.usage.input_tokens"] == 20 assert span["data"]["gen_ai.usage.total_tokens"] == 30 @@ -978,9 +1083,13 @@ def test_embeddings_create( @pytest.mark.asyncio @pytest.mark.parametrize( "send_default_pii, include_prompts", - [(True, True), (True, False), (False, True), (False, False)], + [ + (True, False), + (False, True), + (False, False), + ], ) -async def test_embeddings_create_async( +async def test_embeddings_create_async_no_pii( sentry_init, capture_events, send_default_pii, include_prompts ): sentry_init( @@ -1014,10 +1123,112 @@ async def test_embeddings_create_async( assert tx["type"] == "transaction" span = tx["spans"][0] assert span["op"] == "gen_ai.embeddings" - if send_default_pii and include_prompts: - assert "hello" in span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT] + + 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.asyncio +@pytest.mark.parametrize( + "input", + [ + pytest.param( + "hello", + id="string", + ), + pytest.param( + ["First text", "Second text", "Third text"], + id="string_sequence", + ), + pytest.param( + iter(["First text", "Second text", "Third text"]), + id="string_iterable", + ), + pytest.param( + [5, 8, 13, 21, 34], + id="tokens", + ), + pytest.param( + iter( + [5, 8, 13, 21, 34], + ), + id="token_iterable", + ), + pytest.param( + [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ], + id="tokens_sequence", + ), + pytest.param( + iter( + [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ] + ), + id="tokens_sequence_iterable", + ), + ], +) +async def test_embeddings_create_async(sentry_init, capture_events, input, request): + sentry_init( + integrations=[OpenAIIntegration(include_prompts=True)], + traces_sample_rate=1.0, + send_default_pii=True, + ) + events = capture_events() + + 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) + with start_transaction(name="openai tx"): + response = await client.embeddings.create( + input=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" + + param_id = request.node.callspec.id + if param_id == "string": + assert json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == ["hello"] + elif param_id == "string_sequence" or param_id == "string_iterable": + assert json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == [ + "First text", + "Second text", + "Third text", + ] + elif param_id == "tokens" or param_id == "token_iterable": + assert json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == [ + 5, + 8, + 13, + 21, + 34, + ] else: - assert SPANDATA.GEN_AI_EMBEDDINGS_INPUT not in span["data"] + assert json.loads(span["data"][SPANDATA.GEN_AI_EMBEDDINGS_INPUT]) == [ + [5, 8, 13, 21, 34], + [8, 13, 21, 34, 55], + ] assert span["data"]["gen_ai.usage.input_tokens"] == 20 assert span["data"]["gen_ai.usage.total_tokens"] == 30 From 5c03b6e7803334e342cf374bdf7bba680e7e5811 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Fri, 20 Feb 2026 14:53:14 +0100 Subject: [PATCH 2/4] remove unused type ignore --- sentry_sdk/integrations/openai.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index 766c4f3261..3e51a60e0d 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -438,7 +438,7 @@ def _set_embeddings_input_data( kwargs["input"] = messages if len(messages) > 0: - normalized_messages = normalize_message_roles(messages) # type: ignore + normalized_messages = normalize_message_roles(messages) scope = sentry_sdk.get_current_scope() messages_data = truncate_and_annotate_embedding_inputs( normalized_messages, span, scope From bf46e1eea7041ed7816defb3fc1f74468b550813 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Fri, 20 Feb 2026 15:05:33 +0100 Subject: [PATCH 3/4] add dict edge case --- sentry_sdk/integrations/openai.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index 3e51a60e0d..2d38f15646 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -429,7 +429,8 @@ def _set_embeddings_input_data( return - if not isinstance(messages, Iterable): + # Special case following https://github.com/openai/openai-python/blob/3e0c05b84a2056870abf3bd6a5e7849020209cc3/src/openai/_utils/_transform.py#L194-L197 + if not isinstance(messages, Iterable) or isinstance(messages, dict): set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") _commmon_set_input_data(span, kwargs) return From 8bfb726fddd9f6bcb34048f5d0f9f1019f2385c9 Mon Sep 17 00:00:00 2001 From: Alexander Alderman Webb Date: Fri, 20 Feb 2026 15:11:19 +0100 Subject: [PATCH 4/4] . --- sentry_sdk/integrations/openai.py | 2 +- tests/integrations/openai/test_openai.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index 2d38f15646..70dcda0384 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -429,7 +429,7 @@ def _set_embeddings_input_data( return - # Special case following https://github.com/openai/openai-python/blob/3e0c05b84a2056870abf3bd6a5e7849020209cc3/src/openai/_utils/_transform.py#L194-L197 + # dict special case following https://github.com/openai/openai-python/blob/3e0c05b84a2056870abf3bd6a5e7849020209cc3/src/openai/_utils/_transform.py#L194-L197 if not isinstance(messages, Iterable) or isinstance(messages, dict): set_data_normalized(span, SPANDATA.GEN_AI_OPERATION_NAME, "embeddings") _commmon_set_input_data(span, kwargs) diff --git a/tests/integrations/openai/test_openai.py b/tests/integrations/openai/test_openai.py index 4ac49ff124..2e806cc426 100644 --- a/tests/integrations/openai/test_openai.py +++ b/tests/integrations/openai/test_openai.py @@ -928,7 +928,6 @@ async def test_bad_chat_completion_async(sentry_init, capture_events): assert event["level"] == "error" -@pytest.mark.asyncio @pytest.mark.parametrize( "send_default_pii, include_prompts", [