refactor(fireworks): resolve cost map key through a provider config hook

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

litellm/llms/base_llm/base_utils.py

@@ -50,6 +50,12 @@ class BaseLLMModelInfo(ABC):
         """
         return None
 
+    def get_model_cost_key(self, model: str) -> str | None:
+        """
+        Extra `litellm.model_cost` key to try for this provider's spelling of `model`, after the exact keys miss.
+        """
+        return None
+
     @abstractmethod
     def get_models(self, api_key: str | None = None, api_base: str | None = None) -> list[str]:
         """

litellm/llms/fireworks_ai/chat/transformation.py

@@ -602,6 +602,9 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
             return None
         return max(matches, key=lambda match: len(match[0]))[1]
 
+    def get_model_cost_key(self, model: str) -> str:
+        return f"fireworks_ai/{resolve_fireworks_resource_name(model)}"
+
     def get_provider_info(self, model: str) -> ProviderSpecificModelInfo:
         supports_function_calling_value: Final = self._get_model_cost_capability(
             model=model, capability="supports_function_calling"

litellm/llms/fireworks_ai/common_utils.py

@@ -72,10 +72,6 @@ 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

litellm/utils.py

@@ -5556,17 +5556,21 @@ 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 fireworks_cost_map_key
-
-        provider_prefixed_model_name = fireworks_cost_map_key(split_model)
+    provider_model_info: Final = (
+        ProviderConfigManager.get_provider_model_info(model=split_model, provider=LlmProviders(custom_llm_provider))
+        if custom_llm_provider in LlmProvidersSet
+        else None
+    )
+    provider_cost_key: Final = (
+        provider_model_info.get_model_cost_key(split_model) if provider_model_info is not None else None
+    )
 
     return PotentialModelNamesAndCustomLLMProvider(
         split_model=split_model,
         combined_model_name=combined_model_name,
         stripped_model_name=stripped_model_name,
         combined_stripped_model_name=combined_stripped_model_name,
-        provider_prefixed_model_name=provider_prefixed_model_name,
+        provider_prefixed_model_name=provider_cost_key or provider_prefixed_model_name,
         custom_llm_provider=cast(str, custom_llm_provider),
     )
 

tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py

@@ -1813,3 +1813,15 @@ def test_streaming_preserves_selected_model_for_private_accounting():
         completion_response=assembled,
         custom_llm_provider="fireworks_ai",
     ) == pytest.approx(expected_cost)
+
+
+@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_get_model_cost_key_resolves_short_names_to_long_keys(model: str, expected: str) -> None:
+    assert FireworksAIConfig().get_model_cost_key(model) == expected

tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py

@@ -1,6 +1,8 @@
+
 import pytest
 
-from litellm.llms.fireworks_ai.common_utils import fireworks_cost_map_key, resolve_fireworks_resource_name
+
+from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_name
 
 
 @pytest.mark.parametrize(
@@ -41,15 +43,3 @@ from litellm.llms.fireworks_ai.common_utils import fireworks_cost_map_key, resol
 )
 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