diff --git a/ag2/config/xai/mappers.py b/ag2/config/xai/mappers.py index 5557f327c0d2..8acfe771dc96 100644 --- a/ag2/config/xai/mappers.py +++ b/ag2/config/xai/mappers.py @@ -313,9 +313,15 @@ def normalize_usage(usage: usage_pb2.SamplingUsage | None) -> Usage: if usage is None: return Usage() + # xAI reports reasoning beside the completion (total = prompt + completion + + # reasoning); completion includes it, with thinking_tokens as the breakdown. + completion = float(usage.completion_tokens) if usage.completion_tokens else None + if usage.reasoning_tokens: + completion = (completion or 0.0) + usage.reasoning_tokens + return Usage( prompt_tokens=float(usage.prompt_tokens) if usage.prompt_tokens else None, - completion_tokens=float(usage.completion_tokens) if usage.completion_tokens else None, + completion_tokens=completion, total_tokens=float(usage.total_tokens) if usage.total_tokens else None, cache_read_input_tokens=float(usage.cached_prompt_text_tokens) if usage.cached_prompt_text_tokens else None, thinking_tokens=float(usage.reasoning_tokens) if usage.reasoning_tokens else None, diff --git a/test/ag_ui/test_usage.py b/test/ag_ui/test_usage.py index 19569c210cad..ecdb2be77dbb 100644 --- a/test/ag_ui/test_usage.py +++ b/test/ag_ui/test_usage.py @@ -629,8 +629,8 @@ async def test_bedrock_cache_is_added_to_its_input(self) -> None: @pytest.mark.parametrize("provider", ["openai", "xai"]) async def test_other_providers_are_reported_as_they_count(self, provider: str) -> None: - """OpenAI counts cache and reasoning inside its totals already. xAI's reasoning - is not known to sit outside them, so it is not added in. + """OpenAI and xAI count reasoning inside the normalized completion already, + so it is not added in again. """ usage = Usage( prompt_tokens=100, diff --git a/test/config/xai/test_xai_usage.py b/test/config/xai/test_xai_usage.py index 99f64606b9c3..1b2932872194 100644 --- a/test/config/xai/test_xai_usage.py +++ b/test/config/xai/test_xai_usage.py @@ -32,13 +32,29 @@ def test_cached_prompt_tokens_map_to_cache_read() -> None: def test_reasoning_tokens_map_to_thinking_tokens() -> None: - raw = usage_pb2.SamplingUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15, reasoning_tokens=3) + raw = usage_pb2.SamplingUsage(prompt_tokens=10, completion_tokens=5, total_tokens=18, reasoning_tokens=3) result = normalize_usage(raw) assert result.thinking_tokens == 3.0 +def test_completion_includes_reasoning_tokens() -> None: + # xAI reports reasoning beside the completion: total = prompt + completion + reasoning. + raw = usage_pb2.SamplingUsage(prompt_tokens=204, completion_tokens=32, total_tokens=536, reasoning_tokens=300) + + result = normalize_usage(raw) + + assert result == Usage(prompt_tokens=204.0, completion_tokens=332.0, total_tokens=536.0, thinking_tokens=300.0) + assert result.prompt_tokens + result.completion_tokens == result.total_tokens + + +def test_reasoning_without_completion_tokens() -> None: + raw = usage_pb2.SamplingUsage(prompt_tokens=10, total_tokens=50, reasoning_tokens=40) + + assert normalize_usage(raw).completion_tokens == 40.0 + + def test_zero_or_missing_fields_become_none() -> None: raw = usage_pb2.SamplingUsage()