Merge pull request #79 from DL4DS/regen_tokens

code/app.py

@@ -14,6 +14,8 @@ from modules.config.constants import (
     CHAINLIT_URL,
     GITHUB_REPO,
     DOCS_WEBSITE,
+    ALL_TIME_TOKENS_ALLOCATED,
+    TOKENS_LEFT,
 )
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
@@ -207,7 +209,9 @@ async def cooldown(request: Request):
     user_info = await get_user_info_from_cookie(request)
     user_details = await get_user_details(user_info["email"])
     current_datetime = get_time()
-    cooldown, cooldown_end_time = check_user_cooldown(user_details, current_datetime)
+    cooldown, cooldown_end_time = await check_user_cooldown(
+        user_details, current_datetime
+    )
     print(f"User in cooldown: {cooldown}")
     print(f"Cooldown end time: {cooldown_end_time}")
     if cooldown and "admin" not in get_user_role(user_info["email"]):
@@ -218,9 +222,11 @@ async def cooldown(request: Request):
                 "username": user_info["email"],
                 "role": get_user_role(user_info["email"]),
                 "cooldown_end_time": cooldown_end_time,
+                "tokens_left": user_details.metadata["tokens_left"],
             },
         )
     else:
+        user_details.metadata["in_cooldown"] = False
         await update_user_info(user_details)
         await reset_tokens_for_user(user_details)
         return RedirectResponse("/post-signin")
@@ -236,15 +242,26 @@ async def post_signin(request: Request):
     user_details.metadata["last_login"] = current_datetime
     # 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["tokens_left"] = (
+            TOKENS_LEFT  # set the number of tokens left for the new user
+        )
+        if "all_time_tokens_allocated" not in user_details.metadata:
+            user_details.metadata["all_time_tokens_allocated"] = (
+                ALL_TIME_TOKENS_ALLOCATED
+            )
+        if "in_cooldown" not in user_details.metadata:
+            user_details.metadata["in_cooldown"] = False
+        await update_user_info(user_details)
 
     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)
+        cooldown, _ = await check_user_cooldown(user_details, current_datetime)
         if cooldown:
+            user_details.metadata["in_cooldown"] = True
             return RedirectResponse("/cooldown")
         else:
+            user_details.metadata["in_cooldown"] = False
             await reset_tokens_for_user(user_details)
 
     if user_info:
@@ -259,6 +276,10 @@ async def post_signin(request: Request):
                 "role": role,
                 "jwt_token": jwt_token,
                 "tokens_left": user_details.metadata["tokens_left"],
+                "all_time_tokens_allocated": user_details.metadata[
+                    "all_time_tokens_allocated"
+                ],
+                "total_tokens_allocated": ALL_TIME_TOKENS_ALLOCATED,
             },
         )
     return RedirectResponse("/")
@@ -309,6 +330,19 @@ async def logout(request: Request, response: Response):
     return response
 
 
+@app.get("/get-tokens-left")
+async def get_tokens_left(request: Request):
+    try:
+        user_info = await get_user_info_from_cookie(request)
+        user_details = await get_user_details(user_info["email"])
+        await reset_tokens_for_user(user_details)
+        tokens_left = user_details.metadata["tokens_left"]
+        return {"tokens_left": tokens_left}
+    except Exception as e:
+        print(f"Error getting tokens left: {e}")
+        return {"tokens_left": 0}
+
+
 mount_chainlit(app=app, target="main.py", path=CHAINLIT_PATH)
 
 if __name__ == "__main__":

code/main.py

@@ -20,6 +20,8 @@ from modules.chat_processor.helpers import (
     update_user_info,
     get_time,
     check_user_cooldown,
+    reset_tokens_for_user,
+    get_user_details,
 )
 import copy
 from typing import Optional
@@ -54,6 +56,24 @@ async def setup_data_layer():
     return data_layer
 
 
+async def update_user_from_chainlit(user, token_count=0):
+    if "admin" not in user.metadata["role"]:
+        user.metadata["tokens_left"] = user.metadata["tokens_left"] - token_count
+        user.metadata["all_time_tokens_allocated"] = (
+            user.metadata["all_time_tokens_allocated"] - token_count
+        )
+        user.metadata["tokens_left_at_last_message"] = user.metadata[
+            "tokens_left"
+        ]  # tokens_left will keep regenerating outside of chainlit
+    user.metadata["last_message_time"] = get_time()
+    await update_user_info(user)
+
+    tokens_left = user.metadata["tokens_left"]
+    if tokens_left < 0:
+        tokens_left = 0
+    return tokens_left
+
+
 class Chatbot:
     def __init__(self, config):
         """
@@ -221,17 +241,18 @@ class Chatbot:
             "view_sources": llm_settings.get("view_sources"),
             "follow_up_questions": llm_settings.get("follow_up_questions"),
         }
+        print("Settings Dict: ", settings_dict)
         await cl.Message(
             author=SYSTEM,
             content="LLM settings have been updated. You can continue with your Query!",
-            elements=[
-                cl.Text(
-                    name="settings",
-                    display="side",
-                    content=json.dumps(settings_dict, indent=4),
-                    language="json",
-                ),
-            ],
+            # elements=[
+            #     cl.Text(
+            #         name="settings",
+            #         display="side",
+            #         content=json.dumps(settings_dict, indent=4),
+            #         language="json",
+            #     ),
+            # ],
         ).send()
 
     async def set_starters(self):
@@ -365,6 +386,10 @@ class Chatbot:
 
         # update user info with last message time
         user = cl.user_session.get("user")
+        await reset_tokens_for_user(user)
+        updated_user = await get_user_details(user.identifier)
+        user.metadata = updated_user.metadata
+        cl.user_session.set("user", user)
 
         print("\n\n User Tokens Left: ", user.metadata["tokens_left"])
 
@@ -372,7 +397,9 @@ class Chatbot:
         # if not, return message saying they have run out of tokens
         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)
+            cooldown, cooldown_end_time = await check_user_cooldown(
+                user, current_datetime
+            )
             if cooldown:
                 # get time left in cooldown
                 # convert both to datetime objects
@@ -390,7 +417,6 @@ class Chatbot:
                 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 "
@@ -400,7 +426,20 @@ class Chatbot:
                     ),
                     author=SYSTEM,
                 ).send()
+                user.metadata["in_cooldown"] = True
+                await update_user_info(user)
                 return
+            else:
+                await cl.Message(
+                    content=(
+                        "Ah, seems like you don't have any tokens left...Please wait while we regenerate your tokens. Click "
+                        '<a href="/cooldown" style="color: #0000CD; text-decoration: none;" target="_self">here</a> to view your token credits.'
+                    ),
+                    author=SYSTEM,
+                ).send()
+                return
+
+        user.metadata["in_cooldown"] = False
 
         llm_settings = cl.user_session.get("llm_settings", {})
         view_sources = llm_settings.get("view_sources", False)
@@ -479,10 +518,13 @@ class Chatbot:
             print("Time taken to generate questions: ", time.time() - start_time)
 
         # # update user info with 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)
+        tokens_left = await update_user_from_chainlit(user, token_count)
+
+        answer_with_sources += (
+            '\n\n<footer><span style="font-size: 0.8em; text-align: right; display: block;">Tokens Left: '
+            + str(tokens_left)
+            + "</span></footer>\n"
+        )
 
         await cl.Message(
             content=answer_with_sources,

code/modules/chat_processor/helpers.py

@@ -1,7 +1,7 @@
 import os
 from literalai import AsyncLiteralClient
 from datetime import datetime, timedelta, timezone
-from modules.config.constants import COOLDOWN_TIME, TOKENS_LEFT
+from modules.config.constants import COOLDOWN_TIME, TOKENS_LEFT, REGEN_TIME
 from typing_extensions import TypedDict
 import tiktoken
 from typing import Any, Generic, List, Literal, Optional, TypeVar, Union
@@ -141,7 +141,7 @@ def get_time():
 
 
 async def get_user_details(user_email_id):
-    user_info = await literal_client.api.get_user(identifier=user_email_id)
+    user_info = await literal_client.api.get_or_create_user(identifier=user_email_id)
     return user_info
 
 
@@ -155,11 +155,11 @@ async def update_user_info(user_info):
     )
 
 
-def check_user_cooldown(user_info, current_time):
+async def check_user_cooldown(user_info, current_time):
 
-    # Check if no tokens left
+    # # Check if no tokens left
     tokens_left = user_info.metadata.get("tokens_left", 0)
-    if tokens_left > 0:
+    if tokens_left > 0 and not user_info.metadata.get("in_cooldown", False):
         return False, None
 
     user_info = convert_to_dict(user_info)
@@ -186,13 +186,54 @@ def check_user_cooldown(user_info, current_time):
     if elapsed_time_in_seconds < COOLDOWN_TIME:
         return True, cooldown_end_time_iso  # Return in ISO 8601 format
 
+    user_info["metadata"]["in_cooldown"] = False
+    # If not in cooldown, regenerate tokens
+    await reset_tokens_for_user(user_info)
+
     return False, None
 
 
 async def reset_tokens_for_user(user_info):
     user_info = convert_to_dict(user_info)
-    user_info["metadata"]["tokens_left"] = TOKENS_LEFT
-    await update_user_info(user_info)
+    last_message_time_str = user_info["metadata"].get("last_message_time")
+
+    last_message_time = datetime.fromisoformat(last_message_time_str).replace(
+        tzinfo=timezone.utc
+    )
+    current_time = datetime.fromisoformat(get_time()).replace(tzinfo=timezone.utc)
+
+    # Calculate the elapsed time since the last message
+    elapsed_time_in_seconds = (current_time - last_message_time).total_seconds()
+
+    # Current token count (can be negative)
+    current_tokens = user_info["metadata"].get("tokens_left_at_last_message", 0)
+    current_tokens = min(current_tokens, TOKENS_LEFT)
+
+    # Maximum tokens that can be regenerated
+    max_tokens = user_info["metadata"].get("max_tokens", TOKENS_LEFT)
+
+    # Calculate how many tokens should have been regenerated proportionally
+    if current_tokens < max_tokens:
+
+        # Calculate the regeneration rate per second based on REGEN_TIME for full regeneration
+        regeneration_rate_per_second = max_tokens / REGEN_TIME
+
+        # Calculate how many tokens should have been regenerated based on the elapsed time
+        tokens_to_regenerate = int(
+            elapsed_time_in_seconds * regeneration_rate_per_second
+        )
+
+        # Ensure the new token count does not exceed max_tokens
+        new_token_count = min(current_tokens + tokens_to_regenerate, max_tokens)
+
+        print(
+            f"\n\n Adding {tokens_to_regenerate} tokens to the user, Time elapsed: {elapsed_time_in_seconds} seconds, Tokens after regeneration: {new_token_count}, Tokens before: {current_tokens} \n\n"
+        )
+
+        # Update the user's token count
+        user_info["metadata"]["tokens_left"] = new_token_count
+
+        await update_user_info(user_info)
 
 
 async def get_thread_step_info(thread_id):

code/modules/config/constants.py

@@ -5,7 +5,9 @@ load_dotenv()
 
 TIMEOUT = 60
 COOLDOWN_TIME = 60
-TOKENS_LEFT = 3000
+REGEN_TIME = 180
+TOKENS_LEFT = 2000
+ALL_TIME_TOKENS_ALLOCATED = 1000000
 
 GITHUB_REPO = "https://github.com/DL4DS/dl4ds_tutor"
 DOCS_WEBSITE = "https://dl4ds.github.io/dl4ds_tutor/"

code/templates/cooldown.html

@@ -63,6 +63,13 @@
             margin-bottom: 30px;
         }
 
+        .tokens-left {
+            font-size: 14px;
+            color: #333;
+            margin-bottom: 30px;
+            font-weight: 600;
+        }
+
         .button {
             padding: 12px 0;
             margin: 12px 0;
@@ -70,33 +77,33 @@
             border-radius: 6px;
             cursor: pointer;
             width: 100%;
-            border: 1px solid #4285F4; /* Button border color */
-            background-color: #fff; /* Button background color */
-            color: #4285F4; /* Button text color */
+            border: 1px solid #4285F4;
+            background-color: #fff;
+            color: #4285F4;
             transition: background-color 0.3s ease, border-color 0.3s ease;
-            display: none; /* Initially hidden */
+            display: none;
         }
 
         .button.start-tutor {
-            display: none; /* Initially hidden */
+            display: none;
         }
 
         .button:hover {
             background-color: #e0e0e0;
-            border-color: #357ae8; /* Darker blue for hover */
+            border-color: #357ae8;
         }
 
         .sign-out-button {
             border: 1px solid #FF4C4C;
             background-color: #fff;
             color: #FF4C4C;
-            display: block; /* Ensure this button is always visible */
+            display: block;
         }
 
         .sign-out-button:hover {
-            background-color: #ffe6e6; /* Light red on hover */
-            border-color: #e04343; /* Darker red for hover */
-            color: #e04343; /* Red text on hover */
+            background-color: #ffe6e6;
+            border-color: #e04343;
+            color: #e04343;
         }
 
         #countdown {
@@ -104,6 +111,12 @@
             color: #555;
             margin-bottom: 20px;
         }
+
+        .footer {
+            font-size: 12px;
+            color: #777;
+            margin-top: 20px;
+        }
     </style>
 </head>
 <body>
@@ -113,24 +126,26 @@
         <p>It seems like you need to wait a bit before starting a new session.</p>
         <p class="cooldown-message">Time remaining until the cooldown period ends:</p>
         <p id="countdown"></p>
+        <p class="tokens-left">Tokens Left: <span id="tokensLeft">{{ tokens_left }}</span></p>
         <button id="startTutorBtn" class="button start-tutor" onclick="startTutor()">Start AI Tutor</button>
         <form action="/logout" method="get">
             <button type="submit" class="button sign-out-button">Sign Out</button>
         </form>
+        <div class="footer">Reload the page to update token stats</div>
     </div>
     <script>
         function startCountdown(endTime) {
             const countdownElement = document.getElementById('countdown');
             const startTutorBtn = document.getElementById('startTutorBtn');
-            const endTimeDate = new Date(endTime);  // Parse the cooldown end time
+            const endTimeDate = new Date(endTime);
 
             function updateCountdown() {
-                const now = new Date();  // Get the current time
-                const timeLeft = endTimeDate.getTime() - now.getTime();  // Calculate the remaining time in milliseconds
+                const now = new Date();
+                const timeLeft = endTimeDate.getTime() - now.getTime();
 
                 if (timeLeft <= 0) {
                     countdownElement.textContent = "Cooldown period has ended.";
-                    startTutorBtn.style.display = "block";  // Show the start tutor button
+                    startTutorBtn.style.display = "block";
                 } else {
                     const hours = Math.floor(timeLeft / 1000 / 60 / 60);
                     const minutes = Math.floor((timeLeft / 1000 / 60) % 60);
@@ -139,17 +154,28 @@
                 }
             }
 
-            updateCountdown();  // Initial call to set the countdown
-            setInterval(updateCountdown, 1000);  // Update the countdown every second
+            updateCountdown();
+            setInterval(updateCountdown, 1000);
         }
 
         function startTutor() {
-            // Redirect to AI Tutor session start or any other logic to start the tutor
             window.location.href = "/start-tutor";
         }
 
-        // Pass the cooldown_end_time to the script in ISO format with 'Z' to indicate UTC
+        function updateTokensLeft() {
+            fetch('/get-tokens-left')
+                .then(response => response.json())
+                .then(data => {
+                    document.getElementById('tokensLeft').textContent = data.tokens_left;
+                })
+                .catch(error => console.error('Error fetching tokens:', error));
+        }
+
+        // Start the countdown
         startCountdown("{{ cooldown_end_time }}");
+
+        // Update tokens left when the page loads
+        updateTokensLeft();
     </script>
 </body>
 </html>

code/templates/dashboard.html

@@ -27,7 +27,7 @@
             border-radius: 8px;
             width: 100%;
             max-width: 400px;
-            padding: 50px;
+            padding: 40px;
             box-sizing: border-box;
             text-align: center;
             box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
@@ -39,35 +39,43 @@
             width: 90px;
             height: 90px;
             border-radius: 50%;
-            margin-bottom: 25px;
+            margin-bottom: 20px;
             border: 2px solid #ddd;
         }
 
         .container h1 {
-            margin-bottom: 15px;
-            font-size: 24px;
+            margin-bottom: 20px;
+            font-size: 26px;
             font-weight: 600;
             color: #1a1a1a;
         }
 
         .container p {
-            font-size: 16px;
+            font-size: 15px;
             color: #4a4a4a;
-            margin-bottom: 30px;
+            margin-bottom: 25px;
             line-height: 1.5;
         }
 
         .tokens-left {
-            font-size: 16px;
+            font-size: 17px;
             color: #333;
-            margin-bottom: 30px;
+            margin-bottom: 10px;
             font-weight: 600;
         }
 
+        .all-time-tokens {
+            font-size: 14px; /* Reduced font size */
+            color: #555;
+            margin-bottom: 30px;
+            font-weight: 500;
+            white-space: nowrap; /* Prevents breaking to a new line */
+        }
+
         .button {
             padding: 12px 0;
             margin: 12px 0;
-            font-size: 14px;
+            font-size: 15px;
             border-radius: 6px;
             cursor: pointer;
             width: 100%;
@@ -105,6 +113,12 @@
             border-color: #e04343; /* Darker red for hover */
             color: #e04343; /* Red text on hover */
         }
+
+        .footer {
+            font-size: 12px;
+            color: #777;
+            margin-top: 25px;
+        }
     </style>
 </head>
 <body>
@@ -113,12 +127,14 @@
         <h1>Welcome, {{ username }}</h1>
         <p>Ready to start your AI tutoring session?</p>
         <p class="tokens-left">Tokens Left: {{ tokens_left }}</p>
+        <p class="all-time-tokens">All-Time Tokens Allocated: {{ all_time_tokens_allocated }} / {{ total_tokens_allocated }}</p>
         <form action="/start-tutor" method="post">
             <button type="submit" class="button start-button">Start AI Tutor</button>
         </form>
         <form action="/logout" method="get">
             <button type="submit" class="button sign-out-button">Sign Out</button>
         </form>
+        <div class="footer">Reload the page to update token stats</div>
     </div>
     <script>
         let token = "{{ jwt_token }}";