fix(fireworks): resolve short model names to long cost map keys

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

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,

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``."""