Skip to content

fix(litellm): preserve stream types in streaming wrappers and isolate span attribute extraction errors - #3578

Open
sushantsimpplr wants to merge 1 commit into
Arize-ai:mainfrom
sushantsimpplr:fix/litellm-streaming-type-preservation
Open

fix(litellm): preserve stream types in streaming wrappers and isolate span attribute extraction errors#3578
sushantsimpplr wants to merge 1 commit into
Arize-ai:mainfrom
sushantsimpplr:fix/litellm-streaming-type-preservation

Conversation

@sushantsimpplr

Copy link
Copy Markdown

fixes #3577

Instrumentation must be transparent to callers of the instrumented functions. Three related defects in openinference-instrumentation-litellm break that, with litellm itself as the highest-impact caller — its Responses API → chat-completions bridge re-enters the instrumented litellm.acompletion and type-checks the result, so with this instrumentor active every streamed /v1/responses request to a provider without native Responses support fails with Unexpected response type: <class 'async_generator'>, and a secondary serialization crash masks the error (details in the issue).

Changes

  1. Type-preserving stream wrappers — new _TracedSyncStream / _TracedAsyncStream (wrapt.ObjectProxy, already a dependency) present as the original stream (isinstance passes, attribute access delegates) while iterating the existing finalizing generators, so span accumulation/finalization is unchanged. The four streaming return sites (completion, acompletion, responses, aresponses) now return the proxy instead of the bare generator.
  2. Exact-type gate for aresponses streaming — mirrors the sync wrapper's existing gate: streams that are not ResponsesAPIStreamingIterator pass through untouched (previously the finalizer silently drained them into an empty stream). Both responses wrappers now also end the span on the pass-through path (previously leaked).
  3. Exception-isolated attribute extraction_suppress_extractor_errors decorator on the _instrument_func_type_* extractors and _finalize_span: extraction failures (e.g. ValueError: Circular reference detected from json.dumps on self-referencing kwargs seen on litellm router retries) log and drop attributes instead of raising into the traced call.

Behavior before/after

LiteLLMInstrumentor().instrument()
resp = await litellm.acompletion(model=..., messages=..., stream=True)
# before: type(resp) is async_generator  → litellm's Responses bridge raises
# after:  isinstance(resp, CustomStreamWrapper) is True; iteration + span capture unchanged

Testing

  • 4 new tests in tests/test_streaming_transparency.py: sync + async stream type preservation (content and span capture asserted), foreign-iterator pass-through on aresponses, and extractor isolation on cyclic kwargs.
  • Full existing package suite passes: 115 passed. ruff check / ruff format --check clean on changed files.

Notes / shortcomings

  • The anthropic-messages streaming wrappers likely share the type-swap pattern; left out to keep this PR small — happy to follow up.
  • _json_serialize in openinference-instrumentation core could additionally be made cycle-safe (benefits all instrumentors); also left out of scope here.
  • We currently run this exact fix in production at Simpplr as a pinned monkeypatch shim and will drop it in favor of the released fix.

@sushantsimpplr
sushantsimpplr requested a review from a team as a code owner August 18, 2026 05:59
@dosubot dosubot Bot added the size:L This PR changes 100-499 lines, ignoring generated files. label Aug 18, 2026
@github-actions

Copy link
Copy Markdown
Contributor

CLA Assistant Lite bot:
Thank you for your submission, we really appreciate it. Like many open-source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution. You can sign the CLA by just posting a Pull Request Comment same as the below format.


I have read the CLA Document and I hereby sign the CLA


You can retrigger this bot by commenting recheck in this Pull Request

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

Labels

size:L This PR changes 100-499 lines, ignoring generated files.

Projects

Status: No status

1 participant