refactor(cost): cut the PR back to billing a reported usage.cost, no client-visible rewrite and provider readers as on main

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-10-27 21:39:26 +00:00
parent a41700c297
commit 9c6954c00a
12 changed files with 65 additions and 350 deletions

View file

@ -408,7 +408,14 @@ def _per_second_pricing_cost(
return resolved_cost_per_second * seconds, 0.0
_PROVIDERS_WITH_OWN_REPORTED_COST_READER: Final = frozenset(
{LlmProviders.XAI.value, LlmProviders.PERPLEXITY.value, LlmProviders.OPENCODE_GO.value}
)
def _billable_reported_cost(usage_object: Usage | None, custom_llm_provider: str | None) -> float | None:
if custom_llm_provider in _PROVIDERS_WITH_OWN_REPORTED_COST_READER:
return None
return _billable_header_cost(provider_reported_usage_cost(usage_object, custom_llm_provider))

View file

@ -24,6 +24,22 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
"""
## USE PRE-CALCULATED COST FROM PERPLEXITY IF AVAILABLE
## Perplexity returns accurate cost in usage.cost.total_cost including request fees.
## By the time it reaches here, ResponseAPIUsage.parse_cost has already flattened
## that dict down to a float, so both shapes must be accepted.
cost_info: Final = getattr(usage, "cost", None)
total_cost: float | None = None
if isinstance(cost_info, dict):
total_cost = cost_info.get("total_cost")
elif isinstance(cost_info, (int, float)) and not isinstance(cost_info, bool):
total_cost = float(cost_info)
if total_cost is not None:
# Return total cost as completion_cost (prompt_cost=0) since Perplexity
# doesn't break down by input/output in their cost object
return (0.0, float(total_cost))
## FALLBACK: Calculate cost manually if Perplexity doesn't provide it
## GET MODEL INFO
model_info: Final = get_model_info(model=model, custom_llm_provider="perplexity")

View file

@ -1,8 +1,8 @@
"""
Helper util for handling XAI-specific cost calculation
- Prefers the cost xAI reports on the response over recomputing it locally
- Uses the generic cost calculator which already handles tiered pricing correctly
- Handles XAI-specific reasoning token billing (billed as part of completion tokens)
- Skips the web search surcharge when xAI reported a total that already covers it
"""
import math
@ -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.
@ -61,6 +74,9 @@ def cost_per_token(model: str, usage: Usage) -> tuple[float, float]:
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
"""
reported_cost: Final = _cost_reported_by_xai(usage)
if reported_cost is not None:
return _split_reported_cost(model, usage, reported_cost)
return _token_math_cost_per_token(model, usage)

View file

@ -130,7 +130,6 @@ from litellm.proxy.route_llm_request import (
route_request,
)
from litellm.proxy.telemetry.request_context import bind_call_id as bind_telemetry_call_id
from litellm.proxy.usage_cost_egress import with_client_usage_cost
from litellm.proxy.utils import ( # noqa: F401 # legacy module exports
ProxyLogging,
_check_and_merge_model_level_guardrails, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
@ -274,7 +273,6 @@ from litellm.proxy.litellm_pre_call_utils import (
)
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import (
ModelResponse,
ModelResponseStream,
@ -3327,7 +3325,7 @@ class ProxyBaseLLMRequestProcessing:
user_api_key_dict=user_api_key_dict,
)
return with_client_usage_cost(response, response_cost_for_headers)
return response
@staticmethod
async def _record_container_owners_from_responses_if_needed(
@ -4712,7 +4710,6 @@ class ProxyBaseLLMRequestProcessing:
completion_tokens_details=usage.get("completion_tokens_details"),
prompt_tokens_details=usage.get("prompt_tokens_details"),
server_tool_use=usage.get("server_tool_use"),
cost=usage.get("cost"),
)
@staticmethod
@ -4727,7 +4724,7 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
def _completion_cost_or_none(
model_response: ModelResponse | ResponsesAPIResponse, model_name: str, service_tier: str | None
model_response: ModelResponse, model_name: str, service_tier: str | None
) -> float | None:
try:
return litellm.completion_cost(
@ -4738,7 +4735,7 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
def _logging_obj_cost_or_none(
model_response: ModelResponse | ResponsesAPIResponse, litellm_logging_obj: LiteLLMLoggingObj
model_response: ModelResponse, litellm_logging_obj: LiteLLMLoggingObj
) -> float | None:
# Pricing a frame stamps cost_breakdown and, on failure, the cost-failure debug key onto
# the live logging object. The pass-through handlers never recompute either one, so a
@ -4765,7 +4762,7 @@ class ProxyBaseLLMRequestProcessing:
@staticmethod
def _streamed_usage_cost(
model_response: ModelResponse | ResponsesAPIResponse,
model_response: ModelResponse,
model_name: str,
service_tier: str | None,
litellm_logging_obj: LiteLLMLoggingObj | None,

View file

@ -922,7 +922,6 @@ from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import (
from litellm.proxy.ui_crud_endpoints.user_banner_endpoints import (
router as user_banner_endpoints_router,
)
from litellm.proxy.usage_cost_egress import Pricer, with_client_stream_usage_cost
from litellm.proxy.utils import ( # noqa: F401, RUF100 # legacy module exports
PrismaClient,
ProxyLogging,
@ -10135,22 +10134,6 @@ def _fast_serialize_simple_model_response_stream(
return orjson.dumps(payload)
def _stream_usage_pricer(request_data: Mapping[str, object]) -> Pricer:
"""Bills a usage-bearing stream event the way the request's logging object bills the spend row."""
logging_obj: Final = request_data.get("litellm_logging_obj")
model_name: Final = str(request_data.get("model", ""))
def price(reply: ModelResponse | ResponsesAPIResponse) -> float | None:
return ProxyBaseLLMRequestProcessing._streamed_usage_cost( # pyright: ignore[reportPrivateUsage] # the proxy's one stream pricer
reply,
model_name,
None,
logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None,
)
return price
def _serialize_streaming_chunk(chunk: BaseModel) -> str | bytes:
if isinstance(chunk, ModelResponseStream):
serialized_chunk: Final = _fast_serialize_simple_model_response_stream(chunk)
@ -10489,7 +10472,6 @@ async def async_data_generator(
needs_per_chunk_hook: Final = proxy_logging_obj.needs_per_chunk_streaming_hook()
is_raw_sse_stream: Final = bool(request_data.get("_litellm_raw_sse_stream"))
strip_stream_usage: Final = bool(request_data.get("_litellm_strip_stream_usage"))
price_stream_usage: Final = _stream_usage_pricer(request_data)
raw_sse_buffer = ""
# A stream can start on a deployment with keepalive off and fall back
@ -10576,7 +10558,7 @@ async def async_data_generator(
error_state.observe_chunk(chunk)
raw_passthrough = False
if isinstance(chunk, BaseModel):
chunk = _serialize_streaming_chunk(with_client_stream_usage_cost(chunk, price_stream_usage))
chunk = _serialize_streaming_chunk(chunk)
elif isinstance(chunk, bytes):
chunk = chunk.decode("utf-8", errors="replace")
if is_raw_sse_stream:

View file

@ -1,77 +0,0 @@
"""What a client reads in ``usage.cost``: litellm's bill, written once on the way out of the proxy, onto a copy.
Inside the pipeline ``usage.cost`` stays what the upstream reported, so pricing a reply again (an MCP round, a
logging copy, a stream rebuilt from its chunks) never reads litellm's own bill back as a provider report.
"""
from __future__ import annotations
import math
from collections.abc import Callable
from typing import Final, TypeVar
from pydantic import BaseModel
import litellm
from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse
from litellm.types.utils import ModelResponse, ModelResponseStream, Usage
Pricer = Callable[[ModelResponse | ResponsesAPIResponse], float | None]
"""Prices a reply the way the request's logging object bills it (deployment rates, discount, margin)."""
_RESPONSES_USAGE_EVENTS: Final = frozenset({"response.completed", "response.incomplete"})
_ReplyT = TypeVar("_ReplyT", bound=BaseModel)
def billable(cost: object) -> float | None:
if isinstance(cost, bool) or not isinstance(cost, (int, float)) or not math.isfinite(cost) or cost < 0:
return None
return float(cost)
def client_usage_cost(upstream_cost: object, bill: object, *, wanted: bool = False) -> float | None:
"""The bill, when the upstream put a number in usage.cost or the operator wants it in every usage block
(``litellm.include_cost_in_streaming_usage``). None leaves the usage as the upstream sent it."""
reported: Final = isinstance(upstream_cost, (int, float)) and not isinstance(upstream_cost, bool)
if not reported and not wanted and not litellm.include_cost_in_streaming_usage:
return None
return billable(bill)
def _with_usage_cost(reply: _ReplyT, usage: Usage | ResponseAPIUsage, cost: float) -> _ReplyT:
return reply.model_copy(update={"usage": usage.model_copy(update={"cost": cost})})
def with_client_usage_cost(response: object, bill: object) -> object:
"""A plain chat or Responses reply as the client gets it: a copy with the bill in usage.cost, or the reply
itself when its usage stays as the upstream sent it."""
if not isinstance(response, (ModelResponse, ResponsesAPIResponse)):
return response
usage: Final = getattr(response, "usage", None)
if not isinstance(usage, (Usage, ResponseAPIUsage)):
return response
cost: Final = client_usage_cost(getattr(usage, "cost", None), bill)
if cost is None:
return response
return _with_usage_cost(response, usage, cost)
def with_client_stream_usage_cost(chunk: BaseModel, price: Pricer) -> BaseModel:
"""A streamed event as the client gets it. Only the usage-bearing events are priced: the chat chunk carrying
``usage`` and the Responses ``response.completed`` / ``response.incomplete`` events, which carry the bill
whether or not the upstream reported one, as they did before the bill moved out here."""
if isinstance(chunk, ModelResponseStream):
usage: Final = getattr(chunk, "usage", None)
if not isinstance(usage, Usage):
return chunk
cost: Final = client_usage_cost(getattr(usage, "cost", None), price(ModelResponse(usage=usage)))
return chunk if cost is None else _with_usage_cost(chunk, usage, cost)
if getattr(chunk, "type", None) not in _RESPONSES_USAGE_EVENTS:
return chunk
response: Final = getattr(chunk, "response", None)
if not isinstance(response, ResponsesAPIResponse) or not isinstance(response.usage, ResponseAPIUsage):
return chunk
response_cost: Final = client_usage_cost(getattr(response.usage, "cost", None), price(response), wanted=True)
if response_cost is None:
return chunk
return chunk.model_copy(update={"response": _with_usage_cost(response, response.usage, response_cost)})

View file

@ -55,18 +55,6 @@ def assert_streamed_usage_cost(case_name: str, endpoint: str, response_text: str
assert approx_equal(seen, spend), f"{case_name}: streamed usage.cost {seen} != spend {spend}"
def assert_json_usage_cost(case_name: str, endpoint: str, response: httpx.Response, spend: float) -> None:
"""A JSON reply's usage.cost is what the row billed, the upstream's own figure included, so a client and the
spend log never disagree on a priced deployment."""
if endpoint not in STREAMED_COST_ENDPOINTS:
return
usage: Final = JSON_OBJECT.validate_json(response.text).get("usage")
cost: Final = usage.get("cost") if isinstance(usage, dict) else None
if not isinstance(cost, (int, float)) or isinstance(cost, bool):
return
assert approx_equal(float(cost), spend), f"{case_name}: usage.cost {cost} != spend {spend}"
def assert_breakdown(
case_name: str,
response_content_type: str,
@ -172,8 +160,6 @@ def assert_exact(
assert_total_is_sum_of_components(row, breakdown, case_name)
if response is not None and response_content_type == "text/event-stream":
assert_streamed_usage_cost(case_name, response.request.url.path, response.text, expected.spend)
if response is not None and response_content_type == "application/json":
assert_json_usage_cost(case_name, response.request.url.path, response, expected.spend)
def assert_recount(
@ -215,5 +201,3 @@ def assert_recount(
assert_total_is_sum_of_components(row, row.breakdown, case_name)
if response is not None and response_content_type == "text/event-stream":
assert_streamed_usage_cost(case_name, response.request.url.path, response.text, row.spend)
if response is not None and response_content_type == "application/json":
assert_json_usage_cost(case_name, response.request.url.path, response, row.spend)

View file

@ -15,7 +15,6 @@ import pytest
# Add the project root to Python path
import litellm
from litellm.cost_calculator import cost_per_token as billed_cost_per_token
from litellm.llms.perplexity.cost_calculator import (
cost_per_token as perplexity_cost_per_token,
)
@ -110,13 +109,13 @@ class TestPerplexityCostCalculator:
"total_cost": 0.008,
}
prompt_cost, completion_cost = billed_cost_per_token(
model="perplexity/sonar-pro", custom_llm_provider="perplexity", usage_object=usage
)
prompt_cost, completion_cost = perplexity_cost_per_token(model="sonar-pro", usage=usage)
# When Perplexity provides total_cost, that total is the bill
assert math.isclose(prompt_cost + completion_cost, 0.008, rel_tol=1e-10)
assert prompt_cost > 0.0 and completion_cost > 0.0
# When Perplexity provides total_cost, we use it directly
# prompt_cost should be 0, completion_cost should be total_cost
assert prompt_cost == 0.0
assert completion_cost == 0.008
assert prompt_cost + completion_cost == 0.008
def test_uses_perplexity_provided_cost_when_normalized_to_float(self):
"""
@ -136,11 +135,10 @@ class TestPerplexityCostCalculator:
usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150)
usage.cost = 0.008
prompt_cost, completion_cost = billed_cost_per_token(
model="perplexity/sonar-pro", custom_llm_provider="perplexity", usage_object=usage
)
prompt_cost, completion_cost = perplexity_cost_per_token(model="sonar-pro", usage=usage)
assert math.isclose(prompt_cost + completion_cost, 0.008, rel_tol=1e-10)
assert prompt_cost == 0.0
assert completion_cost == 0.008
OFF_PEAK_MODEL = "sonar-off-peak-test"
OFF_PEAK_WINDOW = "14:00-00:00"
@ -228,9 +226,7 @@ class TestPerplexityCostCalculator:
usage = Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200)
usage.cost = {"total_cost": 0.00501}
prompt_cost, completion_cost = billed_cost_per_token(
model=self.OFF_PEAK_MODEL, custom_llm_provider="perplexity", usage_object=usage
)
prompt_cost, completion_cost = perplexity_cost_per_token(model=self.OFF_PEAK_MODEL, usage=usage)
assert math.isclose(prompt_cost + completion_cost, 0.00501, rel_tol=1e-10)
assert prompt_cost > 0.0 and completion_cost > 0.0
assert prompt_cost == 0.0
assert completion_cost == 0.00501

View file

@ -11,7 +11,7 @@ from unittest.mock import MagicMock, Mock
import httpx
from litellm.cost_calculator import cost_per_token
from litellm.llms.xai.cost_calculator import cost_per_token
from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig
from litellm.responses.utils import ResponseAPILoggingUtils
from litellm.types.llms.openai import (
@ -498,7 +498,7 @@ class TestXAIResponsesReportedCost:
assert usage.cost == 0.0037756
chat_usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(usage)
assert cost_per_token(model="grok-4-latest", custom_llm_provider="xai", usage_object=chat_usage) == (0.0, 0.0037756)
assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756)
def test_streamed_reported_cost_reaches_the_cost_calculator(self):
event = XAIResponsesAPIConfig().transform_streaming_response(
@ -520,7 +520,7 @@ class TestXAIResponsesReportedCost:
assert isinstance(event, ResponseCompletedEvent)
chat_usage = ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(event.response.usage)
assert cost_per_token(model="grok-4-latest", custom_llm_provider="xai", usage_object=chat_usage) == (0.0, 0.0037756)
assert cost_per_token(model="grok-4-latest", usage=chat_usage) == (0.0, 0.0037756)
def test_usage_without_a_reported_cost_is_left_alone(self):
usage = self._transformed_usage({"input_tokens": 100, "output_tokens": 200, "total_tokens": 300})

View file

@ -13,7 +13,7 @@ from litellm.llms.xai.chat.transformation import (
XAIChatCompletionStreamingHandler,
XAIChatConfig,
)
from litellm.cost_calculator import cost_per_token
from litellm.llms.xai.cost_calculator import cost_per_token
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
ModelResponse,
@ -433,7 +433,7 @@ class TestXAIReportedCost:
)
assert usage.cost == 0.0037756
assert cost_per_token(model="grok-4-latest", custom_llm_provider="xai", usage_object=usage) == (0.0, 0.0037756)
assert cost_per_token(model="grok-4-latest", usage=usage) == (0.0, 0.0037756)
def test_usage_without_a_reported_cost_is_left_alone(self):
usage = self._transformed_usage({"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300})
@ -481,7 +481,7 @@ class TestXAIReportedCost:
assembled = litellm.stream_chunk_builder(chunks=[parsed])
assert assembled.usage.cost == 0.0037756
assert cost_per_token(model="grok-4-latest", custom_llm_provider="xai", usage_object=assembled.usage) == (
assert cost_per_token(model="grok-4-latest", usage=assembled.usage) == (
0.0,
0.0037756,
)

View file

@ -6,7 +6,6 @@ import math
import os
import litellm
from litellm.cost_calculator import cost_per_token as billed_cost_per_token
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
StandardBuiltInToolCostTracking,
)
@ -116,9 +115,7 @@ class TestXAICostCalculator:
model="grok-4.7", usage=Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300)
)
prompt_cost, completion_cost = billed_cost_per_token(
model="xai/grok-4.7", custom_llm_provider="xai", usage_object=usage
)
prompt_cost, completion_cost = cost_per_token(model="grok-4.7", usage=usage)
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)
@ -168,10 +165,7 @@ class TestXAICostCalculator:
"""A reported zero is a real answer, not a missing value."""
usage = Usage(prompt_tokens=100, completion_tokens=200, total_tokens=300, cost=0.0)
assert billed_cost_per_token(model="xai/grok-4-latest", custom_llm_provider="xai", usage_object=usage) == (
0.0,
0.0,
)
assert cost_per_token(model="grok-4-latest", usage=usage) == (0.0, 0.0)
def test_custom_pricing_beats_the_reported_cost(self):
response = ModelResponse(

View file

@ -1,200 +0,0 @@
from __future__ import annotations
from typing import Final
import pytest
import litellm
from litellm.proxy.usage_cost_egress import with_client_stream_usage_cost, with_client_usage_cost
from litellm.types.llms.openai import (
ResponseAPIUsage,
ResponseCompletedEvent,
ResponseIncompleteEvent,
ResponsesAPIResponse,
ResponsesAPIStreamEvents,
)
from litellm.types.utils import Choices, Message, ModelResponse, ModelResponseStream, StreamingChoices, Usage
UPSTREAM_COST: Final = 0.037
OPERATOR_BILL: Final = 100 * 2e-6 + 20 * 6e-6
def _chat_reply(upstream_cost: float | None) -> ModelResponse:
reply: Final = ModelResponse(
id="chatcmpl-egress",
model="omni-model",
choices=[Choices(index=0, message=Message(role="assistant", content="hi"), finish_reason="stop")],
usage=Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120, cost=upstream_cost),
)
reply._hidden_params["response_cost"] = OPERATOR_BILL
return reply
def _responses_reply(upstream_cost: float | None) -> ResponsesAPIResponse:
return ResponsesAPIResponse(
id="resp_egress",
created_at=1,
status="completed",
model="omni-model",
object="response",
output=[],
usage=ResponseAPIUsage(input_tokens=100, output_tokens=20, total_tokens=120, cost=upstream_cost),
)
def _usage_chunk(upstream_cost: float | None) -> ModelResponseStream:
return ModelResponseStream(
id="chatcmpl-egress",
model="omni-model",
choices=[],
usage=Usage(prompt_tokens=100, completion_tokens=20, total_tokens=120, cost=upstream_cost),
)
def _bill(_reply: object) -> float | None:
return OPERATOR_BILL
def _unpriceable(_reply: object) -> float | None:
return None
@pytest.mark.parametrize("reply", [_chat_reply(UPSTREAM_COST), _responses_reply(UPSTREAM_COST)])
def test_a_reply_whose_upstream_reported_a_cost_tells_the_client_what_litellm_billed(reply: object) -> None:
billed: Final = with_client_usage_cost(reply, OPERATOR_BILL)
assert billed is not reply, "the pipeline's own object keeps the upstream figure"
assert getattr(getattr(billed, "usage", None), "cost", None) == pytest.approx(OPERATOR_BILL)
assert getattr(getattr(reply, "usage", None), "cost", None) == UPSTREAM_COST
def test_the_billed_copy_of_a_chat_reply_keeps_its_hidden_params() -> None:
reply: Final = _chat_reply(UPSTREAM_COST)
billed: Final = with_client_usage_cost(reply, OPERATOR_BILL)
assert isinstance(billed, ModelResponse)
assert billed._hidden_params["response_cost"] == pytest.approx(OPERATOR_BILL)
assert billed.model_dump()["usage"]["cost"] == pytest.approx(OPERATOR_BILL)
@pytest.mark.parametrize("reply", [_chat_reply(None), _responses_reply(None)])
def test_a_reply_without_an_upstream_cost_goes_out_as_the_upstream_sent_it(
monkeypatch: pytest.MonkeyPatch, reply: object
) -> None:
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
assert with_client_usage_cost(reply, OPERATOR_BILL) is reply
@pytest.mark.parametrize("reply", [_chat_reply(None), _responses_reply(None)])
def test_include_cost_in_streaming_usage_puts_the_bill_on_a_reply_without_an_upstream_cost(
monkeypatch: pytest.MonkeyPatch, reply: object
) -> None:
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
billed: Final = with_client_usage_cost(reply, OPERATOR_BILL)
assert getattr(getattr(billed, "usage", None), "cost", None) == pytest.approx(OPERATOR_BILL)
def test_a_zero_bill_replaces_the_upstream_cost() -> None:
billed: Final = with_client_usage_cost(_chat_reply(UPSTREAM_COST), 0.0)
assert getattr(getattr(billed, "usage", None), "cost", None) == 0.0
@pytest.mark.parametrize("bill", [None, -0.001, float("nan"), float("inf"), True, "0.1"])
def test_an_unpriceable_reply_keeps_the_upstream_cost(bill: object) -> None:
reply: Final = _chat_reply(UPSTREAM_COST)
assert with_client_usage_cost(reply, bill) is reply
assert reply.usage is not None and reply.usage.cost == UPSTREAM_COST
def test_a_reply_that_is_neither_chat_nor_responses_passes_through() -> None:
reply: Final = {"usage": {"cost": UPSTREAM_COST}}
assert with_client_usage_cost(reply, OPERATOR_BILL) is reply
def test_the_streamed_usage_chunk_tells_the_client_what_litellm_billed() -> None:
chunk: Final = _usage_chunk(UPSTREAM_COST)
billed: Final = with_client_stream_usage_cost(chunk, _bill)
assert isinstance(billed, ModelResponseStream) and billed is not chunk
assert billed.usage is not None and billed.usage.cost == pytest.approx(OPERATOR_BILL)
assert chunk.usage is not None and chunk.usage.cost == UPSTREAM_COST
def test_a_content_chunk_is_not_priced() -> None:
chunk: Final = ModelResponseStream(
id="chatcmpl-egress", model="omni-model", choices=[StreamingChoices(index=0, delta={"content": "hi"})]
)
priced: list[object] = []
def price(reply: object) -> float | None:
priced.append(reply)
return OPERATOR_BILL
assert with_client_stream_usage_cost(chunk, price) is chunk
assert priced == []
def test_a_streamed_usage_chunk_without_an_upstream_cost_goes_out_as_sent(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
chunk: Final = _usage_chunk(None)
assert with_client_stream_usage_cost(chunk, _bill) is chunk
def test_include_cost_in_streaming_usage_puts_the_bill_on_the_streamed_usage_chunk(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True)
billed: Final = with_client_stream_usage_cost(_usage_chunk(None), _bill)
assert isinstance(billed, ModelResponseStream)
assert billed.usage is not None and billed.usage.cost == pytest.approx(OPERATOR_BILL)
def test_the_stream_pricer_sees_the_upstream_cost_so_the_bill_matches_the_spend_row() -> None:
seen: list[object] = []
def price(reply: object) -> float | None:
seen.append(getattr(getattr(reply, "usage", None), "cost", None))
return OPERATOR_BILL
with_client_stream_usage_cost(_usage_chunk(UPSTREAM_COST), price)
assert seen == [UPSTREAM_COST]
@pytest.mark.parametrize("upstream_cost", [UPSTREAM_COST, None])
@pytest.mark.parametrize(
"event_type", [ResponsesAPIStreamEvents.RESPONSE_COMPLETED, ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE]
)
def test_a_terminal_responses_event_carries_the_bill_with_or_without_an_upstream_cost(
monkeypatch: pytest.MonkeyPatch, upstream_cost: float | None, event_type: ResponsesAPIStreamEvents
) -> None:
monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False)
event_cls: Final = (
ResponseCompletedEvent if event_type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED else ResponseIncompleteEvent
)
event: Final = event_cls(type=event_type, response=_responses_reply(upstream_cost))
billed: Final = with_client_stream_usage_cost(event, _bill)
assert billed is not event
assert getattr(billed, "response").usage.cost == pytest.approx(OPERATOR_BILL)
assert event.response.usage is not None and event.response.usage.cost == upstream_cost
def test_an_unpriceable_terminal_responses_event_goes_out_as_sent() -> None:
event: Final = ResponseCompletedEvent(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, response=_responses_reply(None)
)
assert with_client_stream_usage_cost(event, _unpriceable) is event