fix(openai): finalize completion spans on early stream close - #4394
Open
CTWalk wants to merge 1 commit into
Open
fix(openai): finalize completion spans on early stream close#4394CTWalk wants to merge 1 commit into
CTWalk wants to merge 1 commit into
Conversation
The Completions streaming wrappers end their spans after the iteration loop. That works when the stream is exhausted, but close() or aclose() exits the generator at its suspended yield, before span.end() is reached, so an early-closed stream exports no openai.completion span at all. Wrap both streaming generator bodies in try/finally and end a still-recording span in finally, so exhaustion, close(), and aclose() converge on one terminal span.end(). Full-consumption response attributes, token usage, events, and StatusCode.OK still run only on normal exhaustion. Because finally also runs on exceptions, a stream that raises mid-iteration now ends and exports its span instead of leaving one recording and never exported. Add three tests covering sync early close, async early close, and an exception raised during consumption, the last mirroring the chat-side coverage added in traceloop#3155. All three reuse the existing streaming cassettes via pytest-recording's default_cassette marker, so no fixture data is added. Most of the diff is re-indentation from the new try block; the semantic change is six lines, visible with git diff -w. Related to traceloop#3151 and traceloop#3155, which established this contract for chat streams. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
|
📝 WalkthroughWalkthroughOpenAI streaming completion wrappers now end spans during normal completion, errors, and early interruption. Tests cover synchronous and asynchronous early close, plus synchronous consumption exceptions and garbage-collection cleanup. ChangesOpenAI streaming span finalization
Estimated code review effort: 3 (Moderate) | ~20 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py (2)
601-621: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winNo async counterpart for the consumption-exception test.
The PR adds a sync consumption-exception test (Lines 419-445) but no equivalent async version alongside this async early-close test. Since
_abuild_from_streaming_responsemirrors the sync implementation, an async exception-during-consumption test would close the coverage gap for that code path.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py` around lines 601 - 621, Add an async consumption-exception test alongside test_async_completion_streaming_early_close, using async_openai_client and the streaming completion response. Consume the response until an exception is raised, close it appropriately, and assert the finished span behavior matches the synchronous consumption-exception test, covering _abuild_from_streaming_response.
397-415: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winEarly-close test only checks span presence, not that success data is absent.
test_completion_streaming_early_closeconfirms the span is exported, but it does not assert that full-consumption-only data (response attributes, token usage,StatusCode.OK) is absent for an early-closed stream. Since that is the specific contract this PR establishes ("Early-closed streams retain collected request attributes without fabricated response data or success status"), a negative assertion here would directly pin down that contract for regression protection.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py` around lines 397 - 415, Extend test_completion_streaming_early_close to assert the exported span retains collected request attributes but has no response attributes or token-usage data and does not have StatusCode.OK after the stream is closed early. Keep the existing span-presence assertion and use the span’s established attribute and status accessors.packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py (1)
211-253: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate finalize logic between sync and async wrappers.
The post-loop sequence (
_set_response_attributes,_set_token_usage, event/completion emission,set_status(OK), and thefinallyguard) is duplicated verbatim between_build_from_streaming_responseand_abuild_from_streaming_response. Only the loop syntax (forvsasync for) differs.Extract the shared finalize step into a helper (e.g.,
_finalize_streaming_span(span, request_kwargs, complete_response)) called from both generators inside thetry. This also gives one place to add the stream-close fix suggested above.♻️ Sketch of the shared helper
+def _finalize_streaming_span(span, request_kwargs, complete_response): + _set_response_attributes(span, complete_response) + _set_token_usage(span, request_kwargs, complete_response) + if should_emit_events(): + _emit_streaming_response_events(complete_response) + elif should_send_prompts(): + _set_completions(span, complete_response.get("choices")) + span.set_status(Status(StatusCode.OK)) + + def _build_from_streaming_response(span, request_kwargs, response): complete_response = {"choices": [], "model": "", "id": ""} try: for item in response: yield item _accumulate_streaming_response(complete_response, item) - - _set_response_attributes(span, 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")) - - span.set_status(Status(StatusCode.OK)) + _finalize_streaming_span(span, request_kwargs, complete_response) finally: if span.is_recording(): span.end()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py` around lines 211 - 253, Extract the duplicated post-stream processing from _build_from_streaming_response and _abuild_from_streaming_response into a shared _finalize_streaming_span helper, including response attributes, token usage, event or completion emission, and the OK status. Call the helper after each sync or async loop within its respective try block, while retaining the existing finally span.is_recording() guard and span.end() behavior in both generators.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In
`@packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py`:
- Around line 209-230: Update both _build_from_streaming_response and
_abuild_from_streaming_response to close the raw response in their cleanup paths
before ending the span: call response.close() for synchronous streams and
guarded response.aclose() for asynchronous streams, awaiting it in the async
wrapper when available. Preserve span.end() execution even if stream cleanup is
unavailable or fails.
---
Nitpick comments:
In
`@packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.py`:
- Around line 211-253: Extract the duplicated post-stream processing from
_build_from_streaming_response and _abuild_from_streaming_response into a shared
_finalize_streaming_span helper, including response attributes, token usage,
event or completion emission, and the OK status. Call the helper after each sync
or async loop within its respective try block, while retaining the existing
finally span.is_recording() guard and span.end() behavior in both generators.
In
`@packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py`:
- Around line 601-621: Add an async consumption-exception test alongside
test_async_completion_streaming_early_close, using async_openai_client and the
streaming completion response. Consume the response until an exception is
raised, close it appropriately, and assert the finished span behavior matches
the synchronous consumption-exception test, covering
_abuild_from_streaming_response.
- Around line 397-415: Extend test_completion_streaming_early_close to assert
the exported span retains collected request attributes but has no response
attributes or token-usage data and does not have StatusCode.OK after the stream
is closed early. Keep the existing span-presence assertion and use the span’s
established attribute and status accessors.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 04afb729-1298-4a21-b0c0-dd00f3b471e5
📒 Files selected for processing (2)
packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/completion_wrappers.pypackages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Finalize
openai.completionspans when a caller closes a sync or asyncCompletions stream before consuming it fully.
This brings the Completions path in line with OpenLLMetry's existing cleanup
behavior for chat and Responses streams in ending the span; it deliberately
does not copy chat's
StatusCode.OK, as noted below. It also adds focused syncand async regression tests.
Related: #3151 and #3155.
Problem
The Completions streaming wrappers end their spans after the iteration loop.
That works when the stream is exhausted, but
close()oraclose()exits thegenerator at its suspended
yield, beforespan.end()is reached.Reproduced with
OpenAIandAsyncOpenAIclients against a local SSE serverand an in-memory span exporter on
main@93429cfb:openai.completionspansclose()aclose()The full-consumption controls show that the instrumentation and exporter are
working; the missing span is specific to early close.
A cancelled or timed-out completion is therefore invisible rather than
truncated: the request that was worth investigating is the one with no span at
all, and #3151 reports the same loss on the chat path as a bug.
Change
try/finallyand enda still-recording span in
finally, so exhaustion,close(), andaclose()converge on one terminal
span.end().events, and
StatusCode.OK— they still run only on normal exhaustion.during consumption — the last mirroring
test_chat_streaming_exception_during_consumptionfrom fix(openai): support for openai non-consumed streams #3155. All three reusethe existing streaming cassettes via
@pytest.mark.default_cassette, so nofixture data is added.
After the change, all four cases above export exactly one span.
Most of the diff in
completion_wrappers.pyis re-indentation from the newtryblock. The semantic change is six lines —git diff -wshows just thetry:,finally:, andif span.is_recording():on each of the two generators.An early-closed stream exports the request attributes already collected. It
does not fabricate response attributes, token usage, or a successful status for
a response that was not fully observed. This is the one place the change does
not follow chat:
ChatStreamreportsStatusCode.OKfor an abandoned stream,and asserting success for a response nobody finished reading seemed worth not
copying. Glad to align it either way if you'd rather the two paths match
exactly.
Because
finallyalso runs on exceptions, the same change covers a stream thatraises mid-iteration: the span is now ended and exported instead of being left
recording and never exported. Measured on
main@93429cfbwith a stream thatraises after one chunk — 0 spans exported before, 1 after, with the exception
propagating unchanged in both. The new consumption-exception test pins the
ending; it deliberately does not assert a status, since which status such a span
should carry is a semantic-convention question rather than a lifecycle one.
Verification
From
packages/opentelemetry-instrumentation-openai:All three new tests fail on the unpatched tree because no finished
openai.completionspan is exported.The 26 includes the two pre-existing invalid-key tests, which carry no cassette
and need a reachable endpoint to see the 401 they assert. They are unrelated to
this change; without network access they fail on both the patched and unpatched
tree alike.
Checklist
screenshots from an observability platform.
Not included: this lifecycle behavior is asserted directly with the
repository's in-memory exporter and does not change an observability UI.
Not applicable: there is no public API, configuration, or documented
output change.
AI assistance was used during investigation and draft preparation. The
resulting patch was reviewed against the source, reproduced on both paths, and
verified with the checks above.
Summary by CodeRabbit