tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-19 22:19:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(ui_sso.py): maintain backwards compatibility for older user id va… (#11106)

Browse Source 
* 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
This commit is contained in:
Krish Dholakia
2025-05-23 23:46:29 -07:00
committed by GitHub
parent ba2d4d080f
commit 1caefb0ce0
8 changed files with 197 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/constants.py
+2
View File
@@ -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"]
litellm/proxy/_new_secret_config.yaml
+7 -3
View File
@@ -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
+2 -2
View File
@@ -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
+26 -20
View File
@@ -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
+33 -22
View File
@@ -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
+10
View File
@@ -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(
tests/litellm/proxy/management_endpoints/test_ui_sso.py
+41
View File
@@ -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
+76
View File
@@ -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",