From 44e94eb9858d30ff26cc991793632122cc078279 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Wed, 14 Oct 2026 17:02:00 -0700 Subject: [PATCH] fix(cost): split xai spend rows across input and output instead of booking the total as output (#46320) * fix(cost): keep the input/output split on streamed xai spend rows Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(cost): split xAI's reported total across input and output by token math Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 14 +++++-- litellm/llms/xai/cost_calculator.py | 18 +++++++- .../integration/cost_calculation/cases/xai.py | 10 ++++- .../unit/llms/xai/test_xai_cost_calculator.py | 14 ++++--- tests/unit/test_main.py | 42 +++++++++++++++++++ 5 files changed, 86 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 2ceb0042f1f..9dde1432fae 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1306,12 +1306,16 @@ def _stream_builder_model_map_cost(response: ModelResponse) -> float | None: return None +def _set_hidden_response_cost(response: ModelResponse, response_cost: float) -> None: + hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor + hidden_params["response_cost"] = response_cost + + def _set_stream_builder_response_cost(response: ModelResponse, logging_obj: Optional["Logging"]) -> None: response_cost: Final = _stream_builder_response_cost(response, logging_obj) if response_cost is None: return - hidden_params: Final = response._hidden_params # pyright: ignore[reportPrivateUsage] # no public accessor - hidden_params["response_cost"] = response_cost + _set_hidden_response_cost(response, response_cost) def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None: @@ -1320,8 +1324,10 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o if isinstance(getattr(usage, "cost", None), (int, float)): return computed_cost: Final = logging_obj.response_cost_calculator(result=response) - if isinstance(computed_cost, (int, float)) and computed_cost > 0: - setattr(usage, "cost", computed_cost) + if not isinstance(computed_cost, (int, float)) or computed_cost <= 0: + return + setattr(usage, "cost", computed_cost) + _set_hidden_response_cost(response, float(computed_cost)) _NON_TEXT_DELTA_FIELDS: Final = ( diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index 164568451d4..de2cbf5f9e3 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -49,6 +49,19 @@ def _cost_reported_by_xai(usage: "Usage") -> float | None: return float(reported_cost) +def _split_reported_cost(model: str, usage: Usage, reported_cost: float) -> tuple[float, float]: + """xAI reports one total, so attribute it to input and output in the ratio our own token math gives.""" + try: + priced_input, priced_output = _token_math_cost_per_token(model, usage) + except Exception: # noqa: BLE001 # generic_cost_per_token raises bare Exception for unmapped models + return 0.0, reported_cost + priced_total: Final = priced_input + priced_output + if priced_total <= 0: + return 0.0, reported_cost + input_cost: Final = reported_cost * priced_input / priced_total + return input_cost, reported_cost - input_cost + + def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: """ Calculates the cost per token for a given XAI model, prompt tokens, and completion tokens. @@ -63,8 +76,11 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]: """ reported_cost: Final = _cost_reported_by_xai(usage) if reported_cost is not None: - return 0.0, reported_cost + return _split_reported_cost(model, usage, reported_cost) + return _token_math_cost_per_token(model, usage) + +def _token_math_cost_per_token(model: str, usage: Usage) -> tuple[float, float]: # XAI-specific completion cost: completion is billed as visible + reasoning # tokens. Detect when the transformation layer already folded them so we # don't double-count; fall back to raw xAI shape for callers that bypass diff --git a/tests/integration/cost_calculation/cases/xai.py b/tests/integration/cost_calculation/cases/xai.py index 2ee676fd4d0..3d0849e1311 100644 --- a/tests/integration/cost_calculation/cases/xai.py +++ b/tests/integration/cost_calculation/cases/xai.py @@ -19,7 +19,6 @@ XAI_GROK_4_7_STREAM: Final = CostTrackingTestCase( name="xai-grok-4.7-stream", covers="quota_management.spend_tracking.cost_matrix.logs_cost", model="xai/grok-4.7", - known_bug="streamed xai responses bill the right total but the spend row splits it as input_cost 0 and output_cost = total: _stamp_streaming_usage_cost writes the stream total onto usage.cost and litellm/llms/xai/cost_calculator.py::cost_per_token reads usage.cost as an xAI reported price", endpoint="/v1/chat/completions", request={ "model": "$MODEL", @@ -67,6 +66,7 @@ XAI_GROK_4_7_STREAM: Final = CostTrackingTestCase( "usage": { "prompt_tokens": 1840, "completion_tokens": 412, + "cost_in_usd_ticks": 48600000, "total_tokens": 2252, "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, "completion_tokens_details": { @@ -83,7 +83,7 @@ XAI_GROK_4_7_STREAM: Final = CostTrackingTestCase( ), ), expected=ExactExpected( - spend=0.006152, input_cost=0.00368, output_cost=0.002472, prompt_tokens=1840, completion_tokens=412 + spend=0.00486, input_cost=0.002907282, output_cost=0.001952718, prompt_tokens=1840, completion_tokens=412 ), ) @@ -148,6 +148,7 @@ XAI_GROK_4_7_TOOL_CALL: Final = CostTrackingTestCase( "usage": { "prompt_tokens": 1840, "completion_tokens": 412, + "cost_in_usd_ticks": 61520000, "total_tokens": 2252, "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, "completion_tokens_details": { @@ -200,6 +201,7 @@ XAI_GROK_4_7_REASONING: Final = CostTrackingTestCase( "usage": { "prompt_tokens": 1840, "completion_tokens": 412, + "cost_in_usd_ticks": 61520000, "total_tokens": 2252, "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, "completion_tokens_details": { @@ -224,6 +226,7 @@ XAI_GROK_4_7_CACHE_READ: Final = with_usage( { "prompt_tokens": 1840, "completion_tokens": 412, + "cost_in_usd_ticks": 51920000, "total_tokens": 2252, "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 640}, "completion_tokens_details": { @@ -242,6 +245,7 @@ XAI_GROK_4_7_TIERED_INPUT_ABOVE_200K: Final = with_usage( { "prompt_tokens": 300000, "completion_tokens": 412, + "cost_in_usd_ticks": 12049440000, "total_tokens": 300412, "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, "completion_tokens_details": { @@ -260,6 +264,7 @@ XAI_GROK_4_7_TIERED_CACHE_READ_ABOVE_200K: Final = with_usage( { "prompt_tokens": 300000, "completion_tokens": 412, + "cost_in_usd_ticks": 6049440000, "total_tokens": 300412, "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 200000}, "completion_tokens_details": { @@ -278,6 +283,7 @@ XAI_GROK_4_7_TIERED_AT_200K_BOUNDARY: Final = with_usage( { "prompt_tokens": 200000, "completion_tokens": 412, + "cost_in_usd_ticks": 8049440000, "total_tokens": 200412, "prompt_tokens_details": {"audio_tokens": 0, "cache_write_tokens": 0, "cached_tokens": 0}, "completion_tokens_details": { diff --git a/tests/unit/llms/xai/test_xai_cost_calculator.py b/tests/unit/llms/xai/test_xai_cost_calculator.py index 6065f1053e7..5219dba42ba 100644 --- a/tests/unit/llms/xai/test_xai_cost_calculator.py +++ b/tests/unit/llms/xai/test_xai_cost_calculator.py @@ -102,8 +102,8 @@ class TestXAICostCalculator: def test_reported_cost_is_preferred_over_token_math(self): """The amount xAI reported, carried on usage.cost by the transformation, is billed. - It lands entirely on completion cost because xAI does not split its total by - direction, the same shape the perplexity calculator returns. + xAI does not split its total by direction, so the split follows the ratio our + own token math gives for the same usage while the sum stays xAI's number. """ usage = Usage( prompt_tokens=100, @@ -111,11 +111,15 @@ class TestXAICostCalculator: total_tokens=300, cost=0.0037756, ) + priced_prompt, priced_completion = cost_per_token( + model="grok-4.7", usage=Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300) + ) - prompt_cost, completion_cost = cost_per_token(model="grok-4-latest", usage=usage) + prompt_cost, completion_cost = cost_per_token(model="grok-4.7", usage=usage) - assert prompt_cost == 0.0 - assert math.isclose(completion_cost, 0.0037756, rel_tol=1e-10) + assert math.isclose(prompt_cost + completion_cost, 0.0037756, rel_tol=1e-10) + assert math.isclose(prompt_cost / completion_cost, priced_prompt / priced_completion, rel_tol=1e-10) + assert not math.isclose(prompt_cost + completion_cost, priced_prompt + priced_completion, rel_tol=1e-3) def test_reported_cost_suppresses_web_search_surcharge(self): """The reported total already covers server-side tool calls. diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 635ee98549e..7dc26fca1f1 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -4250,6 +4250,48 @@ def test_stream_chunk_builder_stamps_streaming_usage_cost_by_default(monkeypatch assert response._hidden_params["response_cost"] == pytest.approx(usage_cost) +def test_stream_chunk_builder_stamped_cost_is_not_re_read_as_xai_reported_cost(): + import time as time_module + + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging + + logging_obj: Final = LiteLLMLogging( + model="xai/grok-4.7", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="completion", + start_time=time_module.time(), + litellm_call_id="stream-builder-xai-stamped-cost", + function_id="1", + ) + logging_obj.model_call_details["custom_llm_provider"] = "xai" + logging_obj.optional_params = {} + usage_chunk: Final = _stream_builder_text_chunk("grok-4.7", "") + usage_chunk.usage = Usage(prompt_tokens=1840, completion_tokens=412, total_tokens=2252) + chunks: Final = [ + _stream_builder_text_chunk("grok-4.7", "Hello.", finish_reason="stop"), + usage_chunk, + ] + + response: Final = litellm.stream_chunk_builder( + chunks=chunks, messages=[{"role": "user", "content": "hi"}], logging_obj=logging_obj + ) + + assert response is not None + usage_cost: Final = getattr(response.usage, "cost", None) + assert usage_cost is not None + assert usage_cost > 0 + + spend_row_cost: Final = logging_obj._response_cost_calculator(result=response) + breakdown: Final = logging_obj.cost_breakdown + assert breakdown is not None + assert breakdown["input_cost"] > 0 + assert breakdown["output_cost"] > 0 + assert breakdown["input_cost"] + breakdown["output_cost"] == pytest.approx(usage_cost) + assert spend_row_cost == pytest.approx(usage_cost) + assert response._hidden_params.get("response_cost") == pytest.approx(usage_cost) + + def test_stream_chunk_builder_skips_stamp_when_cost_is_unpriceable(): import time as time_module