fix(openai): finalize completion spans on early stream close

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):

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(