Skip to content

fix(openai): end the LLM span when a stream is abandoned or closed early - #3796

Open
feiiiiii5 wants to merge 1 commit into
Arize-ai:mainfrom
feiiiiii5:fix/openai-stream-teardown-oi
Open

feiiiiii5 wants to merge 1 commit into
Arize-ai:mainfrom
feiiiiii5:fix/openai-stream-teardown-oi

Conversation

@feiiiiii5

Copy link
Copy Markdown
Contributor

What happens

A streamed chat completion that the caller stops consuming early is missing from the trace
entirely
. _Stream ends the LLM span only when iteration raises StopIteration /
StopAsyncIteration, or when an error propagates out of __next__. The ordinary ways to stop
early — leaving the with block, calling close() / aclose(), or dropping the object — hand
control back without ever finishing tracing. An unended span is never exported, so a call that
really happened, really cost tokens and may have already streamed a partial answer to the user
leaves no record.

with client.chat.completions.create(model=..., messages=..., stream=True) as stream:
    for chunk in stream:
        break          # the span is never ended

langchain_openai is a live consumer of the first path: it wraps streaming in with blocks and
breaks out early.

Exported spans for that snippet

CONTRIBUTING asks for trace output when a change affects emitted spans. This is reproduced
offline against a stubbed HTTP backend, so the evidence below is the exported span rendered from an
InMemorySpanExporter rather than a Phoenix screenshot; one script, same environment, both trees:

unmodified main : exported LLM spans after 'break' inside `with`: 0
this branch     : exported LLM spans after 'break' inside `with`: 1
                    name='ChatCompletion' kind=LLM status=UNSET
                    input.value='{"messages": [{"role": "user", "content": "hi"}]'

Root cause

In python/instrumentation/openinference-instrumentation-openai/src/openinference/instrumentation/openai/_stream.py,
__exit__ and __aexit__ call the wrapped teardown and return. close() is not overridden, so
ObjectProxy forwards it straight to the SDK object, and on openai >= 3.x AsyncStream.aclose()
calls the wrapped close(), which bypasses the wrapper too. None of those paths reaches
_finish_tracing.

Change

Finish the tracing on every teardown path, leaving the status UNSET so a truncated stream stays
distinguishable from a completed one (OK) or a failed one (ERROR, already recorded by
__next__):

  • __exit__ / __aexit__: try/finally, so the span still ends if the wrapped teardown raises.
  • close(): Stream.close() is synchronous while AsyncStream.close() is a coroutine; the
    awaited case finishes once the wrapped close completes.
  • aclose(): covers the SDK method that bypasses close() on the proxy.
  • __del__: last resort for a stream dropped without any explicit teardown; it suppresses
    exceptions so finalization can never raise into user code.

close / aclose / __del__ use the getattr(...) + callable() guard and a finally-based
finish taken from openinference-instrumentation-ollama/src/openinference/instrumentation/ollama/_stream.py,
which already ends abandoned streams with status UNSET and has a committed test for it
(test_chat_stream_abandoned_before_iteration). For contrast: together and cohere finish on
__exit__ with OK, and anthropic and mistralai have no teardown handling. This PR follows
the ollama convention, because UNSET keeps "the caller stopped early" distinguishable from
"the answer arrived".

_WithSpan.finish_tracing in _with_span.py returns as soon as its _is_finished flag is set, so
the added calls are no-ops after a normal completion: no span is ended twice, and a completed
stream keeps its OK status. The last test in the new file pins that.

Tests

tests/openinference/instrumentation/openai/test_stream_lifecycle.py covers the five paths plus a
control asserting that a fully consumed stream still yields exactly one OK span. No network and no
credentials: the SSE body is served through respx, which the package's test-requirements.txt
already pins.

command (same test file, one venv per tree, base = a719562e) unmodified main this branch
pytest tests/openinference/instrumentation/openai/test_stream_lifecycle.py with openai==2.8.0 (the ci-openai factor) 5 failed, 2 passed, 1 skipped 7 passed, 1 skipped
same file with the latest openai (3.16.2, the openai-latest factor) 6 failed, 2 passed 8 passed
pytest tests (whole package, pinned deps) 6 failed, 486 passed, 1 skipped 1 failed, 491 passed, 1 skipped
ruff format --diff . / ruff check --no-fix . (0.9.2, the pinned version) clean / pass clean / pass
mypy . (1.11.2, strict, as tox runs it) pass, 31 files pass, 31 files

The aclose case is skipped under openai==2.8.0 because that release has no
AsyncStream.aclose; it runs on openai-latest.

Each hunk carries its weight. Reverting one at a time (latest openai, same file) fails exactly the
tests that claim it:

mutation tests that then fail
__exit__ back to a plain delegate test_leaving_context_without_exhausting_ends_one_span[False], [True]
__aexit__ back to a plain delegate test_leaving_context_without_exhausting_ends_one_span[True], test_closing_without_exhausting_ends_one_span[True]
close() removed test_closing_without_exhausting_ends_one_span[False], [True], test_aclosing_without_exhausting_ends_one_span, test_stream_never_iterated_ends_one_span
aclose() removed test_aclosing_without_exhausting_ends_one_span, test_stream_never_iterated_ends_one_span
__del__ removed test_stream_never_iterated_ends_one_span

Every test holds a reference to the stream until after its own assertions, so a span can only be
finished by the path the test exercises and never by a later __del__ during collection garbage.

Environment: macOS arm64, Python 3.11.15; dependencies installed from the package's own
test-requirements.txt for the pinned column and uv pip install -U openai for the latest column.

A whole-package failure that predates this change

test_tool_calls.py::test_tool_calls fails in a whole-package run on unmodified main with the
same command and dependency set (1 failed, 484 passed, measured with no new test file present in
the tree), passes when the file is run alone, and shows the same assert 2 == 1 at
test_tool_calls.py:107 on this branch. It looks order-dependent through the session-scoped
exporter. I did not investigate it and this PR does not change it; the numbers above are reported
as measured so the two are not confused.

Scope

Closes #3790. That issue asked two questions. On status: this uses UNSET, reserving ERROR for
exceptions raised through iteration. On the __del__ fallback: included, because the ollama
instrumentor already ships the same fallback in main with a committed test, so this makes the two
packages behave alike rather than adding a new idea — if you would rather keep this PR to the
explicit teardown paths, say so and I will drop that method and its test.

#3784 is the same class of defect in the cohere instrumentor and is not addressed here.

`_Stream` finished the span only when iteration raised StopIteration /
StopAsyncIteration, or when an error propagated out of `__next__`. A caller that leaves
the `with` block early, calls `close()` / `aclose()`, or drops the stream object never
reaches those paths, so the span stayed open, was never exported, and the whole LLM
call disappeared from the trace.

End the tracing from `__exit__`, `__aexit__`, `close`, `aclose` and `__del__`, leaving
the status UNSET so a truncated stream stays distinguishable from a completed one --
the same shape `openinference-instrumentation-ollama` already uses for its streams.
`_WithSpan.finish_tracing` returns early once the span is finished, so these calls are
no-ops after a normal completion.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: No status

Development

Successfully merging this pull request may close these issues.

[bug] openai: abandoning a streamed chat completion never ends the LLM span

1 participant