diff --git a/sentry_sdk/integrations/openai.py b/sentry_sdk/integrations/openai.py index 742b693245..699605166a 100644 --- a/sentry_sdk/integrations/openai.py +++ b/sentry_sdk/integrations/openai.py @@ -148,6 +148,11 @@ def setup_once() -> None: Responses.create = _wrap_responses_create(Responses.create) # type: ignore[method-assign] AsyncResponses.create = _wrap_async_responses_create(AsyncResponses.create) # type: ignore[assignment,method-assign] + # Finishing the span on close() covers streams that are closed early, + # abandoned mid-iteration, or exited via `with` before exhaustion. + Stream.close = _wrap_close(Stream.close) # type: ignore[assignment,method-assign] + AsyncStream.close = _wrap_async_close(AsyncStream.close) # type: ignore[assignment,method-assign] + def count_tokens(self: "OpenAIIntegration", s: str) -> int: if self.tiktoken_encoding is None: return 0 @@ -885,6 +890,8 @@ def _new_sync_chat_completion( old_iterator=response._iterator, finish_span=True, ) + # Attach the span so the patched Stream.close() can finish it on early exit. + response._span = span # type: ignore[attr-defined] else: _set_completions_api_output_data( @@ -969,6 +976,8 @@ async def _new_async_chat_completion( old_iterator=response._iterator, finish_span=True, ) + # Attach the span so the patched Stream.close() can finish it on early exit. + response._span = span # type: ignore[attr-defined] else: _set_completions_api_output_data( span, response, kwargs, integration, finish_span=True @@ -998,6 +1007,59 @@ def _set_completions_api_output_data( ) +def _finish_stream_span( + stream: "Union[Stream[Any], AsyncStream[Any]]", + exc_info: "tuple[Any, Any, Any]", +) -> None: + """ + Finishes the AI Client Span attached to a stream, if any, and detaches it + so the span is finished exactly once (close() and iterator finalization + both route through here). + """ + span = getattr(stream, "_span", None) + if span is None: + return + with capture_internal_exceptions(): + del stream._span # type: ignore[union-attr] + span.__exit__(*exc_info) + + +def _wrap_close( + f: "Callable[[Stream[Any]], None]", +) -> "Callable[[Stream[Any]], None]": + """ + Finishes the stream's AI Client Span when the stream is closed early, + unless the chunk iterator already finished it. + """ + + @wraps(f) + def close(self: "Stream[Any]") -> None: + try: + return f(self) + finally: + _finish_stream_span(self, sys.exc_info()) + + return close + + +def _wrap_async_close( + f: "Callable[[AsyncStream[Any]], Coroutine[Any, Any, None]]", +) -> "Callable[[AsyncStream[Any]], Coroutine[Any, Any, None]]": + """ + Finishes the stream's AI Client Span when the stream is closed early, + unless the chunk iterator already finished it. + """ + + @wraps(f) + async def close(self: "AsyncStream[Any]") -> None: + try: + return await f(self) + finally: + _finish_stream_span(self, sys.exc_info()) + + return close + + def _wrap_synchronous_completions_chunk_iterator( span: "Union[Span, StreamedSpan]", integration: "OpenAIIntegration", @@ -1017,56 +1079,64 @@ def _wrap_synchronous_completions_chunk_iterator( streaming_message_total_token_usage = None client = sentry_sdk.get_client() - for x in old_iterator: - if isinstance(span, StreamedSpan): - span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model) - else: - span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model) + try: + for x in old_iterator: + if isinstance(span, StreamedSpan): + span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model) + else: + span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model) + + with capture_internal_exceptions(): + if hasattr(x, "choices") and x.choices is not None: + choice_index = 0 + for choice in x.choices: + if hasattr(choice, "delta") and hasattr( + choice.delta, "content" + ): + if start_time is not None and ttft is None: + ttft = time.perf_counter() - start_time + content = choice.delta.content + if len(data_buf) <= choice_index: + data_buf.append([]) + data_buf[choice_index].append(content or "") + choice_index += 1 + if hasattr(x, "usage"): + streaming_message_total_token_usage = x.usage + + yield x with capture_internal_exceptions(): - if hasattr(x, "choices") and x.choices is not None: - choice_index = 0 - for choice in x.choices: - if hasattr(choice, "delta") and hasattr(choice.delta, "content"): - if start_time is not None and ttft is None: - ttft = time.perf_counter() - start_time - content = choice.delta.content - if len(data_buf) <= choice_index: - data_buf.append([]) - data_buf[choice_index].append(content or "") - choice_index += 1 - if hasattr(x, "usage"): - streaming_message_total_token_usage = x.usage - - yield x - - with capture_internal_exceptions(): - if ttft is not None: - set_data_normalized( - span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft - ) - all_responses = None - if len(data_buf) > 0: - all_responses = ["".join(chunk) for chunk in data_buf] - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["outputs"]: + if ttft is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft + ) + all_responses = None + if len(data_buf) > 0: + 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 + ) + elif should_send_default_pii() and integration.include_prompts: 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( - messages=messages, - response=response, - span=span, - streaming_message_responses=all_responses, - streaming_message_total_token_usage=streaming_message_total_token_usage, - count_tokens=integration.count_tokens, - ) - - if finish_span: - span.__exit__(None, None, None) + _calculate_completions_token_usage( + messages=messages, + response=response, + span=span, + streaming_message_responses=all_responses, + streaming_message_total_token_usage=streaming_message_total_token_usage, + count_tokens=integration.count_tokens, + ) + finally: + # Finish the span even if the stream is closed early, abandoned, or + # the connection drops mid-stream. close() detaches the span first, + # making this a no-op if it already ran there. + if finish_span: + _finish_stream_span(response, sys.exc_info()) async def _wrap_asynchronous_completions_chunk_iterator( @@ -1088,56 +1158,64 @@ async def _wrap_asynchronous_completions_chunk_iterator( streaming_message_total_token_usage = None client = sentry_sdk.get_client() - async for x in old_iterator: - if isinstance(span, StreamedSpan): - span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model) - else: - span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model) + try: + async for x in old_iterator: + if isinstance(span, StreamedSpan): + span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model) + else: + span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.model) + + with capture_internal_exceptions(): + if hasattr(x, "choices") and x.choices is not None: + choice_index = 0 + for choice in x.choices: + if hasattr(choice, "delta") and hasattr( + choice.delta, "content" + ): + if start_time is not None and ttft is None: + ttft = time.perf_counter() - start_time + content = choice.delta.content + if len(data_buf) <= choice_index: + data_buf.append([]) + data_buf[choice_index].append(content or "") + choice_index += 1 + if hasattr(x, "usage"): + streaming_message_total_token_usage = x.usage + + yield x with capture_internal_exceptions(): - if hasattr(x, "choices") and x.choices is not None: - choice_index = 0 - for choice in x.choices: - if hasattr(choice, "delta") and hasattr(choice.delta, "content"): - if start_time is not None and ttft is None: - ttft = time.perf_counter() - start_time - content = choice.delta.content - if len(data_buf) <= choice_index: - data_buf.append([]) - data_buf[choice_index].append(content or "") - choice_index += 1 - if hasattr(x, "usage"): - streaming_message_total_token_usage = x.usage - - yield x - - with capture_internal_exceptions(): - if ttft is not None: - set_data_normalized( - span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft - ) - all_responses = None - if len(data_buf) > 0: - all_responses = ["".join(chunk) for chunk in data_buf] - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["outputs"]: + if ttft is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft + ) + all_responses = None + if len(data_buf) > 0: + 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 + ) + elif should_send_default_pii() and integration.include_prompts: 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( - messages=messages, - response=response, - span=span, - streaming_message_responses=all_responses, - streaming_message_total_token_usage=streaming_message_total_token_usage, - count_tokens=integration.count_tokens, - ) - - if finish_span: - span.__exit__(None, None, None) + _calculate_completions_token_usage( + messages=messages, + response=response, + span=span, + streaming_message_responses=all_responses, + streaming_message_total_token_usage=streaming_message_total_token_usage, + count_tokens=integration.count_tokens, + ) + finally: + # Finish the span even if the stream is closed early, abandoned, or + # the connection drops mid-stream. close() detaches the span first, + # making this a no-op if it already ran there. + if finish_span: + _finish_stream_span(response, sys.exc_info()) def _wrap_synchronous_responses_event_iterator( @@ -1159,58 +1237,66 @@ def _wrap_synchronous_responses_event_iterator( client = sentry_sdk.get_client() count_tokens_manually = True - for x in old_iterator: - with capture_internal_exceptions(): - if isinstance(x, ResponseTextDeltaEvent): - if start_time is not None and ttft is None: - ttft = time.perf_counter() - start_time - if len(data_buf) == 0: - data_buf.append([]) - data_buf[0].append(x.delta or "") - - elif isinstance(x, ResponseCompletedEvent): - if isinstance(span, StreamedSpan): - span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model) - else: - span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model) - - _calculate_responses_token_usage( - input=input, - response=x.response, - span=span, - streaming_message_responses=None, - count_tokens=integration.count_tokens, - ) - count_tokens_manually = False + try: + for x in old_iterator: + with capture_internal_exceptions(): + if isinstance(x, ResponseTextDeltaEvent): + if start_time is not None and ttft is None: + ttft = time.perf_counter() - start_time + if len(data_buf) == 0: + data_buf.append([]) + data_buf[0].append(x.delta or "") + + elif isinstance(x, ResponseCompletedEvent): + if isinstance(span, StreamedSpan): + span.set_attribute( + SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model + ) + else: + span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model) + + _calculate_responses_token_usage( + input=input, + response=x.response, + span=span, + streaming_message_responses=None, + count_tokens=integration.count_tokens, + ) + count_tokens_manually = False - yield x + yield x - with capture_internal_exceptions(): - if ttft is not None: - set_data_normalized( - span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft - ) - if len(data_buf) > 0: - all_responses = ["".join(chunk) for chunk in data_buf] - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["outputs"]: + with capture_internal_exceptions(): + if ttft is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft + ) + if len(data_buf) > 0: + 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 + ) + elif should_send_default_pii() and integration.include_prompts: 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, - response=response, - span=span, - streaming_message_responses=all_responses, - count_tokens=integration.count_tokens, - ) - if finish_span: - span.__exit__(None, None, None) + if count_tokens_manually: + _calculate_responses_token_usage( + input=input, + response=response, + span=span, + streaming_message_responses=all_responses, + count_tokens=integration.count_tokens, + ) + finally: + # Finish the span even if the stream is closed early, abandoned, or + # the connection drops mid-stream. close() detaches the span first, + # making this a no-op if it already ran there. + if finish_span: + _finish_stream_span(response, sys.exc_info()) async def _wrap_asynchronous_responses_event_iterator( @@ -1232,58 +1318,67 @@ async def _wrap_asynchronous_responses_event_iterator( client = sentry_sdk.get_client() count_tokens_manually = True - async for x in old_iterator: - with capture_internal_exceptions(): - if isinstance(x, ResponseTextDeltaEvent): - if start_time is not None and ttft is None: - ttft = time.perf_counter() - start_time - if len(data_buf) == 0: - data_buf.append([]) - data_buf[0].append(x.delta or "") - - elif isinstance(x, ResponseCompletedEvent): - if isinstance(span, StreamedSpan): - span.set_attribute(SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model) - else: - span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model) - - _calculate_responses_token_usage( - input=input, - response=x.response, - span=span, - streaming_message_responses=None, - count_tokens=integration.count_tokens, - ) - count_tokens_manually = False + try: + async for x in old_iterator: + with capture_internal_exceptions(): + if isinstance(x, ResponseTextDeltaEvent): + if start_time is not None and ttft is None: + ttft = time.perf_counter() - start_time + if len(data_buf) == 0: + data_buf.append([]) + data_buf[0].append(x.delta or "") + + elif isinstance(x, ResponseCompletedEvent): + if isinstance(span, StreamedSpan): + span.set_attribute( + SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model + ) + else: + span.set_data(SPANDATA.GEN_AI_RESPONSE_MODEL, x.response.model) + + _calculate_responses_token_usage( + input=input, + response=x.response, + span=span, + streaming_message_responses=None, + count_tokens=integration.count_tokens, + ) + count_tokens_manually = False - yield x + yield x - with capture_internal_exceptions(): - if ttft is not None: - set_data_normalized( - span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft - ) - if len(data_buf) > 0: - all_responses = ["".join(chunk) for chunk in data_buf] - - if has_data_collection_enabled(client.options): - if client.options["data_collection"]["gen_ai"]["outputs"]: + with capture_internal_exceptions(): + if ttft is not None: + set_data_normalized( + span, SPANDATA.GEN_AI_RESPONSE_TIME_TO_FIRST_TOKEN, ttft + ) + if len(data_buf) > 0: + 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 + ) + elif should_send_default_pii() and integration.include_prompts: 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, - response=response, - span=span, - streaming_message_responses=all_responses, - count_tokens=integration.count_tokens, - ) - if finish_span: - span.__exit__(None, None, None) + + if count_tokens_manually: + _calculate_responses_token_usage( + input=input, + response=response, + span=span, + streaming_message_responses=all_responses, + count_tokens=integration.count_tokens, + ) + finally: + # Finish the span even if the stream is closed early, abandoned, or + # the connection drops mid-stream. close() detaches the span first, + # making this a no-op if it already ran there. + if finish_span: + _finish_stream_span(response, sys.exc_info()) def _set_responses_api_output_data( @@ -1567,6 +1662,8 @@ def _new_sync_responses_create( old_iterator=response._iterator, finish_span=True, ) + # Attach the span so the patched Stream.close() can finish it on early exit. + response._span = span # type: ignore[attr-defined] else: _set_responses_api_output_data( @@ -1643,6 +1740,8 @@ async def _new_async_responses_create( old_iterator=response._iterator, finish_span=True, ) + # Attach the span so the patched Stream.close() can finish it on early exit. + response._span = span # type: ignore[attr-defined] else: _set_responses_api_output_data( span, response, kwargs, integration, finish_span=True diff --git a/tests/integrations/openai/test_openai.py b/tests/integrations/openai/test_openai.py index 96866886ec..80d1dc022a 100644 --- a/tests/integrations/openai/test_openai.py +++ b/tests/integrations/openai/test_openai.py @@ -2568,6 +2568,86 @@ def test_streaming_chat_completion( pass # if tiktoken is not installed, we can't guarantee token usage will be calculated properly +@pytest.mark.parametrize("consume", ["break_then_close", "with_block"]) +def test_streaming_chat_completion_span_finished_on_early_exit( + sentry_init, + capture_events, + get_model_response, + server_side_event_chunks, + consume, +): + # https://github.com/getsentry/sentry-python/issues/7847 + # Closing a stream early (or leaving a `with` block after one chunk) must + # still finish the gen_ai.chat span. + sentry_init( + integrations=[OpenAIIntegration()], + disabled_integrations=[StdlibIntegration], + traces_sample_rate=1.0, + # Keep gen_ai spans on the transaction so the assertion below sees them. + stream_gen_ai_spans=False, + ) + + 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=0, + delta=ChoiceDelta(content="lo"), + finish_reason=None, + ) + ], + created=100000, + model="model-id", + object="chat.completion.chunk", + ), + ], + include_event_type=False, + ) + ) + + with mock.patch.object( + client.chat._client._client, + "send", + return_value=returned_stream, + ): + events = capture_events() + with start_transaction(name="openai tx"): + stream = client.chat.completions.create( + model="some-model", + messages=[{"role": "user", "content": "hello"}], + stream=True, + ) + if consume == "break_then_close": + for _ in stream: + break + stream.close() + else: + with stream: + next(iter(stream)) + + transactions = [event for event in events if event.get("type") == "transaction"] + assert len(transactions) == 1 + span_ops = [span["op"] for span in transactions[0]["spans"]] + assert "gen_ai.chat" in span_ops + + # noinspection PyTypeChecker @pytest.mark.parametrize("span_streaming", [True, False]) @pytest.mark.parametrize("stream_gen_ai_spans", [True, False])