fix(cost): keep the served model's filled-in and size-bucket rates, and cache misses

A served row priced only by a fill_missing_for_providers rule bills those
rates on main, but the price check read only the raw row and billed the
requested alias instead. It now merges the rule's fill-in keys through the
helper model info lookup uses.

A served Fireworks or Together AI id that names its parameter count keeps its
size-bucket price, as on main, instead of switching to the requested row.

The served model is checked first, cost map misses are cached and cleared with
the model info cache, and the rules are read only for a name that resolved
through one. A priced snapshot behind an Azure deployment name takes about
19 us of CPU time against 7 us on the merge base, down from about 300 us.

litellm/cost_calculator.py

@@ -72,6 +72,10 @@ 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,
 )
@@ -152,7 +156,9 @@ from litellm.utils import (
     ProviderConfigManager,
     TextCompletionResponse,
     TranscriptionResponse,
+    cached_cost_map_key,
     cached_get_model_info_helper,
+    mapped_model_fill_missing,
     model_info_with_pricing_overrides,
     runtime_model_cost_registrations,
     token_counter,
@@ -1032,25 +1038,40 @@ def _cost_map_model_name(model: str, custom_llm_provider: str | None, region_nam
     return _strip_unregistered_leading_segments(f"{provider_prefix}/{model}", region_name)
 
 
-def _cost_map_row_key(model: str, custom_llm_provider: str | None) -> str | None:
-    try:
-        return cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)["key"]
-    except Exception:
-        return None
-
-
-def _row_declares_rates(model: str, custom_llm_provider: str | None) -> bool:
-    key: Final = _cost_map_row_key(model, custom_llm_provider)
-    return key is not None and cost_map_key_prices_anything(key)
-
-
 def _rule_declares_rates(model: str, custom_llm_provider: str | None) -> bool:
     rule_info: Final = unmapped_model_generalization(model, custom_llm_provider)
     return rule_info is not None and _cost_map_entry_prices_anything(rule_info)
 
 
 def _resolves_to_rates(model: str, custom_llm_provider: str | None) -> bool:
-    return _row_declares_rates(model, custom_llm_provider) or _rule_declares_rates(model, custom_llm_provider)
+    key: Final = cached_cost_map_key(model, custom_llm_provider)
+    if key is None:
+        return False
+    entry: Final = _raw_cost_map_entry(key)
+    if entry is None:
+        return _rule_declares_rates(model, custom_llm_provider)
+    fill_missing: Final = mapped_model_fill_missing(key, entry)
+    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(
@@ -1063,9 +1084,9 @@ 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 not _resolves_to_rates(requested, pricing_provider):
+    if _served_model_has_rates(served, pricing_provider):
         return served
-    return served if _resolves_to_rates(served, pricing_provider) else requested
+    return requested if _resolves_to_rates(requested, pricing_provider) else served
 
 
 def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> str:

litellm/llms/fireworks_ai/cost_calculator.py

@@ -15,6 +15,8 @@ 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
@@ -53,7 +55,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"
+    return FIREWORKS_AI_DEFAULT_PRICING_MODEL
 
 
 def _resolve_model_info(model: str) -> ModelInfo:

litellm/utils.py

@@ -277,7 +277,7 @@ from litellm.types.utils import (
 )
 
 _CALL_TYPE_ENUM_MAP: Final[dict] = {ct.value: ct for ct in CallTypes}
-_BACKFILL_MODES: Final = frozenset({"chat", "responses"})
+_BACKFILL_MODES: Final[frozenset[str]] = frozenset({"chat", "responses"})
 
 # +-----------------------------------------------+
 # |                                               |
@@ -6221,6 +6221,7 @@ def _invalidate_model_cost_lowercase_map() -> None:
     # Clear LRU caches that depend on model_cost data
     _cached_get_model_info.cache_clear()
     cached_get_model_info_helper.cache_clear()
+    cached_cost_map_key.cache_clear()
 
 
 def _rebuild_model_cost_lowercase_map() -> dict[str, str]:
@@ -6446,6 +6447,14 @@ def unmapped_model_generalization(model: str, custom_llm_provider: str | None) -
     return generalization[1] if generalization is not None else None
 
 
+def mapped_model_fill_missing(key: str, model_info: Mapping[str, object]) -> Mapping[str, object] | None:
+    mode: Final = model_info.get("mode", "chat")
+    provider: Final = model_info.get("litellm_provider", "")
+    if not isinstance(mode, str) or mode not in _BACKFILL_MODES or not isinstance(provider, str):
+        return None
+    return match_fill_missing_generalizations(key, provider)
+
+
 def _strip_mantle_region_prefix(model: str) -> str:
     from litellm.llms.bedrock_mantle.common_utils import split_mantle_region_prefix
 
@@ -6574,6 +6583,14 @@ def cached_get_model_info_helper(
 _cached_get_model_info_helper = cached_get_model_info_helper
 
 
+@lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)
+def cached_cost_map_key(model: str, custom_llm_provider: str | None) -> str | None:
+    try:
+        return cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)["key"]
+    except Exception:
+        return None
+
+
 def get_provider_info(model: str, custom_llm_provider: str | None) -> ProviderSpecificModelInfo | None:
     ## PROVIDER-SPECIFIC INFORMATION
     # if custom_llm_provider == "predibase":
@@ -6729,8 +6746,8 @@ def get_model_info_helper(
             key: str | None = lookup[0]
             _model_info: dict[str, Any] | None = lookup[1]
 
-            if _model_info is not None and key is not None and _model_info.get("mode", "chat") in _BACKFILL_MODES:
-                fill_missing: Final = match_fill_missing_generalizations(key, _model_info.get("litellm_provider", ""))
+            if _model_info is not None and key is not None:
+                fill_missing: Final = mapped_model_fill_missing(key, _model_info)
                 if fill_missing is not None:
                     _model_info = {
                         **{k: v for k, v in fill_missing.items() if k not in _model_info},

tests/unit/test_cost_calculator.py

@@ -37,7 +37,7 @@ from litellm.types.llms.xai import RealtimeMeteredUsage
 from litellm.types.rerank import RerankResponse
 from litellm.types.utils import CacheCreationTokenDetails, CallTypes, Choices, CompletionTokensDetailsWrapper, EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails, LiteLLMRealtimeStreamLoggingObject, Message, ModelInfo, ModelResponse, PromptTokensDetails, PromptTokensDetailsWrapper, Usage
 from litellm.types.videos.main import VideoObject
-from litellm.utils import supports_prompt_caching
+from litellm.utils import _invalidate_model_cost_lowercase_map, supports_prompt_caching
 import os
 import threading
 from unittest.mock import MagicMock, patch
@@ -605,6 +605,101 @@ def test_completion_cost_bills_requested_model_regional_row_when_response_names_
     assert cost == pytest.approx(_token_rates_cost(litellm.model_cost["bedrock/eu-central-1/sol-regional-test"], usage))
 
 
+def test_completion_cost_keeps_served_row_rates_filled_by_a_rule_over_the_requested_alias(
+    _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    monkeypatch.setitem(litellm.model_cost, "sol-fill-snapshot-test", {"litellm_provider": "openai", "mode": "chat"})
+    monkeypatch.setitem(
+        litellm.model_cost,
+        "sol-fill-alias-test",
+        {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6},
+    )
+    filled_rates: Final = {"input_cost_per_token": 5e-6, "output_cost_per_token": 11e-6}
+    usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
+    response: Final = ModelResponse(model="sol-fill-snapshot-test", choices=[], usage=usage)
+    response._hidden_params.update({"litellm_model_name": "openai/sol-fill-alias-test", "custom_llm_provider": "openai"})
+    previous_rules: Final = list(get_fallback_generalization_rules())
+    set_fallback_generalizations(
+        [
+            {
+                "name": "sol-fill-snapshot-rates",
+                "pattern": r"^sol-fill-snapshot-test$",
+                "fill_missing_for_providers": ["openai"],
+                "model_info": filled_rates,
+            }
+        ]
+    )
+    try:
+        cost: Final = completion_cost(completion_response=response)
+    finally:
+        set_fallback_generalizations(previous_rules)
+
+    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"),
+    ],
+)
+def test_completion_cost_keeps_served_size_bucket_rates_over_the_requested_model(
+    _local_model_cost_map: None,
+    monkeypatch: pytest.MonkeyPatch,
+    custom_llm_provider: str,
+    requested_model: str,
+    served_model: str,
+    bucket: str,
+) -> None:
+    monkeypatch.setitem(
+        litellm.model_cost,
+        requested_model,
+        {
+            "litellm_provider": custom_llm_provider,
+            "mode": "chat",
+            "input_cost_per_token": 3e-6,
+            "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)
+
+    assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[bucket], usage))
+
+
+def test_completion_cost_bills_a_requested_model_priced_after_an_unpriced_lookup(
+    _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    requested_model: Final = "fireworks_ai/accounts/fireworks/models/sol-late-priced-test"
+    usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
+
+    def display_name_response() -> ModelResponse:
+        response: Final = ModelResponse(model="Sol Late Display", choices=[], usage=usage)
+        response._hidden_params.update({"litellm_model_name": requested_model, "custom_llm_provider": "fireworks_ai"})
+        return response
+
+    completion_cost(completion_response=display_name_response())
+    monkeypatch.setitem(
+        litellm.model_cost,
+        requested_model,
+        {"litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6},
+    )
+    _invalidate_model_cost_lowercase_map()
+
+    cost: Final = completion_cost(completion_response=display_name_response())
+
+    assert cost == pytest.approx(_token_rates_cost(litellm.model_cost[requested_model], usage))
+
+
 def test_completion_cost_orders_candidates_without_extra_huggingface_config_requests(
     _local_model_cost_map: None,
 ) -> None: