From 1caefb0ce037726168031c5152e199471b75c184 Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Fri, 23 May 2025 23:46:29 -0700 Subject: [PATCH] =?UTF-8?q?fix(ui=5Fsso.py):=20maintain=20backwards=20comp?= =?UTF-8?q?atibility=20for=20older=20user=20id=20va=E2=80=A6=20(#11106)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(ui_sso.py): maintain backwards compatibility for older user id variations Fixes issue in later SSO checks which only checked id from result * fix(internal_user_endpoints.py): handle trailing whitespace in new user email * fix(internal_user_endpoints.py): apply default_internal_user_settings on all new user calls (even when role not set) allows role undefined users to be assigned the correct role on sign up * feat(proxy_server.py): load default user settings from db - update litellm correctly updates the litellm module with default internal user settings ensures updated settings actually apply * test: add unit test * fix(internal_user_endpoints.py): fix internal user default param role * fix(ui_sso.py): fix linting error --- litellm/constants.py | 2 + litellm/proxy/_new_secret_config.yaml | 10 ++- litellm/proxy/auth/user_api_key_auth.py | 4 +- .../internal_user_endpoints.py | 46 ++++++----- litellm/proxy/management_endpoints/ui_sso.py | 55 ++++++++------ litellm/proxy/proxy_server.py | 10 +++ .../proxy/management_endpoints/test_ui_sso.py | 41 ++++++++++ tests/proxy_unit_tests/test_proxy_utils.py | 76 +++++++++++++++++++ 8 files changed, 197 insertions(+), 47 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index bb5d56978b..43a8eba3c1 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -672,3 +672,5 @@ LENGTH_OF_LITELLM_GENERATED_KEY = int(os.getenv("LENGTH_OF_LITELLM_GENERATED_KEY SECRET_MANAGER_REFRESH_INTERVAL = int( os.getenv("SECRET_MANAGER_REFRESH_INTERVAL", 86400) ) +LITELLM_SETTINGS_SAFE_DB_OVERRIDES = ["default_internal_user_params"] +SPECIAL_LITELLM_AUTH_TOKEN = ["ui-token"] diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 90a44b1f27..5fa3e618c2 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -84,6 +84,10 @@ litellm_settings: general_settings: store_model_in_db: true store_prompts_in_spend_logs: true - custom_auth: custom_auth_auto.user_api_key_auth - custom_auth_settings: - mode: "auto" \ No newline at end of file +# litellm_settings: +# # max_internal_user_budget: 10 # max budget for internal users +# # internal_user_budget_duration: "1mo" # reset every month + +# default_internal_user_params: # Default Params used when a new user signs in Via SSO +# user_role: "proxy_admin" # one of "internal_user", "internal_user_viewer", "proxy_admin", "proxy_admin_viewer". New SSO users not in litellm will be created as this user +# max_budget: 1000 \ No newline at end of file diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index e276681ced..31f8437d29 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -725,8 +725,8 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 ) ## Check DB - if isinstance( - api_key, str + if ( + isinstance(api_key, str) and valid_token is None ): # if generated token, make sure it starts with sk-. assert api_key.startswith( "sk-" diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index aeddf8c1e3..cb58f0559e 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -59,28 +59,34 @@ def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> d "table_name" ] = "user" # only create a user, don't create key if 'auto_create_key' set to False - is_internal_user = False - if data.user_role and data.user_role.is_internal_user_role: - is_internal_user = True - if litellm.default_internal_user_params: - for key, value in litellm.default_internal_user_params.items(): - if key == "available_teams": - continue - elif key not in data_json or data_json[key] is None: - data_json[key] = value - elif ( - key == "models" - and isinstance(data_json[key], list) - and len(data_json[key]) == 0 - ): - data_json[key] = value + if litellm.default_internal_user_params: + for key, value in litellm.default_internal_user_params.items(): + if key == "available_teams": + continue + elif key not in data_json or data_json[key] is None: + data_json[key] = value + elif ( + key == "models" + and isinstance(data_json[key], list) + and len(data_json[key]) == 0 + ): + data_json[key] = value - if "max_budget" in data_json and data_json["max_budget"] is None: - if is_internal_user and litellm.max_internal_user_budget is not None: + ## INTERNAL USER ROLE ONLY DEFAULT PARAMS ## + if ( + data.user_role is not None + and data.user_role == LitellmUserRoles.INTERNAL_USER.value + ): + if ( + litellm.max_internal_user_budget is not None + and data_json.get("max_budget") is None + ): data_json["max_budget"] = litellm.max_internal_user_budget - if "budget_duration" in data_json and data_json["budget_duration"] is None: - if is_internal_user and litellm.internal_user_budget_duration is not None: + if ( + litellm.internal_user_budget_duration is not None + and data_json.get("budget_duration") is None + ): data_json["budget_duration"] = litellm.internal_user_budget_duration return data_json @@ -105,7 +111,7 @@ async def _check_duplicate_user_email( raise Exception("Database not connected") existing_user = await prisma_client.db.litellm_usertable.find_first( - where={"user_email": user_email} + where={"user_email": user_email.strip()} ) if existing_user is not None: diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index a85c43a2a6..40127bae79 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -353,31 +353,41 @@ async def get_user_info_from_db( proxy_logging_obj: ProxyLogging, user_email: Optional[str], user_defined_values: Optional[SSOUserDefinedValues], + alternate_user_id: Optional[str] = None, ) -> Optional[Union[LiteLLM_UserTable, NewUserResponse]]: try: - user_info: Optional[ - Union[LiteLLM_UserTable, NewUserResponse] - ] = await get_existing_user_info_from_db( - user_id=cast( - Optional[str], - ( - getattr(result, "id", None) - if not isinstance(result, dict) - else result.get("id", None) - ), - ), - user_email=cast( - Optional[str], - ( - getattr(result, "email", None) - if not isinstance(result, dict) - else result.get("email", None) - ), - ), - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - proxy_logging_obj=proxy_logging_obj, + + potential_user_ids = [] + if alternate_user_id is not None: + potential_user_ids.append(alternate_user_id) + if not isinstance(result, dict): + _id = getattr(result, "id", None) + if _id is not None and isinstance(_id, str): + potential_user_ids.append(_id) + else: + _id = result.get("id", None) + if _id is not None and isinstance(_id, str): + potential_user_ids.append(_id) + + user_email = ( + getattr(result, "email", None) + if not isinstance(result, dict) + else result.get("email", None) ) + + user_info: Optional[Union[LiteLLM_UserTable, NewUserResponse]] = None + + for user_id in potential_user_ids: + user_info = await get_existing_user_info_from_db( + user_id=user_id, + user_email=user_email, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + if user_info is not None: + break + verbose_proxy_logger.debug( f"user_info: {user_info}; litellm.default_internal_user_params: {litellm.default_internal_user_params}" ) @@ -566,6 +576,7 @@ async def auth_callback(request: Request): # noqa: PLR0915 proxy_logging_obj=proxy_logging_obj, user_email=user_email, user_defined_values=user_defined_values, + alternate_user_id=user_id, ) user_defined_values = apply_user_info_values_to_sso_user_defined_values( diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b88e221880..a3add15669 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -29,6 +29,7 @@ from litellm.constants import ( DEFAULT_MAX_RECURSE_DEPTH, DEFAULT_SLACK_ALERTING_THRESHOLD, LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS, + LITELLM_SETTINGS_SAFE_DB_OVERRIDES, ) from litellm.types.utils import ( ModelResponse, @@ -2581,13 +2582,21 @@ class ProxyConfig: Returns: dict: Updated configuration dictionary """ + if param_name == "environment_variables": self._decrypt_and_set_db_env_variables(db_param_value) return current_config + elif param_name == "litellm_settings" and isinstance(db_param_value, dict): + for key, value in db_param_value.items(): + if ( + key in LITELLM_SETTINGS_SAFE_DB_OVERRIDES + ): # params that are safe to override with db values + setattr(litellm, key, value) # If param doesn't exist in config, add it if param_name not in current_config: current_config[param_name] = db_param_value + return current_config # For dictionary values, update only non-empty values @@ -2599,6 +2608,7 @@ class ProxyConfig: current_config[param_name].update(non_empty_values) else: current_config[param_name] = db_param_value + return current_config async def _update_config_from_db( diff --git a/tests/litellm/proxy/management_endpoints/test_ui_sso.py b/tests/litellm/proxy/management_endpoints/test_ui_sso.py index 1877f8dfc9..75bf17cfd6 100644 --- a/tests/litellm/proxy/management_endpoints/test_ui_sso.py +++ b/tests/litellm/proxy/management_endpoints/test_ui_sso.py @@ -588,3 +588,44 @@ async def test_get_user_info_from_db(): user_info = await get_user_info_from_db(**args) mock_get_user_object.assert_called_once() mock_get_user_object.call_args.kwargs["user_id"] = "krrishd" + + +async def test_get_user_info_from_db_alternate_user_id(): + from litellm.proxy.management_endpoints.ui_sso import get_user_info_from_db + + prisma_client = MagicMock() + user_api_key_cache = MagicMock() + proxy_logging_obj = MagicMock() + user_email = "krrishdholakia@gmail.com" + user_defined_values = { + "models": [], + "user_id": "krrishd", + "user_email": "krrishdholakia@gmail.com", + "max_budget": None, + "user_role": None, + "budget_duration": None, + } + args = { + "result": CustomOpenID( + id="krrishd", + email="krrishdholakia@gmail.com", + first_name=None, + last_name=None, + display_name="a3f1c107-04dc-4c93-ae60-7f32eb4b05ce", + picture=None, + provider=None, + team_ids=[], + ), + "prisma_client": prisma_client, + "user_api_key_cache": user_api_key_cache, + "proxy_logging_obj": proxy_logging_obj, + "user_email": user_email, + "user_defined_values": user_defined_values, + "alternate_user_id": "krrishd-email1234", + } + with patch.object( + litellm.proxy.management_endpoints.ui_sso, "get_user_object" + ) as mock_get_user_object: + user_info = await get_user_info_from_db(**args) + mock_get_user_object.assert_called_once() + mock_get_user_object.call_args.kwargs["user_id"] = "krrishd-email1234" diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 4dee0edd94..e02d9acf3c 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -547,6 +547,53 @@ def test_update_internal_user_params(): ) +def test_update_internal_new_user_params_with_no_initial_role_set(): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_internal_new_user_params, + ) + from litellm.proxy._types import NewUserRequest + + litellm.default_internal_user_params = { + "max_budget": 100, + "budget_duration": "30d", + "models": ["gpt-3.5-turbo"], + } + + data = NewUserRequest(user_email="krrish3@berri.ai") + data_json = data.model_dump() + updated_data_json = _update_internal_new_user_params(data_json, data) + assert updated_data_json["models"] == litellm.default_internal_user_params["models"] + assert ( + updated_data_json["max_budget"] + == litellm.default_internal_user_params["max_budget"] + ) + assert ( + updated_data_json["budget_duration"] + == litellm.default_internal_user_params["budget_duration"] + ) + +def test_update_internal_new_user_params_with_user_defined_values(): + from litellm.proxy.management_endpoints.internal_user_endpoints import ( + _update_internal_new_user_params, + ) + from litellm.proxy._types import NewUserRequest + + litellm.default_internal_user_params = { + "max_budget": 100, + "budget_duration": "30d", + "models": ["gpt-3.5-turbo"], + "user_role": "proxy_admin", + } + + data = NewUserRequest(user_email="krrish3@berri.ai", max_budget=1000, budget_duration="1mo") + data_json = data.model_dump() + updated_data_json = _update_internal_new_user_params(data_json, data) + assert updated_data_json["user_email"] == "krrish3@berri.ai" + assert updated_data_json["user_role"] == "proxy_admin" + assert updated_data_json["max_budget"] == 1000 + assert updated_data_json["budget_duration"] == "1mo" + + @pytest.mark.asyncio async def test_proxy_config_update_from_db(): from litellm.proxy.proxy_server import ProxyConfig @@ -956,6 +1003,35 @@ def test_update_config_fields(): assert team_config["langfuse_public_key"] == "my-fake-key" assert team_config["langfuse_secret"] == "my-fake-secret" +def test_update_config_fields_default_internal_user_params(monkeypatch): + from litellm.proxy.proxy_server import ProxyConfig + + proxy_config = ProxyConfig() + + monkeypatch.setattr(litellm, "default_internal_user_params", None) + + + args = { + "current_config": {}, + "param_name": "litellm_settings", + "db_param_value": { + "default_internal_user_params": { + "user_role": "proxy_admin", + "max_budget": 1000, + "budget_duration": "1mo", + }, + }, + } + updated_config = proxy_config._update_config_fields(**args) + + assert litellm.default_internal_user_params == { + "user_role": "proxy_admin", + "max_budget": 1000, + "budget_duration": "1mo", + } + + monkeypatch.setattr(litellm, "default_internal_user_params", None) # reset to default + @pytest.mark.parametrize( "proxy_model_list,provider",