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..c77ba5c0b3 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 @@ -208,43 +208,49 @@ def _set_output_messages(span, choices): @dont_throw 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: + for item in response: + yield item + _accumulate_streaming_response(complete_response, item) - _set_response_attributes(span, complete_response) + _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)) + finally: + if span.is_recording(): + span.end() @dont_throw 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: + async for item in response: + yield item + _accumulate_streaming_response(complete_response, item) - _set_response_attributes(span, complete_response) + _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)) + finally: + if span.is_recording(): + span.end() def _emit_streaming_response_events(complete_response): diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py index 0574de076e..be44fe4a35 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py @@ -1,3 +1,4 @@ +import gc from unittest.mock import patch import httpx @@ -393,6 +394,55 @@ def test_completion_streaming( ), "Assert that it doesn't emit logs when use_legacy_attributes is True" +@pytest.mark.vcr +@pytest.mark.default_cassette("test_completion_streaming.yaml") +def test_completion_streaming_early_close( + instrument_legacy, span_exporter, log_exporter, mock_openai_client +): + response = mock_openai_client.completions.create( + model="davinci-002", + prompt="Tell me a joke about opentelemetry", + stream=True, + ) + + next(response) + response.close() + + spans = span_exporter.get_finished_spans() + assert [span.name for span in spans] == [ + "openai.completion", + ] + + +@pytest.mark.vcr +@pytest.mark.default_cassette("test_completion_streaming.yaml") +def test_completion_streaming_exception_during_consumption( + instrument_legacy, span_exporter, log_exporter, mock_openai_client +): + """The span is still ended when the caller aborts consumption by raising.""" + + response = mock_openai_client.completions.create( + model="davinci-002", + prompt="Tell me a joke about opentelemetry", + stream=True, + ) + + with pytest.raises(ValueError, match="simulated interruption"): + for _ in response: + raise ValueError("simulated interruption") + + del response + gc.collect() + + spans = span_exporter.get_finished_spans() + assert [span.name for span in spans] == [ + "openai.completion", + ] + # Status is deliberately not asserted: what an incompletely observed + # completion should report is a separate question from ending the span. + assert spans[0].end_time is not None + + @pytest.mark.vcr def test_completion_streaming_with_events_with_content( instrument_with_content, span_exporter, log_exporter, openai_client @@ -548,6 +598,27 @@ async def test_async_completion_streaming( ), "Assert that it doesn't emit logs when use_legacy_attributes is True" +@pytest.mark.vcr +@pytest.mark.default_cassette("test_async_completion_streaming.yaml") +@pytest.mark.asyncio +async def test_async_completion_streaming_early_close( + instrument_legacy, span_exporter, log_exporter, async_openai_client +): + response = await async_openai_client.completions.create( + model="davinci-002", + prompt="Tell me a joke about opentelemetry", + stream=True, + ) + + await anext(response) + await response.aclose() + + spans = span_exporter.get_finished_spans() + assert [span.name for span in spans] == [ + "openai.completion", + ] + + @pytest.mark.vcr @pytest.mark.asyncio async def test_async_completion_streaming_with_events_with_content(