fix(kyma): let SDK custom_cost_per_token and custom_cost_per_second override Kyma's usage.cost

This commit is contained in:
mateo-berri 2026-10-28 19:45:52 -07:00
parent 6b335eba7a
commit 2fa7dfc6ba
2 changed files with 37 additions and 1 deletions

View file

@ -1649,6 +1649,9 @@ def completion_cost(
)
explicit_pricing: Final = custom_pricing is True or base_model is not None
caller_set_prices: Final = (
explicit_pricing or custom_cost_per_token is not None or custom_cost_per_second is not None
)
custom_priced_request_model: Final = model if custom_pricing is True else None
configured_cost_per_query: Final = (
_configured_search_cost_per_query(litellm_logging_obj) if custom_pricing else None
@ -2019,7 +2022,7 @@ def completion_cost(
)
model_for_discount = custom_priced_request_model or model
if (
not explicit_pricing
not caller_set_prices
and (reported_cost := billed_usage_cost(priced_response, custom_llm_provider)) is not None
):
return _adjusted_provider_billed_cost(

View file

@ -292,6 +292,39 @@ def test_sdk_completion_cost_prices_kyma_at_the_cost_kyma_reports():
assert litellm.completion_cost(completion_response=response) == pytest.approx(reported_cost)
def _kyma_sdk_response(usage: dict[str, float]) -> litellm.ModelResponse:
with respx.mock() as upstream:
upstream.post(_CHAT_URL).respond(200, json=_kyma_chat_body(usage))
response: Final = litellm.completion(model="kyma/sonar", messages=[{"role": "user", "content": "ping"}], api_key="k")
assert isinstance(response, litellm.ModelResponse)
return response
_KYMA_TOKEN_USAGE: Final = {"prompt_tokens": 4, "completion_tokens": 3, "total_tokens": 7}
_KYMA_REPORTED_COST: Final = 0.00721
def test_sdk_completion_cost_uses_the_callers_token_prices_over_the_cost_kyma_reports():
caller_prices: Final = {"input_cost_per_token": 0.5, "output_cost_per_token": 2.0}
with_reported_cost: Final = _kyma_sdk_response({**_KYMA_TOKEN_USAGE, "cost": _KYMA_REPORTED_COST})
without_reported_cost: Final = _kyma_sdk_response(_KYMA_TOKEN_USAGE)
assert litellm.completion_cost(
completion_response=with_reported_cost, custom_cost_per_token=caller_prices
) == pytest.approx(litellm.completion_cost(completion_response=without_reported_cost, custom_cost_per_token=caller_prices))
assert litellm.completion_cost(
completion_response=with_reported_cost, custom_cost_per_token=caller_prices
) != pytest.approx(_KYMA_REPORTED_COST)
def test_sdk_completion_cost_uses_the_callers_per_second_price_over_the_cost_kyma_reports():
response: Final = _kyma_sdk_response({**_KYMA_TOKEN_USAGE, "cost": _KYMA_REPORTED_COST})
assert litellm.completion_cost(completion_response=response, custom_cost_per_second=50.0) == pytest.approx(
50.0 * response._response_ms / 1000
)
def _with_kyma_discount_and_margin(monkeypatch: pytest.MonkeyPatch) -> float:
monkeypatch.setattr(litellm, "cost_discount_config", {"kyma": 0.5})
monkeypatch.setattr(litellm, "cost_margin_config", {"kyma": 0.1})