From 7225fc066ff5985053b7ea94b809dd8cac0c9c35 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Thu, 20 Nov 2025 16:05:49 -0800 Subject: [PATCH] Fix key model alias (#16896) --- litellm/proxy/common_request_processing.py | 9 +++ litellm/proxy/litellm_pre_call_utils.py | 33 ++++++++++ .../pass_through_endpoints.py | 9 +++ .../proxy/test_litellm_pre_call_utils.py | 63 +++++++++++++++++++ 4 files changed, 114 insertions(+) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 835afbdc238..5a3e0b334b4 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -378,6 +378,15 @@ class ProxyBaseLLMRequestProcessing: ): 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()) ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index ac789671407..0553c36e3c7 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/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]: diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 3eee47f201b..644b5ce9291 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/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" diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 6379ea704ce..fd39b308a7a 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -15,6 +15,7 @@ from litellm.proxy.litellm_pre_call_utils import ( 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