diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py index 7c805d6995..129e1d9350 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py @@ -931,56 +931,69 @@ def _build_from_streaming_response( time_of_first_token = start_time # will be updated when first token is received for item in response: - span.add_event(name=SpanAttributes.GEN_AI_CONTENT_COMPLETION_CHUNK) - item_to_yield = item - if first_token and streaming_time_to_first_token: - time_of_first_token = time.time() - streaming_time_to_first_token.record( - time_of_first_token - start_time) - first_token = False - _accumulate_stream_items(item, complete_response) + try: + span.add_event(name=SpanAttributes.GEN_AI_CONTENT_COMPLETION_CHUNK) - yield item_to_yield + if first_token and streaming_time_to_first_token: + time_of_first_token = time.time() + streaming_time_to_first_token.record( + time_of_first_token - start_time) + first_token = False - shared_attributes = { - GenAIAttributes.GEN_AI_RESPONSE_MODEL: complete_response.get("model") or None, - "server.address": _get_openai_base_url(instance), - "stream": True, - } + _accumulate_stream_items(item, complete_response) + except Exception as e: + logging.warning( + "OpenLLMetry failed to trace a streaming chunk, continuing without it: %s", + e, + ) - _set_streaming_token_metrics( - request_kwargs, complete_response, span, token_counter, shared_attributes - ) + yield item_to_yield - # choice metrics - if choice_counter and complete_response.get("choices"): - _set_choice_counter_metrics( - choice_counter, complete_response.get("choices"), shared_attributes + try: + shared_attributes = { + GenAIAttributes.GEN_AI_RESPONSE_MODEL: complete_response.get("model") or None, + "server.address": _get_openai_base_url(instance), + "stream": True, + } + + _set_streaming_token_metrics( + request_kwargs, complete_response, span, token_counter, shared_attributes ) - # duration metrics - if start_time and isinstance(start_time, (float, int)): - duration = time.time() - start_time - else: - duration = None - if duration and isinstance(duration, (float, int)) and duration_histogram: - duration_histogram.record(duration, attributes=shared_attributes) - if streaming_time_to_generate and time_of_first_token: - streaming_time_to_generate.record(time.time() - time_of_first_token) + # choice metrics + if choice_counter and complete_response.get("choices"): + _set_choice_counter_metrics( + choice_counter, complete_response.get("choices"), shared_attributes + ) - _set_response_attributes(span, complete_response) - if should_emit_events(): - for choice in complete_response.get("choices", []): - emit_event(_parse_choice_event(choice)) - else: - if should_send_prompts(): - _set_completions(span, complete_response.get("choices")) + # duration metrics + if start_time and isinstance(start_time, (float, int)): + duration = time.time() - start_time + else: + duration = None + if duration and isinstance(duration, (float, int)) and duration_histogram: + duration_histogram.record(duration, attributes=shared_attributes) + if streaming_time_to_generate and time_of_first_token: + streaming_time_to_generate.record(time.time() - time_of_first_token) - span.set_status(Status(StatusCode.OK)) - span.end() + _set_response_attributes(span, complete_response) + if should_emit_events(): + for choice in complete_response.get("choices", []): + emit_event(_parse_choice_event(choice)) + else: + if should_send_prompts(): + _set_completions(span, complete_response.get("choices")) + + span.set_status(Status(StatusCode.OK)) + except Exception as e: + logging.warning( + "OpenLLMetry failed to finalize the streaming span: %s", e + ) + finally: + span.end() @dont_throw @@ -1002,56 +1015,69 @@ async def _abuild_from_streaming_response( time_of_first_token = start_time # will be updated when first token is received async for item in response: - span.add_event(name=SpanAttributes.GEN_AI_CONTENT_COMPLETION_CHUNK) - item_to_yield = item - if first_token and streaming_time_to_first_token: - time_of_first_token = time.time() - streaming_time_to_first_token.record( - time_of_first_token - start_time) - first_token = False - _accumulate_stream_items(item, complete_response) + try: + span.add_event(name=SpanAttributes.GEN_AI_CONTENT_COMPLETION_CHUNK) - yield item_to_yield + if first_token and streaming_time_to_first_token: + time_of_first_token = time.time() + streaming_time_to_first_token.record( + time_of_first_token - start_time) + first_token = False - shared_attributes = { - GenAIAttributes.GEN_AI_RESPONSE_MODEL: complete_response.get("model") or None, - "server.address": _get_openai_base_url(instance), - "stream": True, - } + _accumulate_stream_items(item, complete_response) + except Exception as e: + logging.warning( + "OpenLLMetry failed to trace a streaming chunk, continuing without it: %s", + e, + ) - _set_streaming_token_metrics( - request_kwargs, complete_response, span, token_counter, shared_attributes - ) + yield item_to_yield - # choice metrics - if choice_counter and complete_response.get("choices"): - _set_choice_counter_metrics( - choice_counter, complete_response.get("choices"), shared_attributes + try: + shared_attributes = { + GenAIAttributes.GEN_AI_RESPONSE_MODEL: complete_response.get("model") or None, + "server.address": _get_openai_base_url(instance), + "stream": True, + } + + _set_streaming_token_metrics( + request_kwargs, complete_response, span, token_counter, shared_attributes ) - # duration metrics - if start_time and isinstance(start_time, (float, int)): - duration = time.time() - start_time - else: - duration = None - if duration and isinstance(duration, (float, int)) and duration_histogram: - duration_histogram.record(duration, attributes=shared_attributes) - if streaming_time_to_generate and time_of_first_token: - streaming_time_to_generate.record(time.time() - time_of_first_token) + # choice metrics + if choice_counter and complete_response.get("choices"): + _set_choice_counter_metrics( + choice_counter, complete_response.get("choices"), shared_attributes + ) - _set_response_attributes(span, complete_response) - if should_emit_events(): - for choice in complete_response.get("choices", []): - emit_event(_parse_choice_event(choice)) - else: - if should_send_prompts(): - _set_completions(span, complete_response.get("choices")) + # duration metrics + if start_time and isinstance(start_time, (float, int)): + duration = time.time() - start_time + else: + duration = None + if duration and isinstance(duration, (float, int)) and duration_histogram: + duration_histogram.record(duration, attributes=shared_attributes) + if streaming_time_to_generate and time_of_first_token: + streaming_time_to_generate.record(time.time() - time_of_first_token) - span.set_status(Status(StatusCode.OK)) - span.end() + _set_response_attributes(span, complete_response) + if should_emit_events(): + for choice in complete_response.get("choices", []): + emit_event(_parse_choice_event(choice)) + else: + if should_send_prompts(): + _set_completions(span, complete_response.get("choices")) + + span.set_status(Status(StatusCode.OK)) + except Exception as e: + logging.warning( + "OpenLLMetry failed to finalize the streaming span: %s", e + ) + finally: + span.end() # pydantic.BaseModel here is ChatCompletionMessageFunctionToolCall (as of openai 1.99.7) diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py index b4461cd841..e6d6db424f 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py @@ -210,20 +210,33 @@ def _build_from_streaming_response(span, request_kwargs, response): complete_response = {"choices": [], "model": "", "id": ""} for item in response: yield item - _accumulate_streaming_response(complete_response, item) + + try: + _accumulate_streaming_response(complete_response, item) + except Exception as e: + logging.warning( + "OpenLLMetry failed to trace a streaming chunk, continuing without it: %s", + e, + ) - _set_response_attributes(span, complete_response) + try: + _set_response_attributes(span, complete_response) - _set_token_usage(span, request_kwargs, complete_response) + _set_token_usage(span, request_kwargs, complete_response) - if should_emit_events(): - _emit_streaming_response_events(complete_response) - else: - if should_send_prompts(): - _set_completions(span, complete_response.get("choices")) + if should_emit_events(): + _emit_streaming_response_events(complete_response) + else: + if should_send_prompts(): + _set_completions(span, complete_response.get("choices")) - span.set_status(Status(StatusCode.OK)) - span.end() + span.set_status(Status(StatusCode.OK)) + except Exception as e: + logging.warning( + "OpenLLMetry failed to finalize the streaming span: %s", e + ) + finally: + span.end() @dont_throw @@ -231,20 +244,32 @@ async def _abuild_from_streaming_response(span, request_kwargs, response): complete_response = {"choices": [], "model": "", "id": ""} async for item in response: yield item - _accumulate_streaming_response(complete_response, item) + try: + _accumulate_streaming_response(complete_response, item) + except Exception as e: + logging.warning( + "OpenLLMetry failed to trace a streaming chunk, continuing without it: %s", + e, + ) - _set_response_attributes(span, complete_response) + try: + _set_response_attributes(span, complete_response) - _set_token_usage(span, request_kwargs, complete_response) + _set_token_usage(span, request_kwargs, complete_response) - if should_emit_events(): - _emit_streaming_response_events(complete_response) - else: - if should_send_prompts(): - _set_completions(span, complete_response.get("choices")) + if should_emit_events(): + _emit_streaming_response_events(complete_response) + else: + if should_send_prompts(): + _set_completions(span, complete_response.get("choices")) - span.set_status(Status(StatusCode.OK)) - span.end() + span.set_status(Status(StatusCode.OK)) + except Exception as e: + logging.warning( + "OpenLLMetry failed to finalize the streaming span: %s", e + ) + finally: + span.end() def _emit_streaming_response_events(complete_response): diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_streaming_tracing_failure_safety.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_streaming_tracing_failure_safety.py new file mode 100644 index 0000000000..7fe7e0b224 --- /dev/null +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_streaming_tracing_failure_safety.py @@ -0,0 +1,65 @@ +""" +Regression test for the @dont_throw / generator gap. + +_build_from_streaming_response (and its async twin) are generator functions. +@dont_throw only wraps the instant they're *called*, which just creates a +paused generator object and runs none of the body -- it never protects the +`for item in response:` loop, which only actually executes later, when the +caller iterates via `for chunk in stream:`. Before this fix, a real bug in +the tracing bookkeeping (_accumulate_stream_items) during that loop was +uncaught and crashed straight into the caller's own streaming loop, silently +dropping every remaining chunk of the real response -- despite the function +being decorated with @dont_throw. + +This test forces exactly that: one chunk in the middle of the stream is +missing its "choices" key, which makes `for choice in item.get("choices"):` +inside _accumulate_stream_items raise TypeError (`NoneType not iterable`). +""" + +import logging +from unittest.mock import MagicMock + +from opentelemetry.instrumentation.openai.shared.chat_wrappers import ( + _build_from_streaming_response, +) + + +def test_tracing_failure_on_one_chunk_does_not_drop_later_chunks(monkeypatch, caplog): + # Skip the openai-v1 model_as_dict(...) conversion inside + # _accumulate_stream_items -- our fake chunks are plain dicts already. + monkeypatch.setattr( + "opentelemetry.instrumentation.openai.shared.chat_wrappers.is_openai_v1", + lambda: False, + ) + + good_chunk_1 = { + "id": "c1", + "model": "gpt-4", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "Hel"}}], + } + # Malformed: no "choices" key at all -> item.get("choices") is None -> + # `for choice in None:` raises TypeError inside _accumulate_stream_items. + malformed_chunk = {"id": "c2", "model": "gpt-4"} + good_chunk_3 = { + "id": "c3", + "model": "gpt-4", + "choices": [{"index": 0, "delta": {"content": "lo"}}], + } + + fake_response = [good_chunk_1, malformed_chunk, good_chunk_3] + fake_span = MagicMock() + + gen = _build_from_streaming_response(fake_span, fake_response) + + with caplog.at_level(logging.WARNING): + yielded = list(gen) + + # The actual point of the fix: nothing gets dropped just because tracing + # broke on one chunk. All three real chunks still reach the caller, in order. + assert yielded == [good_chunk_1, malformed_chunk, good_chunk_3] + + # The failure was logged, not silently swallowed and not left to crash. + assert any( + "failed to trace a streaming chunk" in record.getMessage() + for record in caplog.records + )