Merge pull request #34257 from heathriel/litellm_fireworks_router_slug_support

fix(fireworks_ai): support router slugs via routers/ prefix

litellm/llms/fireworks_ai/chat/transformation.py

@@ -39,7 +39,11 @@ from ...openai.chat.gpt_transformation import (
     OpenAIChatCompletionStreamingHandler,
     OpenAIGPTConfig,
 )
-from ..common_utils import FireworksAIException, FireworksAIMixin
+from ..common_utils import (
+    FireworksAIException,
+    FireworksAIMixin,
+    resolve_fireworks_resource_name,
+)
 
 
 def _extract_fireworks_hidden_params(payload: dict) -> dict:
@@ -627,12 +631,10 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
         litellm_params: dict,
         headers: dict,
     ) -> dict:
-        if not model.startswith("accounts/") and "#" not in model:
-            if model.endswith("-fast"):
-                model = f"accounts/fireworks/routers/{model}"
-            else:
-                model = f"accounts/fireworks/models/{model}"
-        messages = self._transform_messages_helper(messages=messages, model=model, litellm_params=litellm_params)
+        resolved_model: Final = resolve_fireworks_resource_name(model)
+        messages = self._transform_messages_helper(
+            messages=messages, model=resolved_model, litellm_params=litellm_params
+        )
         if "tools" in optional_params and optional_params["tools"] is not None:
             tools: Final = self._transform_tools(tools=optional_params["tools"])
             optional_params["tools"] = tools
