From a519d805bb35abad5ba736d0e785aebe813cca26 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 13 Sep 2026 03:29:11 +0000 Subject: [PATCH] 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 | 6 ++++++ litellm/llms/fireworks_ai/chat/transformation.py | 3 +++ litellm/llms/fireworks_ai/common_utils.py | 4 ---- litellm/utils.py | 14 +++++++++----- .../test_fireworks_ai_chat_transformation.py | 12 ++++++++++++ .../test_fireworks_ai_common_utils.py | 16 +++------------- 6 files changed, 33 insertions(+), 22 deletions(-) diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index c5290b41f7b..c9ef7d77ed8 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/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]: """ diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index b4e1856f499..05160d83c12 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/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" diff --git a/litellm/llms/fireworks_ai/common_utils.py b/litellm/llms/fireworks_ai/common_utils.py index b5552583862..21a630a76d7 100644 --- a/litellm/llms/fireworks_ai/common_utils.py +++ b/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 diff --git a/litellm/utils.py b/litellm/utils.py index d088ae7264f..edfdde55699 100644 --- a/litellm/utils.py +++ b/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), ) diff --git a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 8479397efa7..fb0311ef39b 100644 --- a/tests/test_litellm/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/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 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 5aa1429f2e7..16226a3ce74 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,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