Spaces:
Configuration error
Configuration error
File size: 3,779 Bytes
447ebeb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
from typing import Any, Union, Optional
from litellm.proxy._types import (
GenerateKeyRequest,
LiteLLM_ManagementEndpoint_MetadataFields_Premium,
LiteLLM_TeamTable,
LitellmUserRoles,
UserAPIKeyAuth,
)
from litellm.proxy.utils import _premium_user_check
def _user_has_admin_view(user_api_key_dict: UserAPIKeyAuth) -> bool:
return (
user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
or user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY
)
def _is_user_team_admin(
user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable
) -> bool:
for member in team_obj.members_with_roles:
if (
member.user_id is not None and member.user_id == user_api_key_dict.user_id
) and member.role == "admin":
return True
return False
def _set_object_metadata_field(
object_data: Union[LiteLLM_TeamTable, GenerateKeyRequest],
field_name: str,
value: Any,
) -> None:
"""
Helper function to set metadata fields that require premium user checks
Args:
object_data: The team data object to modify
field_name: Name of the metadata field to set
value: Value to set for the field
"""
if field_name in LiteLLM_ManagementEndpoint_MetadataFields_Premium:
_premium_user_check()
object_data.metadata = object_data.metadata or {}
object_data.metadata[field_name] = value
async def _upsert_budget_and_membership(
tx,
*,
team_id: str,
user_id: str,
max_budget: Optional[float],
existing_budget_id: Optional[str],
user_api_key_dict: UserAPIKeyAuth,
):
"""
Helper function to Create/Update or Delete the budget within the team membership
Args:
tx: The transaction object
team_id: The ID of the team
user_id: The ID of the user
max_budget: The maximum budget for the team
existing_budget_id: The ID of the existing budget, if any
user_api_key_dict: User API Key dictionary containing user information
If max_budget is None, the user's budget is removed from the team membership.
If max_budget exists, a budget is updated or created and linked to the team membership.
"""
if max_budget is None:
# disconnect the budget since max_budget is None
await tx.litellm_teammembership.update(
where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}},
data={"litellm_budget_table": {"disconnect": True}},
)
return
if existing_budget_id:
# update the existing budget
await tx.litellm_budgettable.update(
where={"budget_id": existing_budget_id},
data={"max_budget": max_budget},
)
return
# create a new budget
new_budget = await tx.litellm_budgettable.create(
data={
"max_budget": max_budget,
"created_by": user_api_key_dict.user_id or "",
"updated_by": user_api_key_dict.user_id or "",
},
include={"team_membership": True},
)
# upsert the team membership with the new/updated budget
await tx.litellm_teammembership.upsert(
where={
"user_id_team_id": {
"user_id": user_id,
"team_id": team_id,
}
},
data={
"create": {
"user_id": user_id,
"team_id": team_id,
"litellm_budget_table": {
"connect": {"budget_id": new_budget.budget_id},
},
},
"update": {
"litellm_budget_table": {
"connect": {"budget_id": new_budget.budget_id},
},
},
},
)
|