updates

code/main.py

@@ -21,6 +21,7 @@ from modules.chat_processor.helpers import (
     get_time,
     check_user_cooldown,
     reset_tokens_for_user,
+    get_user_details,
 )
 import copy
 from typing import Optional
@@ -55,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):
         """
@@ -367,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"])
 
@@ -406,6 +429,15 @@ class Chatbot:
                 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
 
@@ -486,18 +518,8 @@ 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["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 = await update_user_from_chainlit(user, token_count)
 
-        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)

code/modules/chat_processor/helpers.py

@@ -206,46 +206,34 @@ async def reset_tokens_for_user(user_info):
     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)
+    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:
-        # Determine the time required to fully regenerate tokens from the 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 = elapsed_time_in_seconds / time_to_full_regen
-
-        # Ensure the proportion doesn't exceed 1.0
-        proportion_of_time_elapsed = min(proportion_of_time_elapsed, 1.0)
-
-        # Calculate the tokens to regenerate based on the elapsed proportion
+
+        # 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(
-            proportion_of_time_elapsed * (max_tokens - current_tokens)
+            elapsed_time_in_seconds * regeneration_rate_per_second
         )
 
-        # 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)
+        # 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 left for full credits: {max(0, time_to_full_regen - elapsed_time_in_seconds)} \n\n"
-            )
+        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
+        # Update the user's token count
+        user_info["metadata"]["tokens_left"] = new_token_count
 
-            await update_user_info(user_info)
+        await update_user_info(user_info)
 
 
 async def get_thread_step_info(thread_id):

code/modules/config/constants.py

@@ -4,8 +4,8 @@ import os
 load_dotenv()
 
 TIMEOUT = 60
-COOLDOWN_TIME = 180
-REGEN_TIME = 300
+COOLDOWN_TIME = 60
+REGEN_TIME = 180
 TOKENS_LEFT = 2000
 ALL_TIME_TOKENS_ALLOCATED = 1000000