feat(proxy_server.py): allow admin to invite users via invite link

Closes https://github.com/BerriAI/litellm/issues/3863
This commit is contained in:
Krrish Dholakia 2024-05-27 20:32:25 -07:00
parent 6b50e656b8
commit 86b66c13a4
5 changed files with 354 additions and 37 deletions

View file

@ -1154,3 +1154,28 @@ class WebhookEvent(CallInfo):
class SpecialModelNames(enum.Enum):
all_team_models = "all-team-models"
all_proxy_models = "all-proxy-models"
class InvitationNew(LiteLLMBase):
user_id: str
class InvitationUpdate(LiteLLMBase):
invitation_id: str
is_accepted: bool
class InvitationDelete(LiteLLMBase):
invitation_id: str
class InvitationModel(LiteLLMBase):
id: str
user_id: str
is_accepted: bool
accepted_at: Optional[datetime]
expires_at: datetime
created_at: datetime
created_by: str
updated_at: datetime
updated_by: str

View file

@ -10841,6 +10841,236 @@ async def auth_callback(request: Request):
return RedirectResponse(url=litellm_dashboard_ui)
#### INVITATION MANAGEMENT ####
@router.post(
"/invitation/new",
tags=["Invite Links"],
dependencies=[Depends(user_api_key_auth)],
response_model=InvitationModel,
)
async def new_invitation(
data: InvitationNew, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)
):
"""
Allow admin to create invite links, to onboard new users to Admin UI.
```
curl -X POST 'http://localhost:4000/invitation/new' \
-H 'Content-Type: application/json' \
-D '{
"user_id": "1234" // 👈 id of user in 'LiteLLM_UserTable'
}'
```
"""
global prisma_client
if prisma_client is None:
raise HTTPException(
status_code=400,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_role != "proxy_admin":
raise HTTPException(
status_code=400,
detail={
"error": "{}, your role={}".format(
CommonProxyErrors.not_allowed_access.value,
user_api_key_dict.user_role,
)
},
)
current_time = litellm.utils.get_utc_datetime()
expires_at = current_time + timedelta(days=7)
try:
response = await prisma_client.db.litellm_invitationlink.create(
data={
"user_id": data.user_id,
"created_at": current_time,
"expires_at": expires_at,
"created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
"updated_at": current_time,
"updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
} # type: ignore
)
except Exception as e:
if "Foreign key constraint failed on the field" in str(e):
raise HTTPException(
status_code=400,
detail={
"error": "User id does not exist in 'LiteLLM_UserTable'. Fix this by creating user via `/user/new`."
},
)
return response
@router.get(
"/invitation/info",
tags=["Invite Links"],
dependencies=[Depends(user_api_key_auth)],
response_model=InvitationModel,
)
async def invitation_info(
invitation_id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)
):
"""
Allow admin to create invite links, to onboard new users to Admin UI.
```
curl -X POST 'http://localhost:4000/invitation/new' \
-H 'Content-Type: application/json' \
-D '{
"user_id": "1234" // 👈 id of user in 'LiteLLM_UserTable'
}'
```
"""
global prisma_client
if prisma_client is None:
raise HTTPException(
status_code=400,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_role != "proxy_admin":
raise HTTPException(
status_code=400,
detail={
"error": "{}, your role={}".format(
CommonProxyErrors.not_allowed_access.value,
user_api_key_dict.user_role,
)
},
)
response = await prisma_client.db.litellm_invitationlink.find_unique(
where={"id": invitation_id}
)
if response is None:
raise HTTPException(
status_code=400,
detail={"error": "Invitation id does not exist in the database."},
)
return response
@router.post(
"/invitation/update",
tags=["Invite Links"],
dependencies=[Depends(user_api_key_auth)],
response_model=InvitationModel,
)
async def invitation_update(
data: InvitationUpdate,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Update when invitation is accepted
```
curl -X POST 'http://localhost:4000/invitation/update' \
-H 'Content-Type: application/json' \
-D '{
"invitation_id": "1234" // 👈 id of invitation in 'LiteLLM_InvitationTable'
"is_accepted": True // when invitation is accepted
}'
```
"""
global prisma_client
if prisma_client is None:
raise HTTPException(
status_code=400,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_id is None:
raise HTTPException(
status_code=500,
detail={
"error": "Unable to identify user id. Received={}".format(
user_api_key_dict.user_id
)
},
)
current_time = litellm.utils.get_utc_datetime()
response = await prisma_client.db.litellm_invitationlink.update(
where={"id": data.invitation_id},
data={
"id": data.invitation_id,
"is_accepted": data.is_accepted,
"accepted_at": current_time,
"updated_at": current_time,
"updated_by": user_api_key_dict.user_id, # type: ignore
},
)
if response is None:
raise HTTPException(
status_code=400,
detail={"error": "Invitation id does not exist in the database."},
)
return response
@router.post(
"/invitation/delete",
tags=["Invite Links"],
dependencies=[Depends(user_api_key_auth)],
response_model=InvitationModel,
)
async def invitation_delete(
data: InvitationDelete,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Delete invitation link
```
curl -X POST 'http://localhost:4000/invitation/delete' \
-H 'Content-Type: application/json' \
-D '{
"invitation_id": "1234" // 👈 id of invitation in 'LiteLLM_InvitationTable'
}'
```
"""
global prisma_client
if prisma_client is None:
raise HTTPException(
status_code=400,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
if user_api_key_dict.user_role != "proxy_admin":
raise HTTPException(
status_code=400,
detail={
"error": "{}, your role={}".format(
CommonProxyErrors.not_allowed_access.value,
user_api_key_dict.user_role,
)
},
)
response = await prisma_client.db.litellm_invitationlink.delete(
where={"id": data.invitation_id}
)
if response is None:
raise HTTPException(
status_code=400,
detail={"error": "Invitation id does not exist in the database."},
)
return response
#### CONFIG MANAGEMENT ####
@router.post(
"/config/update",

View file

@ -102,6 +102,7 @@ model LiteLLM_UserTable {
user_alias String?
team_id String?
organization_id String?
password String?
teams String[] @default([])
user_role String?
max_budget Float?
@ -117,6 +118,9 @@ model LiteLLM_UserTable {
model_spend Json @default("{}")
model_max_budget Json @default("{}")
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
invitations_created LiteLLM_InvitationLink[] @relation("CreatedBy")
invitations_updated LiteLLM_InvitationLink[] @relation("UpdatedBy")
invitations_user LiteLLM_InvitationLink[] @relation("UserId")
}
// Generate Tokens for Proxy
@ -221,4 +225,22 @@ model LiteLLM_TeamMembership {
budget_id String?
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
@@id([user_id, team_id])
}
model LiteLLM_InvitationLink {
// use this table to track invite links sent by admin for people to join the proxy
id String @id @default(uuid())
user_id String
is_accepted Boolean @default(false)
accepted_at DateTime? // when link is claimed (user successfully onboards via link)
expires_at DateTime // till when is link valid
created_at DateTime // when did admin create the link
created_by String // who created the link
updated_at DateTime // when was invite status updated
updated_by String // who updated the status (admin/user who accepted invite)
// Relations
liteLLM_user_table_user LiteLLM_UserTable @relation("UserId", fields: [user_id], references: [user_id])
liteLLM_user_table_created LiteLLM_UserTable @relation("CreatedBy", fields: [created_by], references: [user_id])
liteLLM_user_table_updated LiteLLM_UserTable @relation("UpdatedBy", fields: [updated_by], references: [user_id])
}

View file

@ -102,6 +102,7 @@ model LiteLLM_UserTable {
user_alias String?
team_id String?
organization_id String?
password String?
teams String[] @default([])
user_role String?
max_budget Float?
@ -117,6 +118,9 @@ model LiteLLM_UserTable {
model_spend Json @default("{}")
model_max_budget Json @default("{}")
litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id])
invitations_created LiteLLM_InvitationLink[] @relation("CreatedBy")
invitations_updated LiteLLM_InvitationLink[] @relation("UpdatedBy")
invitations_user LiteLLM_InvitationLink[] @relation("UserId")
}
// Generate Tokens for Proxy
@ -222,3 +226,21 @@ model LiteLLM_TeamMembership {
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
@@id([user_id, team_id])
}
model LiteLLM_InvitationLink {
// use this table to track invite links sent by admin for people to join the proxy
id String @id @default(uuid())
user_id String
is_accepted Boolean @default(false)
accepted_at DateTime? // when link is claimed (user successfully onboards via link)
expires_at DateTime // till when is link valid
created_at DateTime // when did admin create the link
created_by String // who created the link
updated_at DateTime // when was invite status updated
updated_by String // who updated the status (admin/user who accepted invite)
// Relations
liteLLM_user_table_user LiteLLM_UserTable @relation("UserId", fields: [user_id], references: [user_id])
liteLLM_user_table_created LiteLLM_UserTable @relation("CreatedBy", fields: [created_by], references: [user_id])
liteLLM_user_table_updated LiteLLM_UserTable @relation("UpdatedBy", fields: [updated_by], references: [user_id])
}

View file

@ -31,14 +31,13 @@ interface ModelHubProps {
}
interface ModelInfo {
model_group: string;
mode: string;
supports_function_calling: boolean;
supports_vision: boolean;
max_input_tokens?: number;
max_output_tokens?: number;
supported_openai_params?: string[];
model_group: string;
mode: string;
supports_function_calling: boolean;
supports_vision: boolean;
max_input_tokens?: number;
max_output_tokens?: number;
supported_openai_params?: string[];
// Add other properties if needed
}
@ -116,7 +115,7 @@ const ModelHub: React.FC<ModelHubProps> = ({
</Button>
</div>
<div className="grid grid-cols-2 gap-6 sm:grid-cols-3 lg:grid-cols-4">
<div className="grid grid-cols-2 gap-6 sm:grid-cols-3 lg:grid-cols-4 pr-5">
{modelHubData &&
modelHubData.map((model: ModelInfo) => (
<Card key={model.model_group} className="mt-5 mx-8">
@ -129,14 +128,27 @@ const ModelHub: React.FC<ModelHubProps> = ({
/>
</Tooltip>
</pre>
<div className='my-5'>
<Text>Mode: {model.mode}</Text>
<Text>Supports Function Calling: {model?.supports_function_calling == true ? "Yes" : "No"}</Text>
<Text>Supports Vision: {model?.supports_vision == true ? "Yes" : "No"}</Text>
<Text>Max Input Tokens: {model?.max_input_tokens ? model?.max_input_tokens : "N/A"}</Text>
<Text>Max Output Tokens: {model?.max_output_tokens ? model?.max_output_tokens : "N/A"}</Text>
</div>
<div className="my-5">
<Text>Mode: {model.mode}</Text>
<Text>
Supports Function Calling:{" "}
{model?.supports_function_calling == true ? "Yes" : "No"}
</Text>
<Text>
Supports Vision:{" "}
{model?.supports_vision == true ? "Yes" : "No"}
</Text>
<Text>
Max Input Tokens:{" "}
{model?.max_input_tokens ? model?.max_input_tokens : "N/A"}
</Text>
<Text>
Max Output Tokens:{" "}
{model?.max_output_tokens
? model?.max_output_tokens
: "N/A"}
</Text>
</div>
<div style={{ marginTop: "auto", textAlign: "right" }}>
<a
href="#"
@ -152,7 +164,11 @@ const ModelHub: React.FC<ModelHubProps> = ({
</div>
<Modal
title={selectedModel && selectedModel.model_group ? selectedModel.model_group : "Unknown Model"}
title={
selectedModel && selectedModel.model_group
? selectedModel.model_group
: "Unknown Model"
}
width={800}
visible={isModalVisible}
footer={null}
@ -161,19 +177,21 @@ const ModelHub: React.FC<ModelHubProps> = ({
>
{selectedModel && (
<div>
<p className='mb-4'><strong>Model Information & Usage</strong></p>
<p className="mb-4">
<strong>Model Information & Usage</strong>
</p>
<TabGroup>
<TabList>
<Tab>OpenAI Python SDK</Tab>
<Tab>Supported OpenAI Params</Tab>
<Tab>LlamaIndex</Tab>
<Tab>Langchain Py</Tab>
</TabList>
<TabPanels>
<TabPanel>
<SyntaxHighlighter language="python">
{`
<TabList>
<Tab>OpenAI Python SDK</Tab>
<Tab>Supported OpenAI Params</Tab>
<Tab>LlamaIndex</Tab>
<Tab>Langchain Py</Tab>
</TabList>
<TabPanels>
<TabPanel>
<SyntaxHighlighter language="python">
{`
import openai
client = openai.OpenAI(
api_key="your_api_key",
@ -192,13 +210,13 @@ response = client.chat.completions.create(
print(response)
`}
</SyntaxHighlighter>
</TabPanel>
<TabPanel>
<SyntaxHighlighter language="python">
{`${selectedModel.supported_openai_params?.map((param) => `${param}\n`).join('')}`}
</SyntaxHighlighter>
</TabPanel>
</SyntaxHighlighter>
</TabPanel>
<TabPanel>
<SyntaxHighlighter language="python">
{`${selectedModel.supported_openai_params?.map((param) => `${param}\n`).join("")}`}
</SyntaxHighlighter>
</TabPanel>
<TabPanel>
<SyntaxHighlighter language="python">
{`