diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index 21a630a76d7..b5552583862 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/litellm/llms/fireworks_ai/common_utils.py @@ -72,6 +72,10 @@ def resolve_fireworks_resource_name(model: str) -> str: return f"accounts/fireworks/models/{stripped}" +def fireworks_cost_map_key(model: str) -> str: + return f"fireworks_ai/{resolve_fireworks_resource_name(model)}" + + class FireworksAIMixin: """ Common Base Config functions across Fireworks AI Endpoints diff --git a/litellm/utils.py b/litellm/utils.py index 4a92f9b1244..d088ae7264f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5557,9 +5557,9 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P split_model = strip_bedrock_routing_prefix(split_model) if custom_llm_provider == "fireworks_ai": - from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_name + from litellm.llms.fireworks_ai.common_utils import fireworks_cost_map_key - provider_prefixed_model_name = f"fireworks_ai/{resolve_fireworks_resource_name(split_model)}" + provider_prefixed_model_name = fireworks_cost_map_key(split_model) return PotentialModelNamesAndCustomLLMProvider( split_model=split_model, diff --git a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py index 16226a3ce74..5aa1429f2e7 100644 --- a/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py +++ b/tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py @@ -1,8 +1,6 @@ - import pytest - -from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_name +from litellm.llms.fireworks_ai.common_utils import fireworks_cost_map_key, resolve_fireworks_resource_name @pytest.mark.parametrize( @@ -43,3 +41,15 @@ from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_na ) def test_resolve_fireworks_resource_name(model, expected): assert resolve_fireworks_resource_name(model) == expected + + +@pytest.mark.parametrize( + "model, expected", + [ + ("deepseek-r1", "fireworks_ai/accounts/fireworks/models/deepseek-r1"), + ("glm-5p3-fast", "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast"), + ("accounts/fireworks/models/deepseek-r1", "fireworks_ai/accounts/fireworks/models/deepseek-r1"), + ], +) +def test_fireworks_cost_map_key(model: str, expected: str) -> None: + assert fireworks_cost_map_key(model) == expected