fix(fireworks_ai): route firerouter short names and bill pass-through legs at the routed model's rates (#42814)

* fix(fireworks_ai): route firerouter short names and bill pass-through legs at the routed model's rates

fireworks_ai/firerouter and fireworks_ai/firerouter/<slug> resolve to
accounts/fireworks/routers/... instead of a models/ path, and the cost
calculator falls back to the routed model's own catalog entry before the
Fireworks size buckets so a Claude leg is no longer priced at $0

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

* fix(fireworks_ai): bill routed legs under the routed model's own provider

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

* test(fireworks_ai): require the k suffix when parsing tiered input fields

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

---------

Co-authored-by: kerry <kerry@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/llms/fireworks_ai/common_utils.py

@@ -59,6 +59,7 @@ def resolve_fireworks_api_key(api_key: str | None) -> str | None:
 
 
 AZURE_FOUNDRY_FIREWORKS_MODEL_ID_PREFIX: Final = "FW-"
+FIREROUTER: Final = "firerouter"
 
 
 def resolve_fireworks_resource_name(model: str) -> str:
@@ -67,7 +68,7 @@ def resolve_fireworks_resource_name(model: str) -> str:
         return stripped
     if stripped.startswith(("routers/", "models/")):
         return f"accounts/fireworks/{stripped}"
-    if stripped.endswith("-fast"):
+    if stripped.endswith("-fast") or stripped == FIREROUTER or stripped.startswith(f"{FIREROUTER}/"):
         return f"accounts/fireworks/routers/{stripped}"
     return f"accounts/fireworks/models/{stripped}"
 

litellm/llms/fireworks_ai/cost_calculator.py

@@ -59,6 +59,13 @@ def get_base_model_for_pricing(model_name: str) -> str:
 def _resolve_model_info(model: str) -> ModelInfo:
     try:
         return get_model_info(model=model, custom_llm_provider="fireworks_ai")
+    except Exception:
+        return _resolve_routed_model_info(model)
+
+
+def _resolve_routed_model_info(model: str) -> ModelInfo:
+    try:
+        return get_model_info(model=model.removeprefix("fireworks_ai/"))
     except Exception:
         base_model: Final = get_base_model_for_pricing(model_name=model)
         return get_model_info(model=base_model, custom_llm_provider="fireworks_ai")
@@ -81,7 +88,7 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non
     return generic_cost_per_token(
         model=model,
         usage=usage,
-        custom_llm_provider="fireworks_ai",
+        custom_llm_provider=model_info["litellm_provider"],
         model_info=model_info,
         current_time=current_time,
     )

tests/integration/providers/test_fireworks_ai_router_slug_wire.py

@@ -1,16 +1,70 @@
 import json
+import uuid
+from pathlib import Path
 from typing import Final
 
 import pytest
-from integration._support.client import Gateway
+from integration._support.client import Gateway, eventually
+from integration._support.database import read_rows
 from integration._support.wire import Reply, Request, wire_server
 from pydantic import JsonValue, TypeAdapter
 
 _ROUTER_SLUG: Final = "routers/glm-latest"
 _ROUTER_RESOURCE: Final = "accounts/fireworks/routers/glm-latest"
+_FIREROUTER_SLUGS: Final = ("firerouter", "firerouter/kimi-k3/deepseek-v4")
 _API_KEY: Final = "synthetic-fireworks-key"
 _PROMPT: Final = "route me through the router"
+_COST_MAP_PATH: Final = Path(__file__).resolve().parents[3] / "model_prices_and_context_window.json"
 _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue])
