fix(ui_sso.py): maintain backwards compatibility for older user id formats + fix existing user email w/ trailing whitespace check + ensure default_internal_user_settings runs on all user new calls

litellm/constants.py

@@ -672,3 +672,5 @@
 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"]

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"
+# 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          

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

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:

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(

litellm/proxy/proxy_server.py

@@ -29,6 +29,7 @@
     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 @@ def _update_config_fields(
         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 @@ def _update_config_fields(
             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(

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"

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",