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>
This commit is contained in:
Kerry Lu 2026-10-14 17:02:00 -07:00 • committed by GitHub
parent 6d1a3bce2c
commit 44e94eb985
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 86 additions and 12 deletions

View file

@ -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 = (

View file

@ -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

View file

@ -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": {

View file

@ -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.

View file

@ -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