+_COST_MAP: Final = TypeAdapter(dict[str, dict[str, object]])
+
+
+def _positive_rate(entry: dict[str, object], field: str) -> bool:
+    value: Final = entry.get(field)
+    return isinstance(value, (int, float)) and value > 0
+
+
+def _pick_routed_model() -> str:
+    catalog: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes())
+    return next(
+        key
+        for key, entry in catalog.items()
+        if "/" not in key
+        and entry.get("litellm_provider") == "anthropic"
+        and _positive_rate(entry, "input_cost_per_token")
+        and _positive_rate(entry, "output_cost_per_token")
+        and f"fireworks_ai/{key}" not in catalog
+    )
+
+
+def _catalog_cost(model: str, field: str) -> float:
+    cost_value: Final = _COST_MAP.validate_json(_COST_MAP_PATH.read_bytes())[model][field]
+    assert isinstance(cost_value, (int, float))
+    return float(cost_value)
+
+
+_ROUTED_MODEL: Final = _pick_routed_model()
+
+
+def _approx(value: float) -> object:
+    return pytest.approx(value, rel=1e-6)  # pyright: ignore[reportUnknownMemberType]  # pytest lacks typed approx stubs
+
+
+def _chat_completion(identity: str, model: str, prompt_tokens: int, completion_tokens: int) -> bytes:
+    return json.dumps(
+        {
+            "id": identity,
+            "object": "chat.completion",
+            "created": 1,
+            "model": model,
+            "choices": [{"index": 0, "message": {"role": "assistant", "content": "routed"}, "finish_reason": "stop"}],
+            "usage": {
+                "prompt_tokens": prompt_tokens,
+                "completion_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens,
+            },
+        }
+    ).encode()
 
 
 def _provider_body(request: Request, target: str) -> dict[str, JsonValue]:
@@ -82,3 +136,64 @@ def test_fireworks_router_slug_text_completion_sends_router_resource_not_models_
         payload: Final = _JSON_OBJECT.validate_json(response.content)
         assert payload["choices"] == [{"index": 0, "text": "routed", "finish_reason": "stop", "logprobs": None}]
         assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/completions")]
