From 29e2390a56638e5f33203b829b15f9bb68f1bcd3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 08:37:35 -0700 Subject: [PATCH] fix(cost): bill the requested model's listed rates over a served size-bucket guess A served Fireworks or Together AI name that is not in the cost map but names its parameter count, like a display name with a size in it, now bills at the requested model's listed rates instead of the provider's size-bucket price. The bucket price still applies when the requested model has no rates, as on main --- litellm/cost_calculator.py | 26 +------- litellm/llms/fireworks_ai/cost_calculator.py | 4 +- tests/unit/test_cost_calculator.py | 65 +++++++++++++++----- 3 files changed, 50 insertions(+), 45 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index ffa054d4a08..e80b9bfde39 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -72,10 +72,6 @@ from litellm.llms.databricks.cost_calculator import ( from litellm.llms.deepseek.cost_calculator import ( cost_per_token as deepseek_cost_per_token, ) -from litellm.llms.fireworks_ai.cost_calculator import ( - FIREWORKS_AI_DEFAULT_PRICING_MODEL, - get_base_model_for_pricing, -) from litellm.llms.fireworks_ai.cost_calculator import ( cost_per_token as fireworks_ai_cost_per_token, ) @@ -1054,26 +1050,6 @@ def _resolves_to_rates(model: str, custom_llm_provider: str | None) -> bool: return _cost_map_entry_prices_anything(entry if fill_missing is None else {**fill_missing, **entry}) -def _size_pricing_bucket(model: str, custom_llm_provider: str | None) -> str | None: - match custom_llm_provider: - case "fireworks_ai": - bucket: Final = get_base_model_for_pricing(model) - return None if bucket == FIREWORKS_AI_DEFAULT_PRICING_MODEL else bucket - case "together_ai": - return get_model_params_and_category(model, call_type=CallTypes.completion) - case _: - return None - - -def _size_bucket_declares_rates(model: str, custom_llm_provider: str | None) -> bool: - bucket: Final = _size_pricing_bucket(model, custom_llm_provider) - return bucket is not None and cost_map_key_prices_anything(bucket) - - -def _served_model_has_rates(model: str, custom_llm_provider: str | None) -> bool: - return _resolves_to_rates(model, custom_llm_provider) or _size_bucket_declares_rates(model, custom_llm_provider) - - def _served_or_priced_requested_model( served_model: str, requested_model: str | None, custom_llm_provider: str | None, priced_region: str | None ) -> str: @@ -1084,7 +1060,7 @@ def _served_or_priced_requested_model( if served in (requested, requested.partition("/")[2]): return served pricing_provider: Final = _get_provider_for_cost_calc(requested, custom_llm_provider) - if _served_model_has_rates(served, pricing_provider): + if _resolves_to_rates(served, pricing_provider): return served return requested if _resolves_to_rates(requested, pricing_provider) else served diff --git a/litellm/llms/fireworks_ai/cost_calculator.py b/litellm/llms/fireworks_ai/cost_calculator.py index dda42b52f06..d42b4ddbbf7 100644 --- a/litellm/llms/fireworks_ai/cost_calculator.py +++ b/litellm/llms/fireworks_ai/cost_calculator.py @@ -15,8 +15,6 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_toke from litellm.types.utils import ModelInfo, Usage from litellm.utils import get_model_info -FIREWORKS_AI_DEFAULT_PRICING_MODEL: Final = "fireworks-ai-default" - # Extract the number of billion parameters from the model name # only used for together_computer LLMs @@ -55,7 +53,7 @@ def get_base_model_for_pricing(model_name: str) -> str: return "fireworks-ai-above-16b" # If no matches, return the original model_name - return FIREWORKS_AI_DEFAULT_PRICING_MODEL + return "fireworks-ai-default" def _resolve_model_info(model: str) -> ModelInfo: diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index a4ab108955f..f10baccd1cb 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -637,19 +637,35 @@ def test_completion_cost_keeps_served_row_rates_filled_by_a_rule_over_the_reques assert cost == pytest.approx(_token_rates_cost(filled_rates, usage)) -@pytest.mark.parametrize( - ("custom_llm_provider", "requested_model", "served_model", "bucket"), - [ - ( - "fireworks_ai", - "fireworks_ai/accounts/fireworks/models/sol-bucket-requested-test", - "accounts/fireworks/models/sol-bucket-405b-test", - "fireworks-ai-above-16b", - ), - ("together_ai", "together_ai/sol-bucket-requested-test", "meta-llama/Sol-Bucket-8B-Test", "together-ai-4.1b-8b"), - ], +_SIZE_BUCKET_CASES: Final = ( + pytest.param( + "fireworks_ai", + "fireworks_ai/accounts/fireworks/models/sol-bucket-requested-test", + "Sol Bucket 405B Display", + "fireworks-ai-above-16b", + id="fireworks_ai", + ), + pytest.param( + "together_ai", + "together_ai/sol-bucket-requested-test", + "meta-llama/Sol-Bucket-8B-Test", + "together-ai-4.1b-8b", + id="together_ai", + ), ) -def test_completion_cost_keeps_served_size_bucket_rates_over_the_requested_model( + + +def _size_bucket_response(served_model: str, requested_model: str, custom_llm_provider: str) -> ModelResponse: + response: Final = ModelResponse(model=served_model, choices=[], usage=_SIZE_BUCKET_USAGE) + response._hidden_params.update({"litellm_model_name": requested_model, "custom_llm_provider": custom_llm_provider}) + return response + + +_SIZE_BUCKET_USAGE: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + +@pytest.mark.parametrize(("custom_llm_provider", "requested_model", "served_model", "bucket"), _SIZE_BUCKET_CASES) +def test_completion_cost_bills_requested_listed_rates_over_a_served_size_bucket( _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, custom_llm_provider: str, @@ -667,13 +683,28 @@ def test_completion_cost_keeps_served_size_bucket_rates_over_the_requested_model "output_cost_per_token": 7e-6, }, ) - usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) - response: Final = ModelResponse(model=served_model, choices=[], usage=usage) - response._hidden_params.update({"litellm_model_name": requested_model, "custom_llm_provider": custom_llm_provider}) - cost: Final = completion_cost(completion_response=response) + cost: Final = completion_cost( + completion_response=_size_bucket_response(served_model, requested_model, custom_llm_provider) + ) - assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[bucket], usage)) + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[requested_model], _SIZE_BUCKET_USAGE)) + assert cost != pytest.approx(_token_rates_cost(litellm.model_cost[bucket], _SIZE_BUCKET_USAGE)) + + +@pytest.mark.parametrize(("custom_llm_provider", "requested_model", "served_model", "bucket"), _SIZE_BUCKET_CASES) +def test_completion_cost_keeps_served_size_bucket_rates_when_the_requested_model_is_unpriced( + _local_model_cost_map: None, + custom_llm_provider: str, + requested_model: str, + served_model: str, + bucket: str, +) -> None: + cost: Final = completion_cost( + completion_response=_size_bucket_response(served_model, requested_model, custom_llm_provider) + ) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[bucket], _SIZE_BUCKET_USAGE)) def test_completion_cost_bills_a_requested_model_priced_after_an_unpriced_lookup(