regen_tokens initial commit

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,7 @@ 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 +220,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,12 +240,17 @@ 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:
             return RedirectResponse("/cooldown")
         else:
@@ -259,6 +268,8 @@ 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 +320,20 @@ 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,7 @@ from modules.chat_processor.helpers import (
     update_user_info,
     get_time,
     check_user_cooldown,
+    reset_tokens_for_user,
 )
 import copy
 from typing import Optional
@@ -224,14 +225,14 @@ class Chatbot:
         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):
@@ -372,8 +373,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:
+                user.metadata["in_cooldown"] = True
                 # get time left in cooldown
                 # convert both to datetime objects
                 cooldown_end_time = datetime.fromisoformat(cooldown_end_time).replace(
@@ -481,9 +483,16 @@ class Chatbot:
         # # 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["all_time_tokens_allocated"] = user.metadata["all_time_tokens_allocated"] - token_count
+            await reset_tokens_for_user(user) # regenerate tokens for the user
         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
+        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,
             elements=source_elements,

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,14 +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
 
-    return False, None
+    user.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", 0)
+
+    # 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:
+        # Determine the time required to fully regenerate tokens from current state
+        if current_tokens < 0:
+            time_to_full_regen = REGEN_TIME * (1 + abs(current_tokens) / max_tokens) # more time to regenerate if tokens left is negative
+        else:
+            time_to_full_regen = REGEN_TIME * (1 - current_tokens / max_tokens) # less time to regenerate if tokens left is positive
+
+        # Calculate the proportion of this time that has elapsed
+        proportion_of_time_elapsed = min(elapsed_time_in_seconds / time_to_full_regen, 1.0)
+
+        # Calculate the tokens to regenerate based on the elapsed proportion
+        tokens_to_regenerate = int(proportion_of_time_elapsed * (max_tokens - current_tokens))
+
+        # Ensure tokens_to_regenerate is positive and doesn't exceed the maximum
+        if tokens_to_regenerate > 0:
+            new_token_count = min(current_tokens + tokens_to_regenerate, max_tokens)
+
+            print(f"\n\n Adding {tokens_to_regenerate} tokens to the user, Time left for full credits: {time_to_full_regen - elapsed_time_in_seconds} \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):
     step = await literal_client.api.get_step(thread_id)

code/modules/config/constants.py

@@ -4,8 +4,10 @@ import os
 load_dotenv()
 
 TIMEOUT = 60
-COOLDOWN_TIME = 60
-TOKENS_LEFT = 3000
+COOLDOWN_TIME = 180
+REGEN_TIME = 300
+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 }}";