phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

feat: set max_internal_budget for user w/ sso

Browse source
This commit is contained in:
Krrish Dholakia 2024-08-08 12:44:07 -07:00
parent 5d664c0441
commit ac601a69fc
4 changed files with 40 additions and 28 deletions
Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/proxy_server.py
View file

@ -283,7 +283,6 @@ except Exception as e:
pass
server_root_path = os.getenv("SERVER_ROOT_PATH", "")
print("server root path: ", server_root_path) # noqa
_license_check = LicenseCheck()
premium_user: bool = _license_check.is_premium()
ui_link = f"{server_root_path}/ui/"
@ -8623,8 +8622,12 @@ async def auth_callback(request: Request):
_last_name = getattr(result, "last_name", "") or ""
user_id = _first_name + _last_name
if user_email is not None and (user_id is None or len(user_id) == 0):
user_id = user_email
user_info = None
user_id_models: List = []
max_internal_user_budget = litellm.max_internal_user_budget
# User might not be already created on first generation of key
# But if it is, we want their models preferences
@ -8636,10 +8639,12 @@ async def auth_callback(request: Request):
"spend": 0,
"team_id": "litellm-dashboard",
}
user_defined_values = {
user_defined_values: SSOUserDefinedValues = {
"models": user_id_models,
"user_id": user_id,
"user_email": user_email,
"max_budget": max_internal_user_budget,
"user_role": None,
}
_user_id_from_sso = user_id
try:
@ -8651,10 +8656,13 @@ async def auth_callback(request: Request):
)
if user_info is not None:
user_defined_values = {
"models": getattr(user_info, "models", []),
"models": getattr(user_info, "models", user_id_models),
"user_id": getattr(user_info, "user_id", user_id),
"user_email": getattr(user_info, "user_id", user_email),
"user_role": getattr(user_info, "user_role", None),
"max_budget": getattr(
user_info, "max_budget", max_internal_user_budget
),
}
user_role = getattr(user_info, "user_role", None)
@ -8668,6 +8676,9 @@ async def auth_callback(request: Request):
"user_id": user_id,
"user_email": getattr(user_info, "user_id", user_email),
"user_role": getattr(user_info, "user_role", None),
"max_budget": getattr(
user_info, "max_budget", max_internal_user_budget
),
}
user_role = getattr(user_info, "user_role", None)
@ -8684,7 +8695,12 @@ async def auth_callback(request: Request):
"user_email": litellm.default_user_params.get(
"user_email", user_email
),
"user_role": litellm.default_user_params.get("user_role", None),
"max_budget": litellm.default_user_params.get(
"max_budget", max_internal_user_budget
),
}
except Exception as e:
pass