From 9aab6ec47b7e1b9f76ffacf340a5ae5e6d228d4e Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Fri, 16 Oct 2026 16:54:38 -0700 Subject: [PATCH] fix(azure_ai): count Grok reasoning tokens in completion_tokens (#46969) * fix(azure_ai): count Grok reasoning tokens in completion_tokens Grok on Azure AI Foundry reports reasoning_tokens outside completion_tokens, like xAI does, so azure_ai/grok-* calls billed only max(visible, reasoning) output tokens. Fold them with XAIChatConfig.fold_reasoning_tokens_into_completion in transform_response and in a new streaming iterator, matching the xai provider. Split out of #46377. * fix(azure_ai): fold reasoning tokens for every deployment name, streaming or not Non-streaming checked the deployment name for grok while streaming checked the chunk model field, so a Grok deployment named without grok was only folded when streaming. The fold already only fires when total_tokens equals prompt plus completion plus reasoning, so it now runs on every azure_ai response. --- litellm/llms/azure_ai/chat/transformation.py | 11 +- .../chat/test_azure_ai_transformation.py | 115 ++++++++++++++++++ 2 files changed, 118 insertions(+), 8 deletions(-) diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index fc22e9a8b73..7aae2cb949d 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -46,14 +46,10 @@ NON_OPENAI_SPEC_MESSAGE_FIELDS: Final = ( ) -def _is_grok_model(model: str) -> bool: - return "grok" in model - - class AzureAIChatCompletionResponseIterator(OpenAIChatCompletionResponseIterator): def chunk_parser(self, chunk: dict[str, object]) -> ModelResponseStream: # mutable-ok: mirrors override contract usage: Final = chunk.get("usage") - if _is_grok_model(str(chunk.get("model") or "")) and isinstance(usage, dict): + if isinstance(usage, dict): XAIChatConfig.fold_reasoning_tokens_into_completion(usage) return super().chunk_parser(chunk) @@ -127,7 +123,7 @@ class AzureAIStudioConfig(OpenAIConfig): """ Check if the model supports stop tokens. """ - if _is_grok_model(model): + if "grok" in model: # Reuse Xai method for Grok model xai_config: Final = XAIChatConfig() return xai_config.supports_stop_reason(model) @@ -345,8 +341,7 @@ class AzureAIStudioConfig(OpenAIConfig): api_key=api_key, json_mode=json_mode, ) - if _is_grok_model(model): - XAIChatConfig.fold_reasoning_tokens_into_completion(response) + XAIChatConfig.fold_reasoning_tokens_into_completion(response) return response def get_model_response_iterator( diff --git a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py index a445d951bc2..634228da1bf 100644 --- a/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py +++ b/tests/unit/llms/azure_ai/chat/test_azure_ai_transformation.py @@ -384,6 +384,66 @@ def test_azure_ai_grok_bills_reasoning_tokens_reported_outside_completion_tokens assert response._hidden_params["response_cost"] == pytest.approx(expected_cost, rel=1e-9) +@pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize( + ("upstream_model", "usage", "expected_completion_tokens"), + [ + pytest.param( + "grok-4.7", + { + "prompt_tokens": 14, + "completion_tokens": 10, + "total_tokens": 336, + "completion_tokens_details": {"reasoning_tokens": 312}, + }, + 322, + id="grok-reports-reasoning-outside-completion", + ), + pytest.param( + "gpt-5.4-mini", + { + "prompt_tokens": 14, + "completion_tokens": 322, + "total_tokens": 336, + "completion_tokens_details": {"reasoning_tokens": 312}, + }, + 322, + id="reasoning-already-inside-completion", + ), + ], +) +def test_azure_ai_counts_reasoning_tokens_once_whatever_the_deployment_is_named( + respx_mock: respx.MockRouter, + upstream_model: str, + usage: dict[str, object], + expected_completion_tokens: int, +) -> None: + api_url: Final = "https://example-resource.services.ai.azure.com/models/chat/completions" + respx_mock.post(url__startswith=api_url).mock( + return_value=Response( + 200, + json={ + "id": "chatcmpl-x", + "object": "chat.completion", + "created": 1700000000, + "model": upstream_model, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "42"}, "finish_reason": "stop"}], + "usage": usage, + }, + ) + ) + + response: Final = litellm.completion( + model="azure_ai/team-reasoner", + messages=[{"role": "user", "content": "What is the meaning of life?"}], + api_base="https://example-resource.services.ai.azure.com/models", + api_key="azure-ai-test-key", + ) + + assert response.usage.completion_tokens == expected_completion_tokens + assert response.usage.total_tokens == response.usage.prompt_tokens + response.usage.completion_tokens + + @pytest.mark.respx(assert_all_called=True) def test_azure_ai_grok_stream_folds_reasoning_tokens_into_the_usage_chunk(respx_mock: respx.MockRouter) -> None: api_url: Final = "https://example-resource.services.ai.azure.com/models/chat/completions" @@ -419,6 +479,61 @@ def test_azure_ai_grok_stream_folds_reasoning_tokens_into_the_usage_chunk(respx_ assert usage_chunks[0].total_tokens == usage_chunks[0].prompt_tokens + usage_chunks[0].completion_tokens +@pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize( + ("completion_tokens", "expected_completion_tokens"), + [ + pytest.param(10, 322, id="reasoning-outside-completion"), + pytest.param(322, 322, id="reasoning-already-inside-completion"), + ], +) +def test_azure_ai_stream_counts_reasoning_tokens_once_whatever_the_chunks_call_the_model( + respx_mock: respx.MockRouter, + completion_tokens: int, + expected_completion_tokens: int, +) -> None: + api_url: Final = "https://example-resource.services.ai.azure.com/models/chat/completions" + chunk_base: Final = { + "id": "chatcmpl-x", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "team-reasoner", + } + chunks_sent: Final = ( + {"choices": [{"index": 0, "delta": {"role": "assistant", "content": "42"}, "finish_reason": None}]}, + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + { + "choices": [], + "usage": { + "prompt_tokens": 14, + "completion_tokens": completion_tokens, + "total_tokens": 336, + "completion_tokens_details": {"reasoning_tokens": 312}, + }, + }, + ) + stream_body: Final = "".join(f"data: {json.dumps({**chunk_base, **chunk})}\n\n" for chunk in chunks_sent) + respx_mock.post(url__startswith=api_url).mock( + return_value=Response(200, text=f"{stream_body}data: [DONE]\n\n", headers={"content-type": "text/event-stream"}) + ) + + chunks: Final = tuple( + litellm.completion( + model="azure_ai/team-reasoner", + messages=[{"role": "user", "content": "What is the meaning of life?"}], + api_base="https://example-resource.services.ai.azure.com/models", + api_key="azure-ai-test-key", + stream=True, + stream_options={"include_usage": True}, + ) + ) + usage_chunks: Final = [chunk.usage for chunk in chunks if getattr(chunk, "usage", None) is not None] + + assert len(usage_chunks) == 1 + assert usage_chunks[0].completion_tokens == expected_completion_tokens + assert usage_chunks[0].total_tokens == usage_chunks[0].prompt_tokens + usage_chunks[0].completion_tokens + + def test_azure_model_router_response_shows_actual_model(): """ Test that Azure Model Router returns the actual model used in the response,