diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3fa0ef51f..c8348ca5e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1744,3 +1744,11 @@ class ProxyErrorTypes(str, enum.Enum): auth_error = "auth_error" internal_server_error = "internal_server_error" bad_request_error = "bad_request_error" + + +class SSOUserDefinedValues(TypedDict): + models: List[str] + user_id: str + user_email: Optional[str] + user_role: Optional[str] + max_budget: Optional[float] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 299b390b9..c9e9e3420 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/ui/litellm-dashboard/src/components/user_dashboard.tsx b/ui/litellm-dashboard/src/components/user_dashboard.tsx index 61ec05847..3d1d4ea60 100644 --- a/ui/litellm-dashboard/src/components/user_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/user_dashboard.tsx @@ -315,6 +315,7 @@ const UserDashboard: React.FC = ({ = ({ userID, userRole, accessToken, userSpend, selectedTeam }) => { +const ViewUserSpend: React.FC = ({ userID, userRole, accessToken, userSpend, userMaxBudget, selectedTeam }) => { console.log(`userSpend: ${userSpend}`) let [spend, setSpend] = useState(userSpend !== null ? userSpend : 0.0); const [maxBudget, setMaxBudget] = useState(selectedTeam ? selectedTeam.max_budget : null); + useEffect(() => { + console.log(`Updating team max budget for default team - userMaxBudget - ${userMaxBudget}, selectedTeam.team_alias - ${selectedTeam.team_alias}`) + if (selectedTeam) { + if (selectedTeam.team_alias === "Default Team") { + setMaxBudget(userMaxBudget); + } else { + setMaxBudget(selectedTeam.max_budget); + } + } + }, [selectedTeam, userMaxBudget]); console.log(`maxBudget: ${maxBudget}, selectedTeam.max_budget: ${selectedTeam.max_budget}, selectedTeam: ${JSON.stringify(selectedTeam)}`) const [userModels, setUserModels] = useState([]); useEffect(() => { @@ -53,25 +64,6 @@ const ViewUserSpend: React.FC = ({ userID, userRole, accessT if (!accessToken || !userID || !userRole) { return; } - // if (userRole === "Admin" && userSpend == null) { - // try { - // const globalSpend = await getTotalSpendCall(accessToken); - // if (globalSpend) { - // if (globalSpend.spend) { - // setSpend(globalSpend.spend); - // } else { - // setSpend(0.0); - // } - // if (globalSpend.max_budget) { - // setMaxBudget(globalSpend.max_budget); - // } else { - // setMaxBudget(null); - // } - // } - // } catch (error) { - // console.error("Error fetching global spend data:", error); - // } - // } }; const fetchUserModels = async () => { try { @@ -103,11 +95,6 @@ const ViewUserSpend: React.FC = ({ userID, userRole, accessT setSpend(userSpend) } }, [userSpend]) - useEffect(() => { - if (selectedTeam && selectedTeam.max_budget !== maxBudget) { - setMaxBudget(selectedTeam.max_budget); - } - }, [selectedTeam, maxBudget]); // logic to decide what models to display let modelsToDisplay = [];