From 9c6954c00a3011a23fc194234cc7aad8d09b54c6 Mon Sep 17 00:00:00 2001 From: kerry Date: Tue, 27 Oct 2026 21:39:26 +0000 Subject: [PATCH] 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> --- litellm/cost_calculator.py | 7 + litellm/llms/perplexity/cost_calculator.py | 16 ++ litellm/llms/xai/cost_calculator.py | 18 +- litellm/proxy/common_request_processing.py | 11 +- litellm/proxy/proxy_server.py | 20 +- litellm/proxy/usage_cost_egress.py | 77 ------- .../cost_calculation/assertions.py | 16 -- .../test_perplexity_cost_calculator.py | 28 ++- .../test_xai_responses_transformation.py | 6 +- .../llms/xai/test_xai_chat_transformation.py | 6 +- .../unit/llms/xai/test_xai_cost_calculator.py | 10 +- tests/unit/proxy/test_usage_cost_egress.py | 200 ------------------ 12 files changed, 65 insertions(+), 350 deletions(-) delete mode 100644 litellm/proxy/usage_cost_egress.py delete mode 100644 tests/unit/proxy/test_usage_cost_egress.py diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index fc305c84ac5..d2fdec8688e 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -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)) diff --git a/litellm/llms/perplexity/cost_calculator.py b/litellm/llms/perplexity/cost_calculator.py index a6825eacfad..2d762adf958 100644 --- a/litellm/llms/perplexity/cost_calculator.py +++ b/litellm/llms/perplexity/cost_calculator.py @@ -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") diff --git a/litellm/llms/xai/cost_calculator.py b/litellm/llms/xai/cost_calculator.py index b62ed39c4fa..aec535c2086 100644 --- a/litellm/llms/xai/cost_calculator.py +++ b/litellm/llms/xai/cost_calculator.py @@ -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) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index c5843da28d9..3b9c17baa02 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -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, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 5c7388f7861..bc79becf7bb 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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: diff --git a/litellm/proxy/usage_cost_egress.py b/litellm/proxy/usage_cost_egress.py deleted file mode 100644 index 820fce85974..00000000000 --- a/litellm/proxy/usage_cost_egress.py +++ /dev/null @@ -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)}) diff --git a/tests/integration/cost_calculation/assertions.py b/tests/integration/cost_calculation/assertions.py index f37bb32e3b3..6e2af74d1a2 100644 --- a/tests/integration/cost_calculation/assertions.py +++ b/tests/integration/cost_calculation/assertions.py @@ -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) diff --git a/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py index cd25b7a9c3d..83c71479311 100644 --- a/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py @@ -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 diff --git a/tests/unit/llms/xai/responses/test_xai_responses_transformation.py b/tests/unit/llms/xai/responses/test_xai_responses_transformation.py index 95ff98ac024..188ca453364 100644 --- a/tests/unit/llms/xai/responses/test_xai_responses_transformation.py +++ b/tests/unit/llms/xai/responses/test_xai_responses_transformation.py @@ -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}) diff --git a/tests/unit/llms/xai/test_xai_chat_transformation.py b/tests/unit/llms/xai/test_xai_chat_transformation.py index 3aa35759662..69e620dbc8a 100644 --- a/tests/unit/llms/xai/test_xai_chat_transformation.py +++ b/tests/unit/llms/xai/test_xai_chat_transformation.py @@ -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, ) diff --git a/tests/unit/llms/xai/test_xai_cost_calculator.py b/tests/unit/llms/xai/test_xai_cost_calculator.py index c1d2ee5f63d..5219dba42ba 100644 --- a/tests/unit/llms/xai/test_xai_cost_calculator.py +++ b/tests/unit/llms/xai/test_xai_cost_calculator.py @@ -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( diff --git a/tests/unit/proxy/test_usage_cost_egress.py b/tests/unit/proxy/test_usage_cost_egress.py deleted file mode 100644 index 3c62ddb9308..00000000000 --- a/tests/unit/proxy/test_usage_cost_egress.py +++ /dev/null @@ -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