updates

code/app.py

@@ -43,9 +43,10 @@ templates = Jinja2Templates(directory="templates")
 session_store = {}
 CHAINLIT_PATH = "/chainlit_tutor"
 
+# only admin is given any additional permissions for now -- no limits on tokens
 USER_ROLES = {
     "[email protected]": ["instructor", "bu"],
-    "[email protected]": ["instructor", "bu"],
+    "[email protected]": ["admin", "instructor", "bu"],
     "[email protected]": ["instructor", "bu"],
     "[email protected]": ["guest"],
     # Add more users and roles as needed
@@ -168,6 +169,7 @@ async def auth_google(request: Request):
         email = user_info["email"]
         name = user_info.get("name", "")
         profile_image = user_info.get("picture", "")
+        role = get_user_role(email)
 
         session_token = secrets.token_hex(16)
         session_store[session_token] = {
@@ -180,6 +182,7 @@ async def auth_google(request: Request):
         # add literalai user info to session store to be sent to chainlit
         literalai_user = await get_user_details(email)
         session_store[session_token]["literalai_info"] = literalai_user.to_dict()
+        session_store[session_token]["literalai_info"]["metadata"]["role"] = role
 
         user_info_json = json.dumps(session_store[session_token])
         user_info_encoded = base64.b64encode(user_info_json.encode()).decode()
@@ -204,18 +207,17 @@ async def cooldown(request: Request):
     cooldown, cooldown_end_time = check_user_cooldown(user_details, current_datetime)
     print(f"User in cooldown: {cooldown}")
     print(f"Cooldown end time: {cooldown_end_time}")
-    if cooldown:
+    if cooldown and "admin" not in get_user_role(user_info["email"]):
         return templates.TemplateResponse(
             "cooldown.html",
             {
                 "request": request,
                 "username": user_info["email"],
                 "role": get_user_role(user_info["email"]),
-                "cooldown_end_time": cooldown_end_time,  # Pass as ISO format string
+                "cooldown_end_time": cooldown_end_time,
             },
         )
     else:
-        user_details.metadata["in_cooldown"] = False
         await update_user_info(user_details)
         await reset_tokens_for_user(user_details)
         return RedirectResponse("/post-signin")
@@ -232,14 +234,14 @@ async def post_signin(request: Request):
     # if new user, set the number of tries
     if "tokens_left" not in user_details.metadata:
         await reset_tokens_for_user(user_details)
-        user_details.metadata["all_time_tokens"] = 0  # Initialize all_time_tokens
-    if "last_message_time" in user_details.metadata:
+
+    if "last_message_time" in user_details.metadata and "admin" not in get_user_role(
+        user_info["email"]
+    ):
         cooldown, _ = check_user_cooldown(user_details, current_datetime)
         if cooldown:
             return RedirectResponse("/cooldown")
-        else:
-            user_details.metadata["in_cooldown"] = False
-            await reset_tokens_for_user(user_details)
+
     if user_info:
         username = user_info["email"]
         role = get_user_role(username)

code/main.py

@@ -27,6 +27,7 @@ from chainlit.types import ThreadDict
 import time
 import base64
 from langchain_community.callbacks import get_openai_callback
+from datetime import datetime, timezone
 
 USER_TIMEOUT = 60_000
 SYSTEM = "System"
@@ -365,21 +366,37 @@ class Chatbot:
         # update user info with last message time
         user = cl.user_session.get("user")
 
-        print("\n\n User Token Usage: ", user.metadata["all_time_tokens"])
         print("\n\n User Tokens Left: ", user.metadata["tokens_left"])
 
         # see if user has token credits left
         # if not, return message saying they have run out of tokens
-        if user.metadata["tokens_left"] <= 0:
+        if user.metadata["tokens_left"] <= 0 and "admin" not in user.metadata["role"]:
             current_datetime = get_time()
             cooldown, cooldown_end_time = check_user_cooldown(user, current_datetime)
             if cooldown:
-                user.metadata["in_cooldown"] = True
+                # get time left in cooldown
+                # convert both to datetime objects
+                cooldown_end_time = datetime.fromisoformat(cooldown_end_time).replace(
+                    tzinfo=timezone.utc
+                )
+                current_datetime = datetime.fromisoformat(current_datetime).replace(
+                    tzinfo=timezone.utc
+                )
+                cooldown_time_left = cooldown_end_time - current_datetime
+                # Get the total seconds
+                total_seconds = int(cooldown_time_left.total_seconds())
+                # Calculate hours, minutes, and seconds
+                hours, remainder = divmod(total_seconds, 3600)
+                minutes, seconds = divmod(remainder, 60)
+                # Format the time as 00 hrs 00 mins 00 secs
+                formatted_time = f"{hours:02} hrs {minutes:02} mins {seconds:02} secs"
                 await update_user_info(user)
                 await cl.Message(
                     content=(
                         "Ah, seems like you have run out of tokens...Click "
-                        '<a href="/cooldown" style="color: #0000CD; text-decoration: none;" target="_self">here</a> for more info.'
+                        '<a href="/cooldown" style="color: #0000CD; text-decoration: none;" target="_self">here</a> for more info. Please come back after {}'.format(
+                            formatted_time
+                        )
                     ),
                     author=SYSTEM,
                 ).send()
@@ -462,10 +479,8 @@ class Chatbot:
             print("Time taken to generate questions: ", time.time() - start_time)
 
         # # update user info with token count
-        user.metadata["tokens_left"] = user.metadata["tokens_left"] - token_count
-        user.metadata["all_time_tokens"] = (
-            user.metadata.get("all_time_tokens", 0) + token_count
-        )
+        if "admin" not in user.metadata["role"]:
+            user.metadata["tokens_left"] = user.metadata["tokens_left"] - token_count
         user.metadata["last_message_time"] = get_time()
         await update_user_info(user)
 
@@ -509,6 +524,10 @@ class Chatbot:
         ).decode()
         decoded_user_info = json.loads(decoded_user_info)
 
+        print(
+            f"\n\n USER ROLE: {decoded_user_info['literalai_info']['metadata']['role']} \n\n"
+        )
+
         return cl.User(
             id=decoded_user_info["literalai_info"]["id"],
             identifier=decoded_user_info["literalai_info"]["identifier"],
@@ -516,10 +535,11 @@ class Chatbot:
         )
 
     async def on_follow_up(self, action: cl.Action):
+        user = cl.user_session.get("user")
         message = await cl.Message(
             content=action.description,
             type="user_message",
-            author=self.user["user_id"],
+            author=user.identifier,
         ).send()
         async with cl.Step(
             name="on_follow_up", type="run", parent_id=message.id

code/modules/chat_processor/helpers.py

@@ -156,6 +156,12 @@ async def update_user_info(user_info):
 
 
 def check_user_cooldown(user_info, current_time):
+
+    # Check if no tokens left
+    tokens_left = user_info.metadata.get("tokens_left", 0)
+    if tokens_left > 0:
+        return False, None
+
     user_info = convert_to_dict(user_info)
     last_message_time_str = user_info["metadata"].get("last_message_time")
 

code/templates/cooldown.html

@@ -90,6 +90,7 @@
             border: 1px solid #FF4C4C;
             background-color: #fff;
             color: #FF4C4C;
+            display: block; /* Ensure this button is always visible */
         }
 
         .sign-out-button:hover {