fix(messages): request usage on streaming so the bridge meters tokens - #1180
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
WalkthroughChangesStreaming message-to-completion conversion now requests a trailing usage chunk when enabled. Azure, Groq, Watsonx, and xAI conversions exclude the OpenAI-specific streaming option, with tests covering conversion, bridge propagation, and provider filtering. Streaming usage propagation
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 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 |
Codecov Report✅ All modified and coverable lines are covered by tests.
... and 32 files with indirect coverage changes 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 `@tests/unit/test_messages.py`:
- Line 445: Move the inline import of ChatCompletion, ChatCompletionMessage, and
Choice from the test body to the module-level import section at the top of
tests/unit/test_messages.py, preserving the existing import usage and removing
the duplicate inline import.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: d57dc3da-457a-417f-b251-c030a8d13bc5
📒 Files selected for processing (3)
src/any_llm/utils/messages_compat.pytests/unit/test_messages.pytests/unit/test_messages_compat.py
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/any_llm/providers/watsonx/watsonx.py (1)
74-79: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winApply the filter after merging
kwargs.converted_params.update(kwargs)can reintroducestream_options, so the unsupported field may still reach Watsonx. Filter it after the merge and add coverage for thekwargspath.🤖 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 `@src/any_llm/providers/watsonx/watsonx.py` around lines 74 - 79, Update the parameter-building flow around converted_params so kwargs are merged before filtering unsupported fields, ensuring stream_options cannot be reintroduced by converted_params.update(kwargs). Preserve the existing reasoning_effort handling and add coverage verifying stream_options supplied through kwargs is excluded from the Watsonx request.
🤖 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.
Outside diff comments:
In `@src/any_llm/providers/watsonx/watsonx.py`:
- Around line 74-79: Update the parameter-building flow around converted_params
so kwargs are merged before filtering unsupported fields, ensuring
stream_options cannot be reintroduced by converted_params.update(kwargs).
Preserve the existing reasoning_effort handling and add coverage verifying
stream_options supplied through kwargs is excluded from the Watsonx request.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 11aa25b2-4eca-47b0-b86e-2b81372e4b74
📒 Files selected for processing (2)
src/any_llm/providers/watsonx/watsonx.pytests/unit/providers/test_watsonx_provider.py
The Messages->Completions bridge streamed without setting stream_options.include_usage, so OpenAI-compatible backends omitted token usage from their chunks. The recent trailing-usage-chunk fix (#1151) then had no usage-only chunk to flush, so streamed amessages against those providers reported zero input and output tokens. Native Anthropic was unaffected (it overrides _amessages and streams usage directly). Request stream_options.include_usage in the streaming branch of messages_params_to_completion_params so the backend emits the trailing usage-only chunk that the stream wrapper flushes into the closing message_delta. Providers that do not support stream_options already strip it in their own param conversion, and the native Anthropic provider never reaches this bridge, so the injection is safe across providers. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…payload The Messages bridge now sets stream_options for every streamed request, but Watsonx merges its params dict straight into the chat payload without filtering against TextChatParameters, so the OpenAI-only stream_options field would be forwarded to the Watsonx API. Exclude it in _convert_completion_params, matching the other OpenAI-incompatible providers (Cerebras, Cohere, Mistral, Ollama, Together), so streaming stays unaffected while usage still comes back natively. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Like Watsonx, the Groq and xAI providers pass their params through to a non-OpenAI SDK client that rejects the OpenAI-only stream_options field. The Messages bridge now sets stream_options for every streamed request, so streaming /v1/messages (and any streaming call that carries stream_options) raised TypeError: create() got an unexpected keyword argument 'stream_options'. Exclude it in _convert_completion_params, matching the other providers that already drop it (Cerebras, Cohere, Mistral, Ollama, Together, Watsonx). Caught by the integration suite (test_messages_streaming[groq], [xai]). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
d6331fa to
e3baea6
Compare
The Azure AI Inference SDK does not model the OpenAI-only stream_options field; complete() forwards unknown kwargs down to the transport, which rejects them with TypeError: ClientSession._request() got an unexpected keyword argument 'stream_options'. The Messages bridge now sets stream_options for every streamed request, so streaming /v1/messages through the Azure provider would crash. Exclude it in _convert_completion_params, matching the other providers that already drop it. Reproduced in isolation against the installed SDK; the integration suite could not catch this because Azure is not credentialed in CI. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/any_llm/providers/azure/azure.py (1)
186-193: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winFilter
stream_optionsafter mergingkwargs.params.model_dump()excludes it, butcall_kwargs.update(kwargs)can add it back and send an unsupported field to Azure.🤖 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 `@src/any_llm/providers/azure/azure.py` around lines 186 - 193, Update the call-kwargs construction in the relevant Azure provider method to remove or filter the stream_options key after call_kwargs.update(kwargs), ensuring kwargs cannot reintroduce this unsupported field before returning. Preserve the existing model_dump exclusions and other merged parameters.
🤖 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 `@tests/unit/providers/test_azure_provider.py`:
- Around line 232-243: Extend
test_convert_completion_params_drops_stream_options to pass stream_options
directly through the kwargs argument of
AzureProvider._convert_completion_params, alongside or separately from the
CompletionParams field, and assert the returned mapping excludes it. Preserve
the existing params-based coverage while adding the extra-kwargs path that
exercises call_kwargs.update(kwargs).
---
Outside diff comments:
In `@src/any_llm/providers/azure/azure.py`:
- Around line 186-193: Update the call-kwargs construction in the relevant Azure
provider method to remove or filter the stream_options key after
call_kwargs.update(kwargs), ensuring kwargs cannot reintroduce this unsupported
field before returning. Preserve the existing model_dump exclusions and other
merged parameters.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: fd48acf9-1531-452e-82ed-eda30b8a3bd0
📒 Files selected for processing (2)
src/any_llm/providers/azure/azure.pytests/unit/providers/test_azure_provider.py
There was a problem hiding this comment.
Pull request overview
Fixes missing token metering when streaming through the Messages to Completions compatibility bridge by ensuring OpenAI-compatible backends are asked to emit usage during streaming. This aligns the Messages streaming path with the already-fixed trailing usage-only chunk handling so streamed amessages reports correct input, output, and cache token counts.
Changes:
- Inject
stream_options={"include_usage": True}whenMessagesParams.stream=Trueinmessages_params_to_completion_params. - Filter out
stream_optionsin providers whose SDKs reject unknown OpenAI streaming knobs (xAI, Groq, Watsonx, Azure). - Add unit tests covering streaming vs non-streaming propagation of
stream_optionsthrough both the compat layer and the_amessagesbridge.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit/test_messages.py | Verifies _amessages passes include_usage when streaming and omits it when non-streaming. |
| tests/unit/test_messages_compat.py | Verifies the compat conversion adds stream_options only for stream=True and omits it otherwise. |
| tests/unit/providers/test_xai_provider.py | Asserts xAI conversion drops stream_options. |
| tests/unit/providers/test_watsonx_provider.py | Asserts Watsonx conversion drops stream_options. |
| tests/unit/providers/test_groq_provider.py | Asserts Groq conversion drops stream_options. |
| tests/unit/providers/test_azure_provider.py | Asserts Azure conversion drops stream_options. |
| src/any_llm/utils/messages_compat.py | Implements stream_options.include_usage injection for streaming bridge requests. |
| src/any_llm/providers/xai/xai.py | Drops unsupported stream_options in xAI provider param conversion. |
| src/any_llm/providers/watsonx/watsonx.py | Drops unsupported stream_options in Watsonx provider param conversion. |
| src/any_llm/providers/groq/groq.py | Drops unsupported stream_options in Groq provider param conversion. |
| src/any_llm/providers/azure/azure.py | Drops unsupported stream_options in Azure provider param conversion. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
khaledosman
left a comment
There was a problem hiding this comment.
Reviewed the fix end-to-end: verified the bridge chain into the #1151 trailing-chunk flush, and checked the provider blast radius against actual SDK signatures (groq create and xai chat.create have no stream_options parameter, azure complete() does not model it, huggingface-hub 1.3.2 does, and cohere/together/cerebras/ollama/mistral already strip it; bedrock/sagemaker/gemini/lmstudio build kwargs explicitly and ignore it). Groq and xAI still meter streamed usage despite the strip since both surface usage natively on chunks. Touched test files, ruff check and format pass locally; CI is green.
Two minor inline comments below. One observation that fits neither diff line: every BaseOpenAIProvider subclass (perplexity, deepinfra, dashscope, moonshot, ...) now sends stream_options to its remote endpoint on the streamed messages path. The OpenAI SDK accepts it, but backend tolerance can't be verified statically; one run with the run-integration-tests label would derisk that.
🤖 This review was created by Claude Code.
any_llm.types.completion is a core dependency and the file already imports from it at the top, so ChatCompletion, ChatCompletionMessage and Choice belong in the module-level import block rather than inline in a test body. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Description
Streaming through the Messages to Completions bridge did not set
stream_options.include_usage, so OpenAI-compatible backends omitted token usage from their streamed chunks. The recent trailing-usage-chunk fix (#1151) then had no usage-only chunk to flush, so streamedamessagesagainst those providers reported zero input and output tokens (and no cache). Native Anthropic was unaffected, since it overrides_amessagesand streams usage directly.This requests
stream_options.include_usagein the streaming branch ofmessages_params_to_completion_params, so the backend emits the trailing usage-only chunk that the stream wrapper flushes into the closingmessage_delta. It completes the chain #1151 started: #1151 captures the trailing chunk, this makes sure the trailing chunk is actually produced.Safe across providers:
stream_options(Cerebras, Ollama, Together, Cohere, Mistral) already strip it in their own_convert_completion_params._amessages), so its SDK never sees astream_optionskwarg.Downstream context: this is the root cause of the zero-metering behavior reported in mozilla-ai/otari#256 (streamed
/v1/messagesrecorded zero tokens and zero cost while non-streaming messages and streaming chat completions metered fine). The old any-llm gateway solved the same class of bug for chat completions in #974; this brings the messages path in line.PR Type
Relevant issues
Complements #1151. Root cause of mozilla-ai/otari#256.
Checklist
AI Usage Information
AI Model used: Claude Opus 4.8
AI Developer Tool used: Claude Code
Any other info you'd like to share: Drafted by Claude via back-and-forth with @njbrake. The investigation (tracing the bug across fix(messages): report Anthropic streaming usage from the trailing usage-only chunk #1151, fix(gateway): inject stream_options for streaming usage tracking #974, and the pinned 1.17.0), the diagnosis, and the decisions are his; the code and this prose are Claude's. On the "discuss with the human, not the AI" policy: respected. @njbrake will reply to review questions himself.
I am an AI Agent filling out this form (check box if true)
Testing
Full unit suite passes (1490 passed, 64 skipped), plus ruff and mypy strict clean over
srcand the touched tests. New tests:test_messages_compat.py:stream_options.include_usageis present when streaming, absent when not streaming and whenstreamis unset.test_messages.py: theCompletionParamshanded to_acompletionactually carriesinclude_usageon the streaming path and omits it on the non-streaming path. This is the coverage the trailing-chunk fix lacked (its tests hand-fed usage chunks, so they could not catch that usage was never requested).Summary by CodeRabbit
stream_optionsfrom provider requests where unsupported (Watsonx, Groq, xAI, and Azure).stream_options.