feat: set max_internal_budget for user w/ sso

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]

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
 

ui/litellm-dashboard/src/components/user_dashboard.tsx

@@ -315,6 +315,7 @@ const UserDashboard: React.FC<UserDashboardProps> = ({
           <ViewUserSpend
             userID={userID}
             userRole={userRole}
+            userMaxBudget={userSpendData?.max_budget || null}
             accessToken={accessToken}
             userSpend={teamSpend}
             selectedTeam={selectedTeam ? selectedTeam : null}

ui/litellm-dashboard/src/components/view_user_spend.tsx

@@ -40,12 +40,23 @@ interface ViewUserSpendProps {
     userRole: string | null;
     accessToken: string | null;
     userSpend: number | null;  
+    userMaxBudget: number | null;  
     selectedTeam: any | null;
 }
-const ViewUserSpend: React.FC<ViewUserSpendProps> = ({ userID, userRole, accessToken, userSpend, selectedTeam }) => {
+const ViewUserSpend: React.FC<ViewUserSpendProps> = ({ 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<ViewUserSpendProps> = ({ 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<ViewUserSpendProps> = ({ 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 = [];