diff --git a/litellm/llms/fireworks_ai/cost_calculator.py b/litellm/llms/fireworks_ai/cost_calculator.py index e36c9616021..2d92af6613d 100644 --- a/litellm/llms/fireworks_ai/cost_calculator.py +++ b/litellm/llms/fireworks_ai/cost_calculator.py @@ -11,9 +11,13 @@ from litellm.constants import ( FIREWORKS_AI_56_B_MOE, FIREWORKS_AI_176_B_MOE, ) -from litellm.litellm_core_utils.llm_cost_calc.utils import cost_map_key_prices_anything, generic_cost_per_token +from litellm.litellm_core_utils.llm_cost_calc.utils import ( + cost_map_entry_prices_anything, + cost_map_key_prices_anything, + generic_cost_per_token, +) from litellm.types.utils import ModelInfo, Usage -from litellm.utils import get_model_info +from litellm.utils import get_model_info, unmapped_model_generalization # Extract the number of billion parameters from the model name @@ -65,7 +69,13 @@ def _resolve_model_info(model: str) -> ModelInfo: model_info: Final = get_model_info(model=model, custom_llm_provider="fireworks_ai") except Exception: return _resolve_routed_model_info(model) - return model_info if cost_map_key_prices_anything(model_info["key"]) else _resolve_routed_model_info(model) + declares_rates: Final = cost_map_key_prices_anything(model_info["key"]) or _rule_declares_rates(model) + return model_info if declares_rates else _resolve_routed_model_info(model) + + +def _rule_declares_rates(model: str) -> bool: + rule_info: Final = unmapped_model_generalization(model, "fireworks_ai") + return rule_info is not None and cost_map_entry_prices_anything(rule_info) def _resolve_routed_model_info(model: str) -> ModelInfo: diff --git a/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py index c3aec979aa6..9952441debe 100644 --- a/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py +++ b/tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py @@ -7,6 +7,10 @@ from typing import Final import pytest import litellm +from litellm.litellm_core_utils.fallback_generalizations import ( + get_fallback_generalization_rules, + set_fallback_generalizations, +) from litellm.litellm_core_utils.llm_cost_calc.utils import ( calculate_prompt_caching_savings, generic_cost_per_token, @@ -365,6 +369,36 @@ def test_an_unknown_fireworks_model_still_falls_back_to_the_parameter_size_bucke assert prompt_cost > 0 +@pytest.fixture +def priced_rule_for_prefixed_fireworks_names() -> Generator[None, None, None]: + previous_rules: Final = list(get_fallback_generalization_rules()) + try: + set_fallback_generalizations( + [ + { + "name": "fireworks-sol-ruled-priced", + "pattern": r"fireworks_ai/accounts/fireworks/models/sol-ruled-[a-z]+", + "model_info": {"mode": "chat", "input_cost_per_token": 3.21e-05, "output_cost_per_token": 7.65e-05}, + } + ] + ) + yield + finally: + set_fallback_generalizations(previous_rules) + + +def test_a_model_priced_only_by_a_fallback_rule_bills_at_the_rule_rates( + priced_rule_for_prefixed_fireworks_names: None, +): + prompt_cost, completion_cost = cost_per_token( + model="accounts/fireworks/models/sol-ruled-max", + usage=_usage(prompt_tokens=100, cached_tokens=0, completion_tokens=10), + ) + + assert prompt_cost == pytest.approx(100 * 3.21e-05) + assert completion_cost == pytest.approx(10 * 7.65e-05) + + _TIERED_INPUT_PATTERN: Final = re.compile(r"^input_cost_per_token_above_(\d+)k_tokens$")