fix(fireworks-ai): bill cache-write, reasoning, and audio tokens via the shared cost calculator

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/llms/fireworks_ai/cost_calculator.py

@@ -2,9 +2,9 @@
 For calculating cost of fireworks ai serverless inference models.
 """
 
-import math
+from collections.abc import Mapping
 from datetime import datetime
-from typing import Final
+from typing import Final, cast
 
 from litellm.constants import (
     FIREWORKS_AI_4_B,
@@ -12,12 +12,10 @@ from litellm.constants import (
     FIREWORKS_AI_56_B_MOE,
     FIREWORKS_AI_176_B_MOE,
 )
-from litellm.litellm_core_utils.llm_cost_calc.utils import TokenRates, apply_off_peak_pricing
+from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
 from litellm.types.utils import ModelInfo, Usage
 from litellm.utils import get_model_info
 
-NO_CACHE_READ_RATE: Final = float("nan")
-
 
 # Extract the number of billion parameters from the model name
 # only used for together_computer LLMs
@@ -67,6 +65,30 @@ def _resolve_model_info(model: str) -> ModelInfo:
         return get_model_info(model=base_model, custom_llm_provider="fireworks_ai")
 
 
+def _with_cache_read_fallback(model_info: ModelInfo) -> ModelInfo:
+    """Most fireworks_ai price-map entries publish no cache-read rate, and the provider bills
+    cached reads at the input rate. generic_cost_per_token prices a missing rate at $0, so the
+    fallback is written into a copy of the entry (the shared model-cost dict must not be
+    mutated), including inside off_peak_pricing so cached reads track the off-peak input rate
+    the way the previous calculator did."""
+    if model_info.get("cache_read_input_token_cost") is not None:
+        return model_info
+    input_rate: Final = model_info.get("input_cost_per_token")
+    if input_rate is None:
+        return model_info
+    effective: Final[dict[str, object]] = dict(model_info)
+    effective["cache_read_input_token_cost"] = input_rate
+    off_peak: Final = effective.get("off_peak_pricing")
+    if isinstance(off_peak, Mapping):
+        off_peak_map: Final[Mapping[str, object]] = cast(Mapping[str, object], off_peak)
+        if "cache_read_input_token_cost" not in off_peak_map:
+            effective["off_peak_pricing"] = {
+                **off_peak_map,
+                "cache_read_input_token_cost": off_peak_map.get("input_cost_per_token", input_rate),
+            }
+    return cast(ModelInfo, effective)
+
+
 def cost_per_token(model: str, usage: Usage, current_time: datetime | None = None) -> tuple[float, float]:
     """
     Calculates the cost per token for a given model, prompt tokens, and completion tokens,
