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

Closes #3863
BerriAI · May 28, 2024 · 86b66c1 · 86b66c1
1 parent 6b50e65
commit 86b66c1
Show file tree

Hide file tree

Showing 5 changed files with 354 additions and 37 deletions.
diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
@@ -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
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
@@ -10841,6 +10841,236 @@ def response_convertor(response, client):
     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",

diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
@@ -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])
 }
diff --git a/schema.prisma b/schema.prisma
@@ -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])
+}