+
+
+@pytest.mark.parametrize("slug", _FIREROUTER_SLUGS)
+def test_fireworks_firerouter_short_name_sends_router_resource_not_models_path(gateway: Gateway, slug: str) -> None:
+    resource: Final = f"accounts/fireworks/routers/{slug}"
+
+    def respond(request: Request) -> Reply:
+        body: Final = _provider_body(request, "/chat/completions")
+        assert body["model"] == resource, body
+        return Reply(body=_chat_completion(f"fw-{slug}", resource, 5, 1))
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(model=f"fireworks_ai/{slug}", api_base=wire.url, api_key=_API_KEY)
+        response: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {"model": model, "messages": [{"role": "user", "content": _PROMPT}]},
+        )
+        assert response.status_code == 200, response.text
+        payload: Final = _JSON_OBJECT.validate_json(response.content)
+        assert payload["choices"] == [
+            {"finish_reason": "stop", "index": 0, "message": {"role": "assistant", "content": "routed"}}
+        ]
+        assert [(request.method, request.target) for request in wire.drain()] == [("POST", "/chat/completions")]
+
+
+def test_fireworks_firerouter_claude_leg_is_charged_at_the_routed_models_own_rate(gateway: Gateway) -> None:
+    identity: Final = f"fw-firerouter-claude-{uuid.uuid4().hex}"
+
+    def respond(request: Request) -> Reply:
+        body: Final = _provider_body(request, "/chat/completions")
+        assert body["model"] == "accounts/fireworks/routers/firerouter", body
+        assert request.headers["x-anthropic-api-key"] == "synthetic-anthropic-key"
+        return Reply(body=_chat_completion(identity, _ROUTED_MODEL, 23, 41))
+
+    with wire_server(respond) as wire, gateway.scenario() as scenario:
+        model: Final = scenario.model(
+            model="fireworks_ai/firerouter",
+            api_base=wire.url,
+            api_key=_API_KEY,
+            extra_headers={"x-anthropic-api-key": "synthetic-anthropic-key"},
+        )
+        response: Final = gateway.request(
+            "POST",
+            "/v1/chat/completions",
+            {"model": model, "messages": [{"role": "user", "content": _PROMPT}]},
+        )
+        assert response.status_code == 200, response.text
+        expected_cost: Final = 23 * _catalog_cost(_ROUTED_MODEL, "input_cost_per_token") + 41 * _catalog_cost(
+            _ROUTED_MODEL, "output_cost_per_token"
+        )
+        assert expected_cost > 0
+        assert float(response.headers["x-litellm-response-cost"]) == _approx(expected_cost)
+        rows: Final = eventually(
+            lambda: read_rows('SELECT spend FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)),
+            lambda values: len(values) == 1,
+            seconds=70,
+        )
+        spend: Final = rows[0]["spend"]
+    assert isinstance(spend, (int, float, str))
+    assert float(spend) == _approx(expected_cost)

tests/unit/llms/fireworks_ai/test_fireworks_ai_common_utils.py

@@ -1,7 +1,5 @@
-
 import pytest
 
-
 from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_name
 
 
@@ -16,6 +14,10 @@ from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_na
         ("glm-4p6", "accounts/fireworks/models/glm-4p6"),
         ("fireworks_ai/glm-4p6", "accounts/fireworks/models/glm-4p6"),
         ("kimi-k2p6-fast", "accounts/fireworks/routers/kimi-k2p6-fast"),
+        ("firerouter", "accounts/fireworks/routers/firerouter"),
+        ("fireworks_ai/firerouter", "accounts/fireworks/routers/firerouter"),
+        ("firerouter/kimi-k3/deepseek-v4", "accounts/fireworks/routers/firerouter/kimi-k3/deepseek-v4"),
+        ("firerouter-v2", "accounts/fireworks/models/firerouter-v2"),
         (
             "accounts/fireworks/routers/glm-latest",
             "accounts/fireworks/routers/glm-latest",

tests/unit/llms/fireworks_ai/test_fireworks_ai_cost_calculator.py

@@ -1,4 +1,5 @@
 import math
+import re
 from collections.abc import Generator
 from datetime import datetime, timezone
 from typing import Final
@@ -326,3 +327,71 @@ def test_an_entry_without_an_input_rate_gets_no_cache_read_fallback():
 
     assert prompt_cost == 0
     assert completion_cost == 200 * 2e-06
+
+
+ROUTED_MODEL: Final = next(
+    key
+    for key, info in litellm.model_cost.items()
+    if "/" not in key
+    and info.get("litellm_provider") == "anthropic"
+    and (info.get("input_cost_per_token") or 0) > 0
+    and (info.get("output_cost_per_token") or 0) > 0
+    and f"fireworks_ai/{key}" not in litellm.model_cost
+)
+
+
+@pytest.mark.parametrize("model", [ROUTED_MODEL, f"fireworks_ai/{ROUTED_MODEL}"])
+def test_a_model_routed_to_another_provider_is_billed_at_that_models_own_rates(model: str):
+    own_rates: Final = litellm.get_model_info(model=ROUTED_MODEL, custom_llm_provider="anthropic")
+    usage: Final = _usage(prompt_tokens=23, cached_tokens=0, completion_tokens=41)
+
+    prompt_cost, completion_cost = cost_per_token(model=model, usage=usage)
+
+    assert prompt_cost == pytest.approx(23 * own_rates["input_cost_per_token"])
+    assert completion_cost == pytest.approx(41 * own_rates["output_cost_per_token"])
+    assert prompt_cost > 0 and completion_cost > 0
+
+
+def test_an_unknown_fireworks_model_still_falls_back_to_the_parameter_size_bucket():
+    prompt_cost, completion_cost = cost_per_token(
+        model="accounts/fireworks/models/not-in-the-map-13b",
+        usage=_usage(prompt_tokens=100, cached_tokens=0, completion_tokens=10),
+    )
+    bucket_prompt_cost, bucket_completion_cost = cost_per_token(
+        model="fireworks-ai-4.1b-to-16b", usage=_usage(prompt_tokens=100, cached_tokens=0, completion_tokens=10)
+    )
+
+    assert (prompt_cost, completion_cost) == (bucket_prompt_cost, bucket_completion_cost)
+    assert prompt_cost > 0
+
+
+_TIERED_INPUT_PATTERN: Final = re.compile(r"^input_cost_per_token_above_(\d+)k_tokens$")
+
+
+def _threshold_tokens(field: str) -> int:
+    match: Final = _TIERED_INPUT_PATTERN.match(field)
+    assert match is not None, field
+    return int(match.group(1)) * 1000
+
+
+def test_a_routed_xai_model_keeps_xais_inclusive_token_threshold():
+    candidate: Final = next(
+        (
+            (key, field)
+            for key, info in litellm.model_cost.items()
+            if info.get("litellm_provider") == "xai" and f"fireworks_ai/{key}" not in litellm.model_cost
+            for field in info
+            if _TIERED_INPUT_PATTERN.match(field)
+        ),
+        None,
+    )
+    if candidate is None:
+        pytest.skip("cost map has no xai entry with a tiered input rate")
+    key, field = candidate
+    usage: Final = _usage(prompt_tokens=_threshold_tokens(field), cached_tokens=0, completion_tokens=10)
+
+    routed_prompt_cost, routed_completion_cost = cost_per_token(model=f"fireworks_ai/{key}", usage=usage)
+
+    assert (routed_prompt_cost, routed_completion_cost) == generic_cost_per_token(
+        model=key, usage=usage, custom_llm_provider="xai"
+    )