diff --git a/litellm/utils.py b/litellm/utils.py index 394ab4b4094..4a92f9b1244 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5556,6 +5556,11 @@ 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 + + provider_prefixed_model_name = f"fireworks_ai/{resolve_fireworks_resource_name(split_model)}" + return PotentialModelNamesAndCustomLLMProvider( split_model=split_model, combined_model_name=combined_model_name, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 835e87aff88..02e2eeec467 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -4448,6 +4448,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``."""