diff --git a/litellm/llms/fireworks_ai/cost_calculator.py b/litellm/llms/fireworks_ai/cost_calculator.py index d42b4ddbbf7..c4276166b30 100644 --- a/litellm/llms/fireworks_ai/cost_calculator.py +++ b/litellm/llms/fireworks_ai/cost_calculator.py @@ -38,11 +38,10 @@ def get_base_model_for_pricing(model_name: str) -> str: elif total_billion <= FIREWORKS_AI_176_B_MOE: return "fireworks-ai-56b-to-176b" - # Check for standard models in the form b - re_params_match: Final = re.search(r"(\d+)b", model_name) + # Check for standard models in the form b, where Fireworks writes 1.5b as 1p5b + re_params_match: Final = re.search(r"(\d+(?:p\d+)?)b", model_name) if re_params_match is not None: - params_match: Final = str(re_params_match.group(1)) - params_billion: Final = float(params_match) + params_billion: Final = float(re_params_match.group(1).replace("p", ".")) # Determine the category based on the number of parameters if params_billion <= FIREWORKS_AI_4_B: 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..b50a1d7c767 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 @@ -365,6 +365,17 @@ def test_an_unknown_fireworks_model_still_falls_back_to_the_parameter_size_bucke assert prompt_cost > 0 +@pytest.mark.parametrize("size", ["0p6b", "1p5b", "3p8b"]) +def test_an_unknown_fireworks_model_with_a_fractional_size_falls_back_to_its_own_size_bucket(size: str): + usage: Final = _usage(prompt_tokens=100, cached_tokens=0, completion_tokens=10) + + fractional_cost: Final = cost_per_token(model=f"accounts/acme/models/not-in-the-map-{size}", usage=usage) + whole_cost: Final = cost_per_token(model="accounts/acme/models/not-in-the-map-1b", usage=usage) + + assert fractional_cost == whole_cost + assert fractional_cost[0] > 0 + + _TIERED_INPUT_PATTERN: Final = re.compile(r"^input_cost_per_token_above_(\d+)k_tokens$")