diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index c5290b41f7b..a446a721307 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -50,6 +50,14 @@ class BaseLLMModelInfo(ABC): """ return None + def get_model_cost_key(self, model: str) -> str | None: + """ + Maps the model name a user sends to the key `litellm.model_cost` stores it under, when the two differ. + `get_model_info` tries this key once the exact `model` and `provider/model` keys miss. The default None means + the provider's user-facing names already match the cost map, so there is nothing extra to try. + """ + 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/utils.py b/litellm/utils.py index c6386471f28..7732cd88cb5 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5563,12 +5563,21 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P split_model = strip_bedrock_routing_prefix(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/test_utils.py b/tests/test_litellm/test_utils.py index fc4207fb26b..d89496a0fd0 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4449,6 +4449,75 @@ def test_fireworks_models_in_backup_cost_map(): ), f"short-form {short_key} does not match long-form {long_key}" +@pytest.fixture +def fireworks_short_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr( + litellm, + "model_cost", + { + "fireworks_ai/accounts/fireworks/models/glm-5p3": { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 2e-6, + "litellm_provider": "fireworks_ai", + "mode": "chat", + "max_tokens": 100, + }, + "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast": { + "input_cost_per_token": 2.1e-6, + "output_cost_per_token": 6.6e-6, + "litellm_provider": "fireworks_ai", + "mode": "chat", + }, + "fireworks_ai/nomic-ai/nomic-embed-text-v1.5": { + "input_cost_per_token": 8e-9, + "output_cost_per_token": 0.0, + "litellm_provider": "fireworks_ai", + "mode": "embedding", + }, + }, + ) + litellm.get_model_info.cache_clear() + yield + litellm.get_model_info.cache_clear() + + +def test_fireworks_short_model_names_resolve_to_long_cost_map_keys(fireworks_short_model_cost_map: None) -> None: + model_info = litellm.get_model_info("fireworks_ai/glm-5p3") + assert model_info["key"] == "fireworks_ai/accounts/fireworks/models/glm-5p3" + assert model_info["input_cost_per_token"] == 1e-6 + assert model_info["max_tokens"] == 100 + + model_info = litellm.get_model_info("glm-5p3", custom_llm_provider="fireworks_ai") + assert model_info["key"] == "fireworks_ai/accounts/fireworks/models/glm-5p3" + + model_info = litellm.get_model_info("fireworks_ai/glm-5p3-fast") + assert model_info["key"] == "fireworks_ai/accounts/fireworks/routers/glm-5p3-fast" + assert model_info["input_cost_per_token"] == 2.1e-6 + + model_info = litellm.get_model_info("fireworks_ai/nomic-ai/nomic-embed-text-v1.5") + assert model_info["key"] == "fireworks_ai/nomic-ai/nomic-embed-text-v1.5" + + with pytest.raises(Exception, match="isn't mapped"): + litellm.get_model_info("fireworks_ai/does-not-exist") + + +def test_fireworks_short_model_names_price_with_completion_cost(fireworks_short_model_cost_map: None) -> None: + from litellm.types.utils import ModelResponse + + response = ModelResponse( + model="fireworks_ai/glm-5p3", + usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + cost = litellm.completion_cost( + completion_response=response, + model="fireworks_ai/glm-5p3", + custom_llm_provider="fireworks_ai", + ) + + assert cost == pytest.approx(10 * 1e-6 + 5 * 2e-6) + + class TestBedrockBaseModelLabelKeepsTools: """Regression for #29618: a Bedrock deployment whose ``base_model`` is a friendly label must not silently drop ``tools``/``tool_choice`` under ``drop_params``."""