@@ -80,29 +102,11 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non
     Returns:
         Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd
     """
-    model_info: Final = _resolve_model_info(model)
-    standard_cache_read_rate: Final = model_info.get("cache_read_input_token_cost")
-    rates: Final = apply_off_peak_pricing(
-        model_info,
-        current_time,
-        TokenRates(
-            input_rate=model_info["input_cost_per_token"] or 0.0,
-            output_rate=model_info["output_cost_per_token"] or 0.0,
-            cache_read_rate=standard_cache_read_rate if standard_cache_read_rate is not None else NO_CACHE_READ_RATE,
-            cache_creation_rate=0.0,
-            reasoning_rate=None,
-        ),
+    model_info: Final = _with_cache_read_fallback(_resolve_model_info(model))
+    return generic_cost_per_token(
+        model=model,
+        usage=usage,
+        custom_llm_provider="fireworks_ai",
+        model_info=model_info,
+        current_time=current_time,
     )
-    cache_read_rate: Final[float] = rates.input_rate if math.isnan(rates.cache_read_rate) else rates.cache_read_rate
-
-    prompt_tokens_details: Final = usage.prompt_tokens_details
-    cached_tokens: Final[int] = (
-        prompt_tokens_details.cached_tokens
-        if prompt_tokens_details is not None and prompt_tokens_details.cached_tokens is not None
-        else 0
-    )
-    non_cached_prompt_tokens: Final[int] = max(usage.prompt_tokens - cached_tokens, 0)
-    prompt_cost: Final[float] = non_cached_prompt_tokens * rates.input_rate + cached_tokens * cache_read_rate
-    completion_cost: Final[float] = usage.completion_tokens * rates.output_rate
-
-    return prompt_cost, completion_cost

tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py

@@ -1,13 +1,16 @@
-
 import math
 from datetime import datetime, timezone
 
 import pytest
 
-
 import litellm
 from litellm.llms.fireworks_ai.cost_calculator import cost_per_token
-from litellm.types.utils import OffPeakPricing, PromptTokensDetailsWrapper, Usage
+from litellm.types.utils import (
+    CompletionTokensDetailsWrapper,
+    OffPeakPricing,
+    PromptTokensDetailsWrapper,
+    Usage,
+)
 
 MODEL = "accounts/fireworks/models/glm-5p2"
 INPUT_COST = 1.4e-06
@@ -47,12 +50,8 @@ def test_warm_call_cheaper_than_cold_call():
     prompt_tokens = 7036
     completion_tokens = 8
 
-    cold_prompt_cost, _ = cost_per_token(
-        model=MODEL, usage=_usage(prompt_tokens, 16, completion_tokens)
-    )
-    warm_prompt_cost, _ = cost_per_token(
-        model=MODEL, usage=_usage(prompt_tokens, 7020, completion_tokens)
-    )
+    cold_prompt_cost, _ = cost_per_token(model=MODEL, usage=_usage(prompt_tokens, 16, completion_tokens))
+    warm_prompt_cost, _ = cost_per_token(model=MODEL, usage=_usage(prompt_tokens, 7020, completion_tokens))
 
     assert warm_prompt_cost < cold_prompt_cost
 
@@ -61,9 +60,7 @@ def test_no_cached_tokens_matches_full_input_rate():
     prompt_tokens = 100
     completion_tokens = 10
 
-    prompt_cost, completion_cost = cost_per_token(
-        model=MODEL, usage=_usage(prompt_tokens, 0, completion_tokens)
-    )
+    prompt_cost, completion_cost = cost_per_token(model=MODEL, usage=_usage(prompt_tokens, 0, completion_tokens))
 
     assert prompt_cost == pytest.approx(prompt_tokens * INPUT_COST)
     assert completion_cost == pytest.approx(completion_tokens * OUTPUT_COST)
@@ -78,7 +75,9 @@ STANDARD_OUTPUT_COST = 6e-07
 STANDARD_CACHE_READ_COST = 1.5e-08
 
 
-def _register_off_peak_model(off_peak_pricing: OffPeakPricing, cache_read_cost: float | None = STANDARD_CACHE_READ_COST) -> None:
+def _register_off_peak_model(
+    off_peak_pricing: OffPeakPricing, cache_read_cost: float | None = STANDARD_CACHE_READ_COST
+) -> None:
     litellm.model_cost[f"fireworks_ai/{OFF_PEAK_MODEL}"] = {
         "litellm_provider": "fireworks_ai",
         "mode": "chat",
@@ -151,10 +150,68 @@ def test_off_peak_window_bills_cached_tokens_at_the_off_peak_input_rate_without_
 def test_off_peak_defaults_to_the_current_time():
     """The proxy's cost dispatch passes no clock, so an all-day window has to apply on the
     default current time."""
-    _register_off_peak_model({"hours_utc": "00:00-00:00", "input_cost_per_token": 1e-08, "output_cost_per_token": 2e-08})
+    _register_off_peak_model(
+        {"hours_utc": "00:00-00:00", "input_cost_per_token": 1e-08, "output_cost_per_token": 2e-08}
+    )
     usage = _usage(prompt_tokens=1000, cached_tokens=0, completion_tokens=200)
 
     prompt_cost, completion_cost = cost_per_token(model=OFF_PEAK_MODEL, usage=usage)
 
     assert math.isclose(prompt_cost, 1000 * 1e-08, rel_tol=1e-10)
     assert math.isclose(completion_cost, 200 * 2e-08, rel_tol=1e-10)
+
+
+COMPONENT_MODEL = "accounts/fireworks/models/cost-components-test"
+COMPONENT_INPUT_COST = 1e-06
+COMPONENT_OUTPUT_COST = 2e-06
+COMPONENT_CACHE_READ_COST = 1e-07
+COMPONENT_CACHE_CREATION_COST = 3e-06
+COMPONENT_REASONING_COST = 4e-06
+COMPONENT_AUDIO_IN_COST = 5e-06
+COMPONENT_AUDIO_OUT_COST = 6e-06
+
+
+def test_cache_write_reasoning_and_audio_tokens_are_billed_at_their_component_rates():
+    """Regression (LIT-7837): the hand-rolled fireworks_ai calculator billed every
+    cache-creation, reasoning and audio token at $0. The shared calculator treats the
+    prompt detail counts as subsets of prompt_tokens and the completion detail counts as
+    subsets of completion_tokens, billing each remainder at the text rate."""
+    litellm.model_cost[f"fireworks_ai/{COMPONENT_MODEL}"] = {
+        "litellm_provider": "fireworks_ai",
+        "mode": "chat",
+        "input_cost_per_token": COMPONENT_INPUT_COST,
+        "output_cost_per_token": COMPONENT_OUTPUT_COST,
+        "cache_read_input_token_cost": COMPONENT_CACHE_READ_COST,
+        "cache_creation_input_token_cost": COMPONENT_CACHE_CREATION_COST,
+        "output_cost_per_reasoning_token": COMPONENT_REASONING_COST,
+        "input_cost_per_audio_token": COMPONENT_AUDIO_IN_COST,
+        "output_cost_per_audio_token": COMPONENT_AUDIO_OUT_COST,
+    }
+    usage = Usage(
+        prompt_tokens=1000,
+        completion_tokens=500,
+        total_tokens=1500,
+        prompt_tokens_details=PromptTokensDetailsWrapper(
+            cached_tokens=300,
+            cache_creation_tokens=200,
+            audio_tokens=100,
+        ),
+        completion_tokens_details=CompletionTokensDetailsWrapper(
+            reasoning_tokens=200,
+            audio_tokens=50,
+        ),
+    )
+
+    prompt_cost, completion_cost = cost_per_token(model=COMPONENT_MODEL, usage=usage)
+
+    expected_prompt_cost = (
+        400 * COMPONENT_INPUT_COST
+        + 300 * COMPONENT_CACHE_READ_COST
+        + 200 * COMPONENT_CACHE_CREATION_COST
+        + 100 * COMPONENT_AUDIO_IN_COST
+    )
+    expected_completion_cost = (
+        250 * COMPONENT_OUTPUT_COST + 200 * COMPONENT_REASONING_COST + 50 * COMPONENT_AUDIO_OUT_COST
+    )
+    assert prompt_cost == pytest.approx(expected_prompt_cost)
+    assert completion_cost == pytest.approx(expected_completion_cost)