Implement user roles

Author: bystepii

libs/ktem/ktem/db/base_models.py

@@ -1,6 +1,7 @@
 import datetime
 import uuid
 from typing import Optional
+from enum import Enum
 
 from sqlalchemy import JSON, Column
 from sqlmodel import Field, SQLModel
@@ -43,6 +44,10 @@ class BaseConversation(SQLModel):
         default_factory=lambda: datetime.datetime.now(get_localzone())
     )
 
+class Role(str, Enum):
+    CHAT_USER = "chat_user"
+    AGENT_CREATOR = "agent_creator"
+    ADMIN = "admin"
 
 class BaseUser(SQLModel):
     """Store the user information
@@ -61,7 +66,7 @@ class BaseUser(SQLModel):
     username: str = Field(unique=True)
     username_lower: str = Field(unique=True)
     password: str
-    admin: bool = Field(default=False)
+    role: Role = Field(default=Role.CHAT_USER)
 
 
 class BaseSettings(SQLModel):

libs/ktem/ktem/main.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from decouple import config
 from ktem.app import BaseApp
+from ktem.db.base_models import Role
 from ktem.pages.chat import ChatPage
 from ktem.pages.help import HelpPage
 from ktem.pages.resources import ResourcesTab
@@ -153,7 +154,7 @@ def toggle_login_visibility(user_id):
                             for k in self._tabs.keys()
                         )
 
-                    is_admin = user.admin
+                    is_admin = user.role == Role.ADMIN
 
                 tabs_update = []
                 for k in self._tabs.keys():

libs/ktem/ktem/pages/resources/__init__.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from ktem.app import BasePage
+from ktem.db.base_models import Role
 from ktem.db.models import User, engine
 from ktem.embeddings.ui import EmbeddingManagement
 from ktem.index.ui import IndexManagement
@@ -58,7 +59,7 @@ def toggle_user_management(self, user_id):
         """Show/hide the user management, depending on the user's role"""
         with Session(engine) as session:
             user = session.exec(select(User).where(User.id == user_id)).first()
-            if user and user.admin:
+            if user and user.role == Role.ADMIN:
                 return gr.update(visible=True)
 
             return gr.update(visible=False)

libs/ktem/ktem/pages/resources/user.py

@@ -3,6 +3,7 @@
 import gradio as gr
 import pandas as pd
 from ktem.app import BasePage
+from ktem.db.base_models import Role
 from ktem.db.models import User, engine
 from sqlmodel import Session, select
 from theflow.settings import settings as flowsettings
@@ -109,7 +110,7 @@ def create_user(usn, pwd, user_id=None, is_admin=True) -> bool:
                 username=usn,
                 username_lower=usn.lower(),
                 password=hashed_password,
-                admin=is_admin,
+                role=Role.ADMIN if is_admin else Role.CHAT_USER,
             )
             session.add(user)
             session.commit()
@@ -136,7 +137,7 @@ def on_building_ui(self):
         with gr.Tab(label="User list"):
             self.state_user_list = gr.State(value=None)
             self.user_list = gr.DataFrame(
-                headers=["id", "name", "admin"],
+                headers=["id", "name", "role"],
                 column_widths=[0, 50, 50],
                 interactive=False,
             )
@@ -150,7 +151,11 @@ def on_building_ui(self):
                         label="Confirm change password",
                         type="password",
                     )
-                self.admin_edit = gr.Checkbox(label="Admin")
+                self.role_edit = gr.Dropdown(
+                    label="Role",
+                    choices=[role.value for role in Role],
+                    value=Role.CHAT_USER.value,
+                )
 
             with gr.Row(visible=False) as self._selected_panel_btn:
                 with gr.Column():
@@ -208,7 +213,7 @@ def on_register_events(self):
                 self.usn_edit,
                 self.pwd_edit,
                 self.pwd_cnf_edit,
-                self.admin_edit,
+                self.role_edit,
             ],
             show_progress="hidden",
         )
@@ -245,7 +250,7 @@ def on_register_events(self):
                 self.usn_edit,
                 self.pwd_edit,
                 self.pwd_cnf_edit,
-                self.admin_edit,
+                self.role_edit,
             ],
             outputs=[self.pwd_edit, self.pwd_cnf_edit],
             show_progress="hidden",
@@ -315,27 +320,27 @@ def create_user(self, usn, pwd, pwd_cnf):
     def list_users(self, user_id):
         if user_id is None:
             return [], pd.DataFrame.from_records(
-                [{"id": "-", "username": "-", "admin": "-"}]
+                [{"id": "-", "username": "-", "role": "-"}]
             )
 
         with Session(engine) as session:
             statement = select(User).where(User.id == user_id)
             user = session.exec(statement).one()
-            if not user.admin:
+            if not user.role == Role.ADMIN:
                 return [], pd.DataFrame.from_records(
-                    [{"id": "-", "username": "-", "admin": "-"}]
+                    [{"id": "-", "username": "-", "role": "-"}]
                 )
 
             statement = select(User)
             results = [
-                {"id": user.id, "username": user.username, "admin": user.admin}
+                {"id": user.id, "username": user.username, "role": user.role.value}
                 for user in session.exec(statement).all()
             ]
             if results:
                 user_list = pd.DataFrame.from_records(results)
             else:
                 user_list = pd.DataFrame.from_records(
-                    [{"id": "-", "username": "-", "admin": "-"}]
+                    [{"id": "-", "username": "-", "role": "-"}]
                 )
 
         return results, user_list
@@ -360,7 +365,7 @@ def on_selected_user_change(self, selected_user_id):
             usn_edit = gr.update(value="")
             pwd_edit = gr.update(value="")
             pwd_cnf_edit = gr.update(value="")
-            admin_edit = gr.update(value=False)
+            role_edit = gr.update(value=Role.CHAT_USER.value)
         else:
             _selected_panel = gr.update(visible=True)
             _selected_panel_btn = gr.update(visible=True)
@@ -375,7 +380,7 @@ def on_selected_user_change(self, selected_user_id):
             usn_edit = gr.update(value=user.username)
             pwd_edit = gr.update(value="")
             pwd_cnf_edit = gr.update(value="")
-            admin_edit = gr.update(value=user.admin)
+            role_edit = gr.update(value=user.role.value)
 
         return (
             _selected_panel,
@@ -386,7 +391,7 @@ def on_selected_user_change(self, selected_user_id):
             usn_edit,
             pwd_edit,
             pwd_cnf_edit,
-            admin_edit,
+            role_edit,
         )
 
     def on_btn_delete_click(self, selected_user_id):
@@ -403,7 +408,7 @@ def on_btn_delete_click(self, selected_user_id):
 
         return btn_delete, btn_delete_yes, btn_delete_no
 
-    def save_user(self, selected_user_id, usn, pwd, pwd_cnf, admin):
+    def save_user(self, selected_user_id, usn, pwd, pwd_cnf, role):
         errors = validate_username(usn)
         if errors:
             gr.Warning(errors)
@@ -420,7 +425,7 @@ def save_user(self, selected_user_id, usn, pwd, pwd_cnf, admin):
             user = session.exec(statement).one()
             user.username = usn
             user.username_lower = usn.lower()
-            user.admin = admin
+            user.role = Role(role)
             if pwd:
                 user.password = hashlib.sha256(pwd.encode()).hexdigest()
             session.commit()