diff --git a/litellm/llms/fireworks_ai/cost_calculator.py b/litellm/llms/fireworks_ai/cost_calculator.py index ed936f6233a..ffa04f377a5 100644 --- a/litellm/llms/fireworks_ai/cost_calculator.py +++ b/litellm/llms/fireworks_ai/cost_calculator.py @@ -10,15 +10,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 generic_cost_per_token from litellm.types.utils import Usage -from litellm.utils import get_model_info -# Extract the number of billion parameters from the model name -# only used for together_computer LLMs def get_base_model_for_pricing(model_name: str) -> str: """ - Helper function for calculating together ai pricing. + Resolves a Fireworks model name to its pricing category based on parameter count. Returns: - str: model pricing category if mapped else received model name @@ -27,7 +25,6 @@ def get_base_model_for_pricing(model_name: str) -> str: model_name = model_name.lower() - # Check for MoE models in the form xb moe_match = re.search(r"(\d+)x(\d+)b", model_name) if moe_match: total_billion = int(moe_match.group(1)) * int(moe_match.group(2)) @@ -36,13 +33,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 = re.search(r"(\d+)b", model_name) if re_params_match is not None: - params_match = str(re_params_match.group(1)) - params_billion = float(params_match) + params_billion = float(re_params_match.group(1)) - # Determine the category based on the number of parameters if params_billion <= FIREWORKS_AI_4_B: return "fireworks-ai-up-to-4b" elif params_billion <= FIREWORKS_AI_16_B: @@ -50,7 +44,6 @@ def get_base_model_for_pricing(model_name: str) -> str: elif params_billion > FIREWORKS_AI_16_B: return "fireworks-ai-above-16b" - # If no matches, return the original model_name return "fireworks-ai-default" @@ -60,25 +53,13 @@ def cost_per_token(model: str, usage: Usage) -> Tuple[float, float]: Input: - model: str, the model name without provider prefix - - usage: LiteLLM Usage block, containing anthropic caching information + - usage: LiteLLM Usage block, containing prompt caching information Returns: Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd """ - ## check if model mapped, else use default pricing try: - model_info = get_model_info(model=model, custom_llm_provider="fireworks_ai") + return generic_cost_per_token(model=model, usage=usage, custom_llm_provider="fireworks_ai") except Exception: base_model = get_base_model_for_pricing(model_name=model) - - ## GET MODEL INFO - model_info = get_model_info(model=base_model, custom_llm_provider="fireworks_ai") - - ## CALCULATE INPUT COST - - prompt_cost: float = usage["prompt_tokens"] * model_info["input_cost_per_token"] - - ## CALCULATE OUTPUT COST - completion_cost = usage["completion_tokens"] * model_info["output_cost_per_token"] - - return prompt_cost, completion_cost + return generic_cost_per_token(model=base_model, usage=usage, custom_llm_provider="fireworks_ai") diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py new file mode 100644 index 00000000000..8400cf2ba83 --- /dev/null +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py @@ -0,0 +1,117 @@ +""" +Regression tests for litellm/llms/fireworks_ai/cost_calculator.py + +Covers: +- Cache-read tokens billed at cache_read_input_token_cost (issue #31714) +- Fallback to parameter-size-based pricing for unknown models +- Basic non-cached cost calculation still works +""" + +import pytest + +from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token +from litellm.llms.fireworks_ai.cost_calculator import ( + cost_per_token, + get_base_model_for_pricing, +) +from litellm.types.utils import PromptTokensDetails, Usage + + +class TestCacheReadTokenCost: + """Regression: cached prompt tokens must be billed at cache_read_input_token_cost.""" + + def test_cached_tokens_billed_at_reduced_rate(self): + """ + Given a model with cache_read_input_token_cost defined and a usage object + containing cached_tokens in prompt_tokens_details, the prompt cost must + reflect the reduced rate for cached tokens rather than charging all + prompt tokens at the full input rate. + + Reproduces issue #31714. + """ + usage = Usage( + prompt_tokens=1000, + completion_tokens=100, + total_tokens=1100, + prompt_tokens_details=PromptTokensDetails(cached_tokens=800), + ) + + prompt_cost, completion_cost = cost_per_token( + model="accounts/fireworks/models/deepseek-v4-flash", + usage=usage, + ) + + from litellm.utils import get_model_info + + model_info = get_model_info( + model="accounts/fireworks/models/deepseek-v4-flash", + custom_llm_provider="fireworks_ai", + ) + input_rate = model_info["input_cost_per_token"] + cache_rate = model_info["cache_read_input_token_cost"] + + non_cached_tokens = 1000 - 800 + expected_prompt_cost = non_cached_tokens * input_rate + 800 * cache_rate + expected_completion_cost = 100 * model_info["output_cost_per_token"] + + assert prompt_cost == pytest.approx(expected_prompt_cost) + assert completion_cost == pytest.approx(expected_completion_cost) + + naive_full_rate_cost = 1000 * input_rate + assert prompt_cost < naive_full_rate_cost + + def test_no_cached_tokens_charges_full_rate(self): + """Without cached tokens, all prompt tokens are billed at the full input rate.""" + usage = Usage( + prompt_tokens=1000, + completion_tokens=100, + total_tokens=1100, + ) + + prompt_cost, completion_cost = cost_per_token( + model="accounts/fireworks/models/deepseek-v4-flash", + usage=usage, + ) + + from litellm.utils import get_model_info + + model_info = get_model_info( + model="accounts/fireworks/models/deepseek-v4-flash", + custom_llm_provider="fireworks_ai", + ) + + assert prompt_cost == pytest.approx(1000 * model_info["input_cost_per_token"]) + assert completion_cost == pytest.approx(100 * model_info["output_cost_per_token"]) + + +class TestBaseModelFallback: + """Verify that unknown models fall back to parameter-size-based pricing.""" + + @pytest.mark.parametrize( + "model_name,expected_category", + [ + ("custom-3b-chat", "fireworks-ai-up-to-4b"), + ("custom-8b-instruct", "fireworks-ai-4.1b-to-16b"), + ("custom-70b-instruct", "fireworks-ai-above-16b"), + ("custom-8x7b-moe", "fireworks-ai-moe-up-to-56b"), + ("custom-8x22b-moe", "fireworks-ai-56b-to-176b"), + ], + ) + def test_base_model_resolution(self, model_name: str, expected_category: str): + assert get_base_model_for_pricing(model_name) == expected_category + + def test_unknown_model_still_calculates_cost(self): + """An unregistered model with a recognizable param count yields a valid cost.""" + usage = Usage( + prompt_tokens=500, + completion_tokens=50, + total_tokens=550, + ) + + prompt_cost, completion_cost = cost_per_token( + model="my-custom-70b-chat", + usage=usage, + ) + + assert prompt_cost > 0 + assert completion_cost > 0