@@ -646,7 +648,7 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig):
                     "include_usage": True,
                 }
         return super().transform_request(
-            model=model,
+            model=resolved_model,
             messages=messages,
             optional_params=optional_params,
             litellm_params=litellm_params,

litellm/llms/fireworks_ai/common_utils.py

@@ -29,6 +29,17 @@ def get_fireworks_session_id(litellm_params: dict) -> str | None:
     return None
 
 
+def resolve_fireworks_resource_name(model: str) -> str:
+    stripped: Final = model.removeprefix("fireworks_ai/")
+    if stripped.startswith("accounts/") or "#" in stripped:
+        return stripped
+    if stripped.startswith(("routers/", "models/")):
+        return f"accounts/fireworks/{stripped}"
+    if stripped.endswith("-fast"):
+        return f"accounts/fireworks/routers/{stripped}"
+    return f"accounts/fireworks/models/{stripped}"
+
+
 class FireworksAIMixin:
     """
     Common Base Config functions across Fireworks AI Endpoints

litellm/llms/fireworks_ai/completion/transformation.py

@@ -13,7 +13,7 @@ from ..chat.transformation import (
     FireworksAIConfig,
     effort_from_chat_template_kwargs,
 )
-from ..common_utils import FireworksAIMixin
+from ..common_utils import FireworksAIMixin, resolve_fireworks_resource_name
 
 _TEXT_COMPLETION_STRIP_PARAMS: Final = (
     frozenset({"truncate_prompt_tokens", "prompt_truncate_len"}) | NIM_VLLM_STRIP_PARAMS
@@ -167,11 +167,8 @@ class FireworksAITextCompletionConfig(FireworksAIMixin, BaseTextCompletionConfig
         translated_params: Final = self.map_extra_body_params(optional_params=optional_params, model=model)
         prompt: Final = _transform_prompt(messages=messages)
 
-        if not model.startswith("accounts/") and "#" not in model:
-            model = f"accounts/fireworks/models/{model}"
-
         data: Final = {
-            "model": model,
+            "model": resolve_fireworks_resource_name(model),
             "prompt": prompt,
             **translated_params,
         }

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

@@ -1295,6 +1295,49 @@ def test_streaming_surfaces_fireworks_response_fields():
     assert surfaced["fireworks_prompt_token_ids"] == [1, 2, 3]
 
 
+def test_transform_request_routes_router_slug():
+    config = FireworksAIConfig()
+
+    data = config.transform_request(
+        model="routers/glm-latest",
+        messages=[{"role": "user", "content": "hi"}],
+        optional_params={},
+        litellm_params={},
+        headers={},
+    )
+
+    assert data["model"] == "accounts/fireworks/routers/glm-latest"
+
+
+def test_transform_request_bare_slug_stays_model():
+    config = FireworksAIConfig()
+
+    data = config.transform_request(
+        model="glm-4p6",
+        messages=[{"role": "user", "content": "hi"}],
+        optional_params={},
+        litellm_params={},
+        headers={},
+    )
+
+    assert data["model"] == "accounts/fireworks/models/glm-4p6"
+
+
+def test_transform_request_direct_route_passthrough():
+    config = FireworksAIConfig()
+    model = "accounts/fireworks/models/qwen2p5-coder-7b#accounts/gitlab/deployments/2fb7764c"
+
+    data = config.transform_request(
+        model=model,
+        messages=[{"role": "user", "content": "hi"}],
+        optional_params={},
+        litellm_params={},
+        headers={},
+    )
+
+    assert data["model"] == model
+
+
 def test_map_extra_body_params_translates_truncate_prompt_tokens():
     config = FireworksAIConfig()
     result = config.map_extra_body_params(

tests/test_litellm/llms/fireworks_ai/completion/test_fireworks_ai_completion_transformation.py

@@ -0,0 +1,34 @@
+import os
+import sys
+
+sys.path.insert(0, os.path.abspath("../../../../.."))
+
+from litellm.llms.fireworks_ai.completion.transformation import (
+    FireworksAITextCompletionConfig,
+)
+
+
+def test_transform_text_completion_request_routes_router_slug():
+    config = FireworksAITextCompletionConfig()
+
+    data = config.transform_text_completion_request(
+        model="routers/glm-latest",
+        messages=[{"role": "user", "content": "hi"}],
+        optional_params={},
+        headers={},
+    )
+
+    assert data["model"] == "accounts/fireworks/routers/glm-latest"
+
+
+def test_transform_text_completion_request_bare_slug_stays_model():
+    config = FireworksAITextCompletionConfig()
+
+    data = config.transform_text_completion_request(
+        model="glm-4p6",
+        messages=[{"role": "user", "content": "hi"}],
+        optional_params={},
+        headers={},
+    )
+
+    assert data["model"] == "accounts/fireworks/models/glm-4p6"

tests/test_litellm/llms/fireworks_ai/test_fireworks_ai_common_utils.py

@@ -0,0 +1,45 @@
+import os
+import sys
+
+import pytest
+
+sys.path.insert(0, os.path.abspath("../../../.."))
+
+from litellm.llms.fireworks_ai.common_utils import resolve_fireworks_resource_name
+
+
+@pytest.mark.parametrize(
+    "model, expected",
+    [
+        ("routers/glm-latest", "accounts/fireworks/routers/glm-latest"),
+        ("routers/firerouter", "accounts/fireworks/routers/firerouter"),
+        ("fireworks_ai/routers/glm-latest", "accounts/fireworks/routers/glm-latest"),
+        ("models/glm-4p6", "accounts/fireworks/models/glm-4p6"),
+        ("fireworks_ai/models/glm-4p6", "accounts/fireworks/models/glm-4p6"),
+        ("glm-4p6", "accounts/fireworks/models/glm-4p6"),
+        ("fireworks_ai/glm-4p6", "accounts/fireworks/models/glm-4p6"),
+        ("kimi-k2p6-fast", "accounts/fireworks/routers/kimi-k2p6-fast"),
+        (
+            "accounts/fireworks/routers/glm-latest",
+            "accounts/fireworks/routers/glm-latest",
+        ),
+        (
+            "accounts/fireworks/models/glm-4p6",
+            "accounts/fireworks/models/glm-4p6",
+        ),
+        (
+            "fireworks_ai/accounts/fireworks/routers/glm-latest",
+            "accounts/fireworks/routers/glm-latest",
+        ),
+        (
+            "accounts/fireworks/models/qwen2p5-coder-7b#accounts/gitlab/deployments/2fb7764c",
+            "accounts/fireworks/models/qwen2p5-coder-7b#accounts/gitlab/deployments/2fb7764c",
+        ),
+        (
+            "glm-4p6#accounts/gitlab/deployments/2fb7764c",
+            "glm-4p6#accounts/gitlab/deployments/2fb7764c",
+        ),
+    ],
+)
+def test_resolve_fireworks_resource_name(model, expected):
+    assert resolve_fireworks_resource_name(model) == expected

type-discipline-budget.json

@@ -30,7 +30,7 @@
     "limit": 16715
   },
   "LIT011": {
-    "limit": 5596
+    "limit": 5593
   },
   "LIT012": {
     "limit": 4519