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

8
litellm/proxy/_types.py
View file

@ -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]

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

1
ui/litellm-dashboard/src/components/user_dashboard.tsx
View file

@ -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}

37
ui/litellm-dashboard/src/components/view_user_spend.tsx
View file

@ -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 = [];