fix: litellm async calls with stream reports token usage and span status

[GustavoCaso]

While testing locally Arize, I noticed that when using the acompletion function with stream=True, stream_options={"include_usage": True} options, the reported traces did not included the count token, and the status was undefined.

I chekced the code and noticed that the token count issue, was that we pass the usage_stats object to the _set_token_counts_from_usage directly. When the function checks if the object has the usage attribute, it return false and exit. Causing the span to not report token count.

I modified the code so the object includes the usage attribute, and called the _set_span_status function.

After those changes the token count and the status is reported:

---

Note

Ensure acompletion/completion streaming reports token counts and sets span status, and add tests covering async streaming and usage reporting.

• Instrumentation (LiteLLM):
  • Wrap streaming usage with SimpleNamespace(usage=...) before calling _set_token_counts_from_usage in _finalize_sync_streaming_span and _finalize_streaming_span.
  • Set span status via _set_span_status(span, aggregated_output) after async streaming completes.
• Tests:
  • Add async streaming tests for acompletion validating output value, status OK, and token counts when stream_options={"include_usage": True}.
  • Cover context attributes in both standard and usage-included async streaming paths.

^{Written by Cursor Bugbot for commit 78956e8. This will update automatically on new commits. Configure here.}

[github-actions]

CLA Assistant Lite bot All contributors have signed the CLA ✍️ ✅

[GustavoCaso]

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

[GustavoCaso]

@codefromthecrypt could I get some feedback for theses changes 😄 ?

[codefromthecrypt]

I don't have workflow approval, but I think the main step is to add unit test, or another row in the unit test, for this case.

Main thing is you need to verify the changes are needed, esp the status thing. so you can look at litellm's existing test_completion_sync_streaming which needs to be ported for async

there is some code in litellm for async, but basically always the openai code is in the best shape, as that's the most used. you can steal patterns from this python/instrumentation/openinference-instrumentation-openai/tests/openinference/instrumentation/openai/test_instrumentor.py

then you need to run tests the same way they run in CI

[codefromthecrypt]

ps one tip is that to run the tests, enter the python directory and do uvx --with tox-uv tox run -e py313-ci-litellm this will be the main thing that will ensure CI is happy

[GustavoCaso]

Thanks for the feedback @codefromthecrypt.

I added tests 😄 and validated that are passing using uvx --with tox-uv tox run -e py313-ci-litellm

[GustavoCaso]

@caroger would it be possible to get some feedback on these changes.

Thanks in advance 😄

[hewliyang]

Bump @caroger. This is a pretty straightforward bug affecting streaming sync completions as well.

Here's a monkeypatch for reference.

def _patch_litellm_streaming_token_counts() -> None:
    """
    Patch openinference-instrumentation-litellm to fix streaming token counts.

    The library has a bug where streaming handlers pass `litellm.Usage` directly to
    `_set_token_counts_from_usage`, but that function expects an object WITH a `.usage`
    attribute. This causes token counts to be missing for all streaming responses.

    Fix pending upstream: https://github.com/Arize-ai/openinference/pull/2480
    """
    try:
        import openinference.instrumentation.litellm as litellm_instr
        from litellm.types.utils import Usage

        original_set_token_counts = litellm_instr._set_token_counts_from_usage

        def patched_set_token_counts(span, result):
            # Streaming handlers pass Usage directly, but function expects object with .usage
            if isinstance(result, Usage):
                result = SimpleNamespace(usage=result)
            return original_set_token_counts(span, result)

        litellm_instr._set_token_counts_from_usage = patched_set_token_counts
    except Exception:
        pass