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

@@ -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
 

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:

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(