fix(messages): request usage on streaming so the bridge meters tokens

src/any_llm/providers/azure/azure.py

@@ -180,7 +180,12 @@ def _convert_completion_params(params: CompletionParams, **kwargs: Any) -> dict[
         if params.response_format:
             azure_response_format = _convert_response_format(params.response_format)
 
-        call_kwargs = params.model_dump(exclude_none=True, exclude={"model_id", "messages", "response_format"})
+        # stream_options is an OpenAI-only knob (the Messages bridge sets it to
+        # request streaming usage); the Azure AI Inference SDK does not model it
+        # and forwards unknown kwargs to the transport, which rejects it.
+        call_kwargs = params.model_dump(
+            exclude_none=True, exclude={"model_id", "messages", "response_format", "stream_options"}
+        )
         if azure_response_format:
             call_kwargs["response_format"] = azure_response_format
 

src/any_llm/providers/groq/groq.py

@@ -71,8 +71,10 @@ class GroqProvider(AnyLLM):
     @override
     def _convert_completion_params(params: CompletionParams, **kwargs: Any) -> dict[str, Any]:
         """Convert CompletionParams to kwargs for Groq API."""
-        # Groq does not support providing reasoning effort
-        converted_params = params.model_dump(exclude_none=True, exclude={"model_id", "messages"})
+        # Groq does not support providing reasoning effort.
+        # stream_options is an OpenAI-only knob (the Messages bridge sets it to
+        # request streaming usage); the Groq SDK rejects it, so drop it here.
+        converted_params = params.model_dump(exclude_none=True, exclude={"model_id", "messages", "stream_options"})
         if converted_params.get("reasoning_effort") in ("auto", "none"):
             converted_params.pop("reasoning_effort")
         converted_params.update(kwargs)

src/any_llm/providers/watsonx/watsonx.py

@@ -66,9 +66,13 @@ class WatsonxProvider(AnyLLM):
     @override
     def _convert_completion_params(params: CompletionParams, **kwargs: Any) -> dict[str, Any]:
         """Convert CompletionParams to kwargs for Watsonx API."""
-        # Watsonx does not support providing reasoning effort
+        # Watsonx does not support providing reasoning effort.
+        # stream_options is an OpenAI-only knob (the Messages bridge sets it to
+        # request streaming usage); Watsonx merges the params dict straight into
+        # its chat payload, so drop it here as the OpenAI-incompatible providers
+        # already do, rather than forward an unsupported field.
         converted_params = params.model_dump(
-            exclude_none=True, exclude={"model_id", "messages", "response_format", "stream"}
+            exclude_none=True, exclude={"model_id", "messages", "response_format", "stream", "stream_options"}
         )
         if converted_params.get("reasoning_effort") in ("auto", "none"):
             converted_params.pop("reasoning_effort")

src/any_llm/providers/xai/xai.py

@@ -59,13 +59,16 @@ class XaiProvider(AnyLLM):
     @override
     def _convert_completion_params(params: CompletionParams, **kwargs: Any) -> dict[str, Any]:
         """Convert CompletionParams to kwargs for xAI API."""
-        # xAI does not support providing reasoning effort
+        # xAI does not support providing reasoning effort.
+        # stream_options is an OpenAI-only knob (the Messages bridge sets it to
+        # request streaming usage); the xAI SDK rejects it, so drop it here.
         converted_params = params.model_dump(
             exclude_none=True,
             exclude={
                 "model_id",
                 "messages",
                 "stream",
+                "stream_options",
                 "response_format",
                 "tools",
                 "tool_choice",

src/any_llm/utils/messages_compat.py

@@ -70,6 +70,15 @@ def messages_params_to_completion_params(params: MessagesParams) -> dict[str, An
         result["stop"] = params.stop_sequences
     if params.stream is not None:
         result["stream"] = params.stream
+        if params.stream:
+            # OpenAI-compatible backends omit token usage from streamed chunks
+            # unless asked for it, so the streamed Messages bridge would report
+            # zero tokens. Request the trailing usage-only chunk that the
+            # streaming wrapper flushes into the closing ``message_delta``.
+            # Providers that don't support ``stream_options`` strip it in their
+            # own param conversion, and the native Anthropic provider never
+            # reaches this bridge (it overrides ``_amessages``).
+            result["stream_options"] = {"include_usage": True}
 
     if params.output_format is not None:
         if is_structured_output_type(params.output_format):

tests/unit/providers/test_azure_provider.py

@@ -227,3 +227,17 @@ def test_azure_no_api_key_no_credential_raises() -> None:
     with mock_azure_provider():
         with pytest.raises(MissingApiKeyError):
             AzureProvider(api_base=custom_endpoint)
+
+
+def test_convert_completion_params_drops_stream_options() -> None:
+    """stream_options is an OpenAI-only knob (set by the Messages bridge for
+    streaming usage); the Azure AI Inference SDK does not model it and forwards
+    unknown kwargs to the transport, which rejects them, so it must be dropped."""
+    params = CompletionParams(
+        model_id="test-model",
+        messages=[{"role": "user", "content": "Hello"}],
+        stream=True,
+        stream_options={"include_usage": True},
+    )
+    result = AzureProvider._convert_completion_params(params)
+    assert "stream_options" not in result

tests/unit/providers/test_groq_provider.py

@@ -108,6 +108,23 @@ async def test_reasoning_effort_filtered_out(reasoning_effort: str) -> None:
         assert "reasoning_effort" not in call_kwargs
 
 
+def test_stream_options_filtered_out() -> None:
+    """stream_options is an OpenAI-only knob (set by the Messages bridge for
+    streaming usage); the Groq SDK rejects it, so it must be dropped."""
+    pytest.importorskip("groq")
+    from any_llm.providers.groq.groq import GroqProvider
+
+    result = GroqProvider._convert_completion_params(
+        CompletionParams(
+            model_id="llama-3.3-70b-versatile",
+            messages=[{"role": "user", "content": "Hello"}],
+            stream=True,
+            stream_options={"include_usage": True},
+        )
+    )
+    assert "stream_options" not in result
+
+
 def test_to_chat_completion_extracts_cached_tokens() -> None:
     """Test that cached tokens from Groq usage are extracted into prompt_tokens_details."""
     from groq.types.chat import ChatCompletion as GroqChatCompletion

tests/unit/providers/test_watsonx_provider.py

@@ -176,6 +176,20 @@ def test_convert_completion_params_filters_reasoning_effort(reasoning_effort: st
     assert "reasoning_effort" not in result
 
 
+def test_convert_completion_params_drops_stream_options() -> None:
+    """stream_options is an OpenAI-only knob (set by the Messages bridge for
+    streaming usage). Watsonx merges the params dict straight into its chat
+    payload, so it must be dropped rather than forwarded as an unsupported field."""
+    params = CompletionParams(
+        model_id="test-model",
+        messages=[{"role": "user", "content": "Hello"}],
+        stream=True,
+        stream_options={"include_usage": True},
+    )
+    result = WatsonxProvider._convert_completion_params(params)
+    assert "stream_options" not in result
+
+
 def test_convert_streaming_chunk_with_tool_calls() -> None:
     """Test streaming chunk conversion with tool calls."""
     chunk = {

tests/unit/providers/test_xai_provider.py

@@ -146,3 +146,19 @@ async def test_reasoning_effort_filtered_out(reasoning_effort: str) -> None:
         )
         _, call_kwargs = mock_xai.return_value.chat.create.call_args
         assert "reasoning_effort" not in call_kwargs
+
+
+def test_stream_options_filtered_out() -> None:
+    """stream_options is an OpenAI-only knob (set by the Messages bridge for
+    streaming usage); the xAI SDK rejects it, so it must be dropped."""
+    from any_llm.providers.xai.xai import XaiProvider
+
+    result = XaiProvider._convert_completion_params(
+        CompletionParams(
+            model_id="model",
+            messages=[{"role": "user", "content": "Hello"}],
+            stream=True,
+            stream_options={"include_usage": True},
+        )
+    )
+    assert "stream_options" not in result

tests/unit/test_messages.py

@@ -9,7 +9,10 @@
 from any_llm.any_llm import AnyLLM
 from any_llm.api import amessages
 from any_llm.types.completion import (
+    ChatCompletion,
     ChatCompletionChunk,
+    ChatCompletionMessage,
+    Choice,
     ChoiceDelta,
     ChunkChoice,
     CompletionUsage,
@@ -405,6 +408,71 @@ async def mock_stream() -> Any:
     assert "message_stop" in types
 
 
+@pytest.mark.asyncio
+async def test_default_amessages_streaming_requests_include_usage() -> None:
+    """Streaming through the bridge must ask the backend for usage, otherwise
+    OpenAI-compatible providers emit no usage-only chunk and the trailing-chunk
+    capture has nothing to report."""
+
+    async def mock_stream() -> AsyncIterator[ChatCompletionChunk]:
+        yield ChatCompletionChunk(
+            id="chunk-1",
+            model="gpt-4",
+            created=0,
+            object="chat.completion.chunk",
+            choices=[ChunkChoice(index=0, delta=ChoiceDelta(), finish_reason="stop")],
+        )
+
+    mock_provider = Mock()
+    mock_provider._acompletion = AsyncMock(return_value=mock_stream())
+
+    params = MessagesParams(
+        model="gpt-4",
+        messages=[{"role": "user", "content": "Hello"}],
+        max_tokens=100,
+        stream=True,
+    )
+    result = await AnyLLM._amessages(mock_provider, params)
+    assert not isinstance(result, (MessageResponse, ParsedMessage))
+    async for _ in result:  # drive the generator so _acompletion is invoked
+        pass
+
+    completion_params = mock_provider._acompletion.call_args.args[0]
+    assert completion_params.stream is True
+    assert completion_params.stream_options == {"include_usage": True}
+
+
+@pytest.mark.asyncio
+async def test_default_amessages_non_streaming_omits_include_usage() -> None:
+    """A non-streaming bridge call has no usage-only chunk to request."""
+    mock_completion = ChatCompletion(
+        id="c",
+        model="gpt-4",
+        created=0,
+        object="chat.completion",
+        choices=[
+            Choice(
+                index=0,
+                message=ChatCompletionMessage(role="assistant", content="hi"),
+                finish_reason="stop",
+            )
+        ],
+    )
+    mock_provider = Mock()
+    mock_provider._acompletion = AsyncMock(return_value=mock_completion)
+
+    params = MessagesParams(
+        model="gpt-4",
+        messages=[{"role": "user", "content": "Hello"}],
+        max_tokens=100,
+        stream=False,
+    )
+    await AnyLLM._amessages(mock_provider, params)
+
+    completion_params = mock_provider._acompletion.call_args.args[0]
+    assert completion_params.stream_options is None
+
+
 @pytest.mark.asyncio
 async def test_default_amessages_streaming_usage_from_trailing_chunk() -> None:
     """Usage (incl. cache) from a trailing usage-only chunk after finish_reason is reported in message_delta."""

tests/unit/test_messages_compat.py

@@ -493,6 +493,45 @@ def test_optional_params_not_included_when_none() -> None:
     assert "tools" not in result
 
 
+def test_stream_requests_include_usage() -> None:
+    """Streaming requests ask the backend for usage, otherwise OpenAI-compatible
+    providers omit it and the streamed bridge reports zero tokens."""
+    params = MessagesParams(
+        model="claude-3-5-sonnet",
+        messages=[{"role": "user", "content": "Hi"}],
+        max_tokens=1024,
+        stream=True,
+    )
+    result = messages_params_to_completion_params(params)
+    assert result["stream"] is True
+    assert result["stream_options"] == {"include_usage": True}
+
+
+def test_non_stream_omits_stream_options() -> None:
+    """A non-streaming request has no usage-only chunk to request."""
+    params = MessagesParams(
+        model="claude-3-5-sonnet",
+        messages=[{"role": "user", "content": "Hi"}],
+        max_tokens=1024,
+        stream=False,
+    )
+    result = messages_params_to_completion_params(params)
+    assert result["stream"] is False
+    assert "stream_options" not in result
+
+
+def test_unset_stream_omits_stream_options() -> None:
+    """When stream is unset, neither stream nor stream_options is included."""
+    params = MessagesParams(
+        model="claude-3-5-sonnet",
+        messages=[{"role": "user", "content": "Hi"}],
+        max_tokens=1024,
+    )
+    result = messages_params_to_completion_params(params)
+    assert "stream" not in result
+    assert "stream_options" not in result
+
+
 def test_temperature_and_top_p_passed_through() -> None:
     """Test that temperature and top_p are passed when set."""
     params = MessagesParams(