diff --git a/headroom/proxy/savings_tracker.py b/headroom/proxy/savings_tracker.py index fbf36822f..5f5bfa599 100644 --- a/headroom/proxy/savings_tracker.py +++ b/headroom/proxy/savings_tracker.py @@ -238,7 +238,12 @@ def _estimate_output_savings_usd(model: str, tokens_saved: int) -> float: resolved = _resolve_litellm_model(model) info = litellm.model_cost.get(resolved, {}) output_cost_per_token = info.get("output_cost_per_token") - if not output_cost_per_token: + # Distinguish "price unknown" (missing key -> fall back to the estimate) + # from a model that is legitimately free (output_cost_per_token == 0.0). + # `if not ...` treated a real 0.0 as unavailable and billed the fallback + # rate -> phantom output savings for a model that costs nothing. Mirrors + # the fix already applied to `_estimate_compression_savings_usd`. + if output_cost_per_token is None: raise RuntimeError("output cost unavailable") return float(tokens_saved) * float(output_cost_per_token) except Exception: diff --git a/tests/test_savings_tracker_zero_price.py b/tests/test_savings_tracker_zero_price.py index 3fc3e91c3..8ac1975e2 100644 --- a/tests/test_savings_tracker_zero_price.py +++ b/tests/test_savings_tracker_zero_price.py @@ -15,8 +15,10 @@ import types from headroom.proxy import savings_tracker as st from headroom.proxy.savings_tracker import ( DEFAULT_FALLBACK_INPUT_COST_PER_TOKEN, + DEFAULT_FALLBACK_OUTPUT_COST_PER_TOKEN, _estimate_compression_savings_usd, _estimate_input_cost_usd, + _estimate_output_savings_usd, ) @@ -61,3 +63,30 @@ def test_input_cost_zero_for_free_model(monkeypatch): lambda: _fake_litellm({"free-model": {"input_cost_per_token": 0.0}}), ) assert _estimate_input_cost_usd("free-model", 500_000) == 0.0 + + +def test_output_savings_zero_for_free_model(monkeypatch): + # output_cost_per_token == 0.0 (free model) must yield $0, not the fallback. + monkeypatch.setattr( + st, + "_get_litellm_module", + lambda: _fake_litellm({"free-model": {"output_cost_per_token": 0.0}}), + ) + assert _estimate_output_savings_usd("free-model", 1_000_000) == 0.0 + + +def test_output_savings_falls_back_for_unknown_model(monkeypatch): + # Model absent from litellm → output_cost_per_token is None → fall back. + monkeypatch.setattr(st, "_get_litellm_module", lambda: _fake_litellm({})) + got = _estimate_output_savings_usd("unknown-model", 1_000_000) + assert got == 1_000_000 * DEFAULT_FALLBACK_OUTPUT_COST_PER_TOKEN + + +def test_output_savings_uses_real_price_for_paid_model(monkeypatch): + price = 15.0 / 1_000_000 + monkeypatch.setattr( + st, + "_get_litellm_module", + lambda: _fake_litellm({"paid-model": {"output_cost_per_token": price}}), + ) + assert _estimate_output_savings_usd("paid-model", 1_000_000) == 1_000_000 * price