Fix key model alias (#16896)

Author: yuneng-jiang

litellm/proxy/common_request_processing.py

@@ -378,6 +378,15 @@ async def common_processing_pre_call_logic(
         ):
             self.data["model"] = litellm.model_alias_map[self.data["model"]]
 
+        # Check key-specific aliases
+        if (
+            isinstance(self.data["model"], str)
+            and user_api_key_dict.aliases
+            and isinstance(user_api_key_dict.aliases, dict)
+            and self.data["model"] in user_api_key_dict.aliases
+        ):
+            self.data["model"] = user_api_key_dict.aliases[self.data["model"]]
+
         self.data["litellm_call_id"] = request.headers.get(
             "x-litellm-call-id", str(uuid.uuid4())
         )

litellm/proxy/litellm_pre_call_utils.py

@@ -1053,6 +1053,12 @@ async def add_litellm_data_to_request(  # noqa: PLR0915
         user_api_key_dict=user_api_key_dict,
     )
 
+    # Key Model Aliases
+    _update_model_if_key_alias_exists(
+        data=data,
+        user_api_key_dict=user_api_key_dict,
+    )
+
     verbose_proxy_logger.debug(
         "[PROXY] returned data from litellm_pre_call_utils: %s", data
     )
@@ -1108,6 +1114,33 @@ def _update_model_if_team_alias_exists(
     return
 
 
+def _update_model_if_key_alias_exists(
+    data: dict,
+    user_api_key_dict: UserAPIKeyAuth,
+) -> None:
+    """
+    Update the model if the key alias exists
+
+    If an alias map has been set on a key, then we want to make the request with the model the key alias is pointing to
+
+    eg.
+        - user calls `modelAlias`
+        - key.aliases = {
+            "modelAlias": "xai/grok-4-fast-non-reasoning"
+        }
+        - requested_model = "xai/grok-4-fast-non-reasoning"
+    """
+    _model = data.get("model")
+    if (
+        _model
+        and user_api_key_dict.aliases
+        and isinstance(user_api_key_dict.aliases, dict)
+        and _model in user_api_key_dict.aliases
+    ):
+        data["model"] = user_api_key_dict.aliases[_model]
+    return
+
+
 def _get_enforced_params(
     general_settings: Optional[dict], user_api_key_dict: UserAPIKeyAuth
 ) -> Optional[list]:

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -197,6 +197,15 @@ async def chat_completion_pass_through_endpoint(  # noqa: PLR0915
         if data["model"] in litellm.model_alias_map:
             data["model"] = litellm.model_alias_map[data["model"]]
 
+        # Check key-specific aliases
+        if (
+            isinstance(data["model"], str)
+            and user_api_key_dict.aliases
+            and isinstance(user_api_key_dict.aliases, dict)
+            and data["model"] in user_api_key_dict.aliases
+        ):
+            data["model"] = user_api_key_dict.aliases[data["model"]]
+
         ### CALL HOOKS ### - modify incoming data before calling the model
         data = await proxy_logging_obj.pre_call_hook(  # type: ignore
             user_api_key_dict=user_api_key_dict, data=data, call_type="text_completion"

tests/test_litellm/proxy/test_litellm_pre_call_utils.py

@@ -15,6 +15,7 @@
     LiteLLMProxyRequestSetup,
     _get_dynamic_logging_metadata,
     _get_enforced_params,
+    _update_model_if_key_alias_exists,
     add_litellm_data_to_request,
     check_if_token_is_service_account,
 )
@@ -1234,3 +1235,65 @@ async def test_request_guardrails_do_not_override_key_guardrails():
     assert "guardrails" not in updated_data_empty
     assert "key-guardrail-1" in requested_guardrails
     assert len(requested_guardrails) == 1
+
+
+def test_update_model_if_key_alias_exists():
+    """
+    Test that _update_model_if_key_alias_exists properly updates the model when a key alias exists.
+    """
+    # Test case 1: Key alias exists and matches model
+    data = {"model": "modelAlias", "messages": [{"role": "user", "content": "Hello"}]}
+    user_api_key_dict = UserAPIKeyAuth(
+        api_key="test-key",
+        aliases={"modelAlias": "xai/grok-4-fast-non-reasoning"},
+    )
+    _update_model_if_key_alias_exists(data=data, user_api_key_dict=user_api_key_dict)
+    assert data["model"] == "xai/grok-4-fast-non-reasoning"
+
+    # Test case 2: Key alias doesn't exist
+    data = {"model": "unknown-model", "messages": [{"role": "user", "content": "Hello"}]}
+    user_api_key_dict = UserAPIKeyAuth(
+        api_key="test-key",
+        aliases={"modelAlias": "xai/grok-4-fast-non-reasoning"},
+    )
+    original_model = data["model"]
+    _update_model_if_key_alias_exists(data=data, user_api_key_dict=user_api_key_dict)
+    assert data["model"] == original_model  # Should remain unchanged
+
+    # Test case 3: Model is None
+    data = {"model": None, "messages": [{"role": "user", "content": "Hello"}]}
+    user_api_key_dict = UserAPIKeyAuth(
+        api_key="test-key",
+        aliases={"modelAlias": "xai/grok-4-fast-non-reasoning"},
+    )
+    _update_model_if_key_alias_exists(data=data, user_api_key_dict=user_api_key_dict)
+    assert data["model"] is None  # Should remain None
+
+    # Test case 4: Model key doesn't exist in data
+    data = {"messages": [{"role": "user", "content": "Hello"}]}
+    user_api_key_dict = UserAPIKeyAuth(
+        api_key="test-key",
+        aliases={"modelAlias": "xai/grok-4-fast-non-reasoning"},
+    )
+    _update_model_if_key_alias_exists(data=data, user_api_key_dict=user_api_key_dict)
+    assert "model" not in data  # Should not add model if it doesn't exist
+
+    # Test case 5: Multiple aliases, matching one
+    data = {"model": "alias1", "messages": [{"role": "user", "content": "Hello"}]}
+    user_api_key_dict = UserAPIKeyAuth(
+        api_key="test-key",
+        aliases={
+            "alias1": "model1",
+            "alias2": "model2",
+            "alias3": "model3",
+        },
+    )
+    _update_model_if_key_alias_exists(data=data, user_api_key_dict=user_api_key_dict)
+    assert data["model"] == "model1"
+
+    # Test case 6: Empty aliases dict
+    data = {"model": "modelAlias", "messages": [{"role": "user", "content": "Hello"}]}
+    user_api_key_dict = UserAPIKeyAuth(api_key="test-key", aliases={})
+    original_model = data["model"]
+    _update_model_if_key_alias_exists(data=data, user_api_key_dict=user_api_key_dict)
+    assert data["model"] == original_model  # Should remain unchanged