Merge pull request #75 from DL4DS/tracking_tokens

code/.chainlit/config.toml

@@ -20,7 +20,7 @@ allow_origins = ["*"]
 
 [features]
 # Process and display HTML in messages. This can be a security risk (see https://stackoverflow.com/questions/19603097/why-is-it-dangerous-to-render-user-generated-html-or-javascript)
-unsafe_allow_html = false
+unsafe_allow_html = true
 
 # Process and display mathematical expressions. This can clash with "$" characters in messages.
 latex = true

code/app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, Request, Response
+from fastapi import FastAPI, Request, Response, HTTPException
 from fastapi.responses import HTMLResponse, RedirectResponse
 from fastapi.templating import Jinja2Templates
 from google.oauth2 import id_token
@@ -12,9 +12,18 @@ from modules.config.constants import (
     OAUTH_GOOGLE_CLIENT_ID,
     OAUTH_GOOGLE_CLIENT_SECRET,
     CHAINLIT_URL,
+    GITHUB_REPO,
+    DOCS_WEBSITE,
 )
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
+from modules.chat_processor.helpers import (
+    get_user_details,
+    get_time,
+    reset_tokens_for_user,
+    check_user_cooldown,
+    update_user_info,
+)
 
 GOOGLE_CLIENT_ID = OAUTH_GOOGLE_CLIENT_ID
 GOOGLE_CLIENT_SECRET = OAUTH_GOOGLE_CLIENT_SECRET
@@ -34,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
@@ -71,7 +81,7 @@ def get_user_role(username: str):
     return USER_ROLES.get(username, ["student"])  # Default to "student" role
 
 
-def get_user_info_from_cookie(request: Request):
+async def get_user_info_from_cookie(request: Request):
     user_info_encoded = request.cookies.get("X-User-Info")
     if user_info_encoded:
         try:
@@ -83,6 +93,14 @@ def get_user_info_from_cookie(request: Request):
     return None
 
 
+async def del_user_info_from_cookie(request: Request, response: Response):
+    response.delete_cookie("X-User-Info")
+    response.delete_cookie("session_token")
+    session_token = request.cookies.get("session_token")
+    if session_token:
+        del session_store[session_token]
+
+
 def get_user_info(request: Request):
     session_token = request.cookies.get("session_token")
     if session_token and session_token in session_store:
@@ -92,33 +110,35 @@ def get_user_info(request: Request):
 
 @app.get("/", response_class=HTMLResponse)
 async def login_page(request: Request):
-    user_info = get_user_info_from_cookie(request)
+    user_info = await get_user_info_from_cookie(request)
     if user_info and user_info.get("google_signed_in"):
         return RedirectResponse("/post-signin")
-    return templates.TemplateResponse("login.html", {"request": request})
-
-
-@app.get("/login/guest")
-@app.post("/login/guest")
-async def login_guest():
-    username = "guest"
-    session_token = secrets.token_hex(16)
-    unique_session_id = secrets.token_hex(8)
-    username = f"{username}_{unique_session_id}"
-    session_store[session_token] = {
-        "email": username,
-        "name": "Guest",
-        "profile_image": "",
-        "google_signed_in": False,  # Ensure guest users do not have this flag
-    }
-    user_info_json = json.dumps(session_store[session_token])
-    user_info_encoded = base64.b64encode(user_info_json.encode()).decode()
-
-    # Set cookies
-    response = RedirectResponse(url="/post-signin", status_code=303)
-    response.set_cookie(key="session_token", value=session_token)
-    response.set_cookie(key="X-User-Info", value=user_info_encoded, httponly=True)
-    return response
+    return templates.TemplateResponse(
+        "login.html",
+        {"request": request, "GITHUB_REPO": GITHUB_REPO, "DOCS_WEBSITE": DOCS_WEBSITE},
+    )
+
+
+# @app.get("/login/guest")
+# async def login_guest():
+#     username = "guest"
+#     session_token = secrets.token_hex(16)
+#     unique_session_id = secrets.token_hex(8)
+#     username = f"{username}_{unique_session_id}"
+#     session_store[session_token] = {
+#         "email": username,
+#         "name": "Guest",
+#         "profile_image": "",
+#         "google_signed_in": False,  # Ensure guest users do not have this flag
+#     }
+#     user_info_json = json.dumps(session_store[session_token])
+#     user_info_encoded = base64.b64encode(user_info_json.encode()).decode()
+
+#     # Set cookies
+#     response = RedirectResponse(url="/post-signin", status_code=303)
+#     response.set_cookie(key="session_token", value=session_token)
+#     response.set_cookie(key="X-User-Info", value=user_info_encoded, httponly=True)
+#     return response
 
 
 @app.get("/login/google")
@@ -128,8 +148,7 @@ async def login_google(request: Request):
     response.delete_cookie(key="session_token")
     response.delete_cookie(key="X-User-Info")
 
-    user_info = get_user_info_from_cookie(request)
-    print(f"User info: {user_info}")
+    user_info = await get_user_info_from_cookie(request)
     # Check if user is already signed in using Google
     if user_info and user_info.get("google_signed_in"):
         return RedirectResponse("/post-signin")
@@ -150,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] = {
@@ -159,25 +179,71 @@ async def auth_google(request: Request):
             "google_signed_in": True,  # Set this flag to True for Google-signed users
         }
 
+        # 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()
 
         # Set cookies
         response = RedirectResponse(url="/post-signin", status_code=303)
         response.set_cookie(key="session_token", value=session_token)
-        response.set_cookie(key="X-User-Info", value=user_info_encoded, httponly=True)
+        response.set_cookie(
+            key="X-User-Info", value=user_info_encoded
+        )  # TODO: is the flag httponly=True necessary?
         return response
     except Exception as e:
         print(f"Error during Google OAuth callback: {e}")
         return RedirectResponse(url="/", status_code=302)
 
 
+@app.get("/cooldown")
+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)
+    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"]):
+        return templates.TemplateResponse(
+            "cooldown.html",
+            {
+                "request": request,
+                "username": user_info["email"],
+                "role": get_user_role(user_info["email"]),
+                "cooldown_end_time": cooldown_end_time,
+            },
+        )
+    else:
+        await update_user_info(user_details)
+        await reset_tokens_for_user(user_details)
+        return RedirectResponse("/post-signin")
+
+
 @app.get("/post-signin", response_class=HTMLResponse)
 async def post_signin(request: Request):
-    user_info = get_user_info_from_cookie(request)
+    user_info = await get_user_info_from_cookie(request)
     if not user_info:
         user_info = get_user_info(request)
-    # if user_info and user_info.get("google_signed_in"):
+    user_details = await get_user_details(user_info["email"])
+    current_datetime = get_time()
+    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)
+
+    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:
+            await reset_tokens_for_user(user_details)
+
     if user_info:
         username = user_info["email"]
         role = get_user_role(username)
@@ -189,14 +255,16 @@ async def post_signin(request: Request):
                 "username": username,
                 "role": role,
                 "jwt_token": jwt_token,
+                "tokens_left": user_details.metadata["tokens_left"],
             },
         )
     return RedirectResponse("/")
 
 
+@app.get("/start-tutor")
 @app.post("/start-tutor")
 async def start_tutor(request: Request):
-    user_info = get_user_info_from_cookie(request)
+    user_info = await get_user_info_from_cookie(request)
     if user_info:
         user_info_json = json.dumps(user_info)
         user_info_encoded = base64.b64encode(user_info_json.encode()).decode()
@@ -208,6 +276,19 @@ async def start_tutor(request: Request):
     return RedirectResponse(url="/")
 
 
+@app.exception_handler(HTTPException)
+async def http_exception_handler(request: Request, exc: HTTPException):
+    if exc.status_code == 404:
+        return templates.TemplateResponse(
+            "error_404.html", {"request": request}, status_code=404
+        )
+    return templates.TemplateResponse(
+        "error.html",
+        {"request": request, "error": str(exc)},
+        status_code=exc.status_code,
+    )
+
+
 @app.exception_handler(Exception)
 async def exception_handler(request: Request, exc: Exception):
     return templates.TemplateResponse(
@@ -215,17 +296,14 @@ async def exception_handler(request: Request, exc: Exception):
     )
 
 
-@app.get("/chainlit_tutor/logout", response_class=HTMLResponse)
-@app.post("/chainlit_tutor/logout", response_class=HTMLResponse)
-async def app_logout(request: Request, response: Response):
-    # Clear session cookies
-    response.delete_cookie("session_token")
-    response.delete_cookie("X-User-Info")
-
-    print("logout_page called")
-
-    # Redirect to the logout page with embedded JavaScript
-    return RedirectResponse(url="/", status_code=302)
+@app.get("/logout", response_class=HTMLResponse)
+async def logout(request: Request, response: Response):
+    await del_user_info_from_cookie(request=request, response=response)
+    response = RedirectResponse(url="/", status_code=302)
+    # Set cookies to empty values and expire them immediately
+    response.set_cookie(key="session_token", value="", expires=0)
+    response.set_cookie(key="X-User-Info", value="", expires=0)
+    return response
 
 
 mount_chainlit(app=app, target="main.py", path=CHAINLIT_PATH)

code/chainlit.md

@@ -1,10 +1,5 @@
 # Welcome to DL4DS Tutor! 🚀🤖
 
-Hi there, this is an LLM chatbot designed to help answer questions on the course content, built using Langchain and Chainlit.
-This is still very much a Work in Progress.
+Hi there, this is an LLM chatbot designed to help answer questions on the course content.
 
 ### --- Please wait while the Tutor loads... ---
-
-## Useful Links 🔗
-
-- **Documentation:**  [Chainlit Documentation](https://docs.chainlit.io) 📚

code/main.py

@@ -16,11 +16,18 @@ from modules.chat.helpers import (
     get_history_setup_llm,
     get_last_config,
 )
+from modules.chat_processor.helpers import (
+    update_user_info,
+    get_time,
+    check_user_cooldown,
+)
 import copy
 from typing import Optional
 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"
@@ -291,9 +298,9 @@ class Chatbot:
 
         await self.make_llm_settings_widgets(self.config)  # Reload the settings widgets
 
-        await self.make_llm_settings_widgets(self.config)
         user = cl.user_session.get("user")
 
+        # TODO: remove self.user with cl.user_session.get("user")
         try:
             self.user = {
                 "user_id": user.identifier,
@@ -307,8 +314,6 @@ class Chatbot:
             }
 
         memory = cl.user_session.get("memory", [])
-
-        cl.user_session.set("user", self.user)
         self.llm_tutor = LLMTutor(self.config, user=self.user)
 
         self.chain = self.llm_tutor.qa_bot(
@@ -353,6 +358,49 @@ class Chatbot:
         start_time = time.time()
 
         chain = cl.user_session.get("chain")
+        token_count = 0  # initialize token count
+        if not chain:
+            await self.start()  # start the chatbot if the chain is not present
+            chain = cl.user_session.get("chain")
+
+        # update user info with last message time
+        user = cl.user_session.get("user")
+
+        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 and "admin" not in user.metadata["role"]:
+            current_datetime = get_time()
+            cooldown, cooldown_end_time = check_user_cooldown(user, current_datetime)
+            if cooldown:
+                # 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. Please come back after {}'.format(
+                            formatted_time
+                        )
+                    ),
+                    author=SYSTEM,
+                ).send()
+                return
 
         llm_settings = cl.user_session.get("llm_settings", {})
         view_sources = llm_settings.get("view_sources", False)
@@ -360,6 +408,7 @@ class Chatbot:
         stream = False  # Fix streaming
         user_query_dict = {"input": message.content}
         # Define the base configuration
+        cb = cl.AsyncLangchainCallbackHandler()
         chain_config = {
             "configurable": {
                 "user_id": self.user["user_id"],
@@ -367,20 +416,22 @@ class Chatbot:
                 "memory_window": self.config["llm_params"]["memory_window"],
             },
             "callbacks": (
-                [cl.LangchainCallbackHandler()]
+                [cb]
                 if cl_data._data_layer and self.config["chat_logging"]["callbacks"]
                 else None
             ),
         }
 
-        if stream:
-            res = chain.stream(user_query=user_query_dict, config=chain_config)
-            res = await self.stream_response(res)
-        else:
-            res = await chain.invoke(
-                user_query=user_query_dict,
-                config=chain_config,
-            )
+        with get_openai_callback() as token_count_cb:
+            if stream:
+                res = chain.stream(user_query=user_query_dict, config=chain_config)
+                res = await self.stream_response(res)
+            else:
+                res = await chain.invoke(
+                    user_query=user_query_dict,
+                    config=chain_config,
+                )
+        token_count += token_count_cb.total_tokens
 
         answer = res.get("answer", res.get("result"))
 
@@ -395,21 +446,24 @@ class Chatbot:
 
         if self.config["llm_params"]["generate_follow_up"]:
             start_time = time.time()
+            cb_follow_up = cl.AsyncLangchainCallbackHandler()
             config = {
                 "callbacks": (
-                    [cl.LangchainCallbackHandler()]
+                    [cb_follow_up]
                     if cl_data._data_layer and self.config["chat_logging"]["callbacks"]
                     else None
                 )
             }
+            with get_openai_callback() as token_count_cb:
+                list_of_questions = await self.question_generator.generate_questions(
+                    query=user_query_dict["input"],
+                    response=answer,
+                    chat_history=res.get("chat_history"),
+                    context=res.get("context"),
+                    config=config,
+                )
 
-            list_of_questions = await self.question_generator.generate_questions(
-                query=user_query_dict["input"],
-                response=answer,
-                chat_history=res.get("chat_history"),
-                context=res.get("context"),
-                config=config,
-            )
+            token_count += token_count_cb.total_tokens
 
             for question in list_of_questions:
 
@@ -424,6 +478,12 @@ 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)
+
         await cl.Message(
             content=answer_with_sources,
             elements=source_elements,
@@ -445,15 +505,6 @@ class Chatbot:
         cl.user_session.set("memory", conversation_list)
         await self.start(config=thread_config)
 
-    # @cl.oauth_callback
-    # def auth_callback(
-    #     provider_id: str,
-    #     token: str,
-    #     raw_user_data: Dict[str, str],
-    #     default_user: cl.User,
-    # ) -> Optional[cl.User]:
-    #     return default_user
-
     @cl.header_auth_callback
     def header_auth_callback(headers: dict) -> Optional[cl.User]:
 
@@ -473,19 +524,22 @@ 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(
-            identifier=decoded_user_info["email"],
-            metadata={
-                "name": decoded_user_info["name"],
-                "avatar": decoded_user_info["profile_image"],
-            },
+            id=decoded_user_info["literalai_info"]["id"],
+            identifier=decoded_user_info["literalai_info"]["identifier"],
+            metadata=decoded_user_info["literalai_info"]["metadata"],
         )
 
     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/langchain/utils.py

@@ -299,7 +299,7 @@ async def return_questions(query, response, chat_history_str, context, config):
     prompt = ChatPromptTemplate.from_messages(
         [
             ("system", system),
-            ("human", "{chat_history_str}, {context}, {query}, {response}"),
+            # ("human", "{chat_history_str}, {context}, {query}, {response}"),
         ]
     )
     llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

code/modules/chat_processor/helpers.py

@@ -0,0 +1,206 @@
+import os
+from literalai import AsyncLiteralClient
+from datetime import datetime, timedelta, timezone
+from modules.config.constants import COOLDOWN_TIME, TOKENS_LEFT
+from typing_extensions import TypedDict
+import tiktoken
+from typing import Any, Generic, List, Literal, Optional, TypeVar, Union
+
+Field = TypeVar("Field")
+Operators = TypeVar("Operators")
+Value = TypeVar("Value")
+
+BOOLEAN_OPERATORS = Literal["is", "nis"]
+STRING_OPERATORS = Literal["eq", "neq", "ilike", "nilike"]
+NUMBER_OPERATORS = Literal["eq", "neq", "gt", "gte", "lt", "lte"]
+STRING_LIST_OPERATORS = Literal["in", "nin"]
+DATETIME_OPERATORS = Literal["gte", "lte", "gt", "lt"]
+
+OPERATORS = Union[
+    BOOLEAN_OPERATORS,
+    STRING_OPERATORS,
+    NUMBER_OPERATORS,
+    STRING_LIST_OPERATORS,
+    DATETIME_OPERATORS,
+]
+
+
+class Filter(Generic[Field], TypedDict, total=False):
+    field: Field
+    operator: OPERATORS
+    value: Any
+    path: Optional[str]
+
+
+class OrderBy(Generic[Field], TypedDict):
+    column: Field
+    direction: Literal["ASC", "DESC"]
+
+
+threads_filterable_fields = Literal[
+    "id",
+    "createdAt",
+    "name",
+    "stepType",
+    "stepName",
+    "stepOutput",
+    "metadata",
+    "tokenCount",
+    "tags",
+    "participantId",
+    "participantIdentifiers",
+    "scoreValue",
+    "duration",
+]
+threads_orderable_fields = Literal["createdAt", "tokenCount"]
+threads_filters = List[Filter[threads_filterable_fields]]
+threads_order_by = OrderBy[threads_orderable_fields]
+
+steps_filterable_fields = Literal[
+    "id",
+    "name",
+    "input",
+    "output",
+    "participantIdentifier",
+    "startTime",
+    "endTime",
+    "metadata",
+    "parentId",
+    "threadId",
+    "error",
+    "tags",
+]
+steps_orderable_fields = Literal["createdAt"]
+steps_filters = List[Filter[steps_filterable_fields]]
+steps_order_by = OrderBy[steps_orderable_fields]
+
+users_filterable_fields = Literal[
+    "id",
+    "createdAt",
+    "identifier",
+    "lastEngaged",
+    "threadCount",
+    "tokenCount",
+    "metadata",
+]
+users_filters = List[Filter[users_filterable_fields]]
+
+scores_filterable_fields = Literal[
+    "id",
+    "createdAt",
+    "participant",
+    "name",
+    "tags",
+    "value",
+    "type",
+    "comment",
+]
+scores_orderable_fields = Literal["createdAt"]
+scores_filters = List[Filter[scores_filterable_fields]]
+scores_order_by = OrderBy[scores_orderable_fields]
+
+generation_filterable_fields = Literal[
+    "id",
+    "createdAt",
+    "model",
+    "duration",
+    "promptLineage",
+    "promptVersion",
+    "tags",
+    "score",
+    "participant",
+    "tokenCount",
+    "error",
+]
+generation_orderable_fields = Literal[
+    "createdAt",
+    "tokenCount",
+    "model",
+    "provider",
+    "participant",
+    "duration",
+]
+generations_filters = List[Filter[generation_filterable_fields]]
+generations_order_by = OrderBy[generation_orderable_fields]
+
+literal_client = AsyncLiteralClient(api_key=os.getenv("LITERAL_API_KEY_LOGGING"))
+
+
+# For consistency, use dictionary for user_info
+def convert_to_dict(user_info):
+    # if already a dictionary, return as is
+    if isinstance(user_info, dict):
+        return user_info
+    if hasattr(user_info, "__dict__"):
+        user_info = user_info.__dict__
+    return user_info
+
+
+def get_time():
+    return datetime.now(timezone.utc).isoformat()
+
+
+async def get_user_details(user_email_id):
+    user_info = await literal_client.api.get_user(identifier=user_email_id)
+    return user_info
+
+
+async def update_user_info(user_info):
+    # if object type, convert to dictionary
+    user_info = convert_to_dict(user_info)
+    await literal_client.api.update_user(
+        id=user_info["id"],
+        identifier=user_info["identifier"],
+        metadata=user_info["metadata"],
+    )
+
+
+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")
+
+    # Convert from ISO format string to datetime object and ensure UTC timezone
+    last_message_time = datetime.fromisoformat(last_message_time_str).replace(
+        tzinfo=timezone.utc
+    )
+    current_time = datetime.fromisoformat(current_time).replace(tzinfo=timezone.utc)
+
+    # Calculate the elapsed time
+    elapsed_time = current_time - last_message_time
+    elapsed_time_in_seconds = elapsed_time.total_seconds()
+
+    # Calculate when the cooldown period ends
+    cooldown_end_time = last_message_time + timedelta(seconds=COOLDOWN_TIME)
+    cooldown_end_time_iso = cooldown_end_time.isoformat()
+
+    # Debug: Print the cooldown end time
+    print(f"Cooldown end time (ISO): {cooldown_end_time_iso}")
+
+    # Check if the user is still in cooldown
+    if elapsed_time_in_seconds < COOLDOWN_TIME:
+        return True, cooldown_end_time_iso  # Return in ISO 8601 format
+
+    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)
+
+
+async def get_thread_step_info(thread_id):
+    step = await literal_client.api.get_step(thread_id)
+    return step
+
+
+def get_num_tokens(text, model):
+    encoding = tiktoken.encoding_for_model(model)
+    tokens = encoding.encode(text)
+    return len(tokens)

code/modules/config/constants.py

@@ -4,6 +4,11 @@ import os
 load_dotenv()
 
 TIMEOUT = 60
+COOLDOWN_TIME = 60
+TOKENS_LEFT = 3000
+
+GITHUB_REPO = "https://github.com/DL4DS/dl4ds_tutor"
+DOCS_WEBSITE = "https://dl4ds.github.io/dl4ds_tutor/"
 
 # API Keys - Loaded from the .env file
 

code/templates/cooldown.html

@@ -0,0 +1,155 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Cooldown Period | Terrier Tutor</title>
+    <style>
+        @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600&display=swap');
+
+        body, html {
+            margin: 0;
+            padding: 0;
+            font-family: 'Inter', sans-serif;
+            background-color: #f7f7f7;
+            background-image: url('https://www.transparenttextures.com/patterns/cubes.png');
+            background-repeat: repeat;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            height: 100vh;
+            color: #333;
+        }
+
+        .container {
+            background: rgba(255, 255, 255, 0.9);
+            border: 1px solid #ddd;
+            border-radius: 8px;
+            width: 100%;
+            max-width: 400px;
+            padding: 50px;
+            box-sizing: border-box;
+            text-align: center;
+            box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
+            backdrop-filter: blur(10px);
+            -webkit-backdrop-filter: blur(10px);
+        }
+
+        .avatar {
+            width: 90px;
+            height: 90px;
+            border-radius: 50%;
+            margin-bottom: 25px;
+            border: 2px solid #ddd;
+        }
+
+        .container h1 {
+            margin-bottom: 15px;
+            font-size: 24px;
+            font-weight: 600;
+            color: #1a1a1a;
+        }
+
+        .container p {
+            font-size: 16px;
+            color: #4a4a4a;
+            margin-bottom: 30px;
+            line-height: 1.5;
+        }
+
+        .cooldown-message {
+            font-size: 16px;
+            color: #333;
+            margin-bottom: 30px;
+        }
+
+        .button {
+            padding: 12px 0;
+            margin: 12px 0;
+            font-size: 14px;
+            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 */
+            transition: background-color 0.3s ease, border-color 0.3s ease;
+            display: none; /* Initially hidden */
+        }
+
+        .button.start-tutor {
+            display: none; /* Initially hidden */
+        }
+
+        .button:hover {
+            background-color: #e0e0e0;
+            border-color: #357ae8; /* Darker blue for hover */
+        }
+
+        .sign-out-button {
+            border: 1px solid #FF4C4C;
+            background-color: #fff;
+            color: #FF4C4C;
+            display: block; /* Ensure this button is always visible */
+        }
+
+        .sign-out-button:hover {
+            background-color: #ffe6e6; /* Light red on hover */
+            border-color: #e04343; /* Darker red for hover */
+            color: #e04343; /* Red text on hover */
+        }
+
+        #countdown {
+            font-size: 14px;
+            color: #555;
+            margin-bottom: 20px;
+        }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <img src="/public/avatars/ai_tutor.png" alt="AI Tutor Avatar" class="avatar">
+        <h1>Hello, {{ username }}</h1>
+        <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>
+        <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>
+    <script>
+        function startCountdown(endTime) {
+            const countdownElement = document.getElementById('countdown');
+            const startTutorBtn = document.getElementById('startTutorBtn');
+            const endTimeDate = new Date(endTime);  // Parse the cooldown end time
+
+            function updateCountdown() {
+                const now = new Date();  // Get the current time
+                const timeLeft = endTimeDate.getTime() - now.getTime();  // Calculate the remaining time in milliseconds
+
+                if (timeLeft <= 0) {
+                    countdownElement.textContent = "Cooldown period has ended.";
+                    startTutorBtn.style.display = "block";  // Show the start tutor button
+                } else {
+                    const hours = Math.floor(timeLeft / 1000 / 60 / 60);
+                    const minutes = Math.floor((timeLeft / 1000 / 60) % 60);
+                    const seconds = Math.floor((timeLeft / 1000) % 60);
+                    countdownElement.textContent = `${hours}h ${minutes}m ${seconds}s`;
+                }
+            }
+
+            updateCountdown();  // Initial call to set the countdown
+            setInterval(updateCountdown, 1000);  // Update the countdown every second
+        }
+
+        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
+        startCountdown("{{ cooldown_end_time }}");
+    </script>
+</body>
+</html>

code/templates/dashboard.html

@@ -3,7 +3,7 @@
 <head>
     <meta charset="UTF-8">
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>Dashboard</title>
+    <title>Dashboard | Terrier Tutor</title>
     <style>
         @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600&display=swap');
 
@@ -11,94 +11,99 @@
             margin: 0;
             padding: 0;
             font-family: 'Inter', sans-serif;
-            background: url('./public/space.jpg') no-repeat center center fixed;
-            background-size: cover;
-            background-color: #1a1a1a;
+            background-color: #f7f7f7; /* Light gray background */
+            background-image: url('https://www.transparenttextures.com/patterns/cubes.png'); /* Subtle geometric pattern */
+            background-repeat: repeat;
             display: flex;
             align-items: center;
             justify-content: center;
             height: 100vh;
-            color: #f5f5f5;
+            color: #333;
         }
 
         .container {
-            background: linear-gradient(145deg, rgba(255, 255, 255, 0.85), rgba(240, 240, 240, 0.85));
-            border: 1px solid rgba(255, 255, 255, 0.3);
-            border-radius: 12px;
-            box-shadow: 0 8px 32px rgba(0, 0, 0, 0.37);
-            backdrop-filter: blur(10px);
-            -webkit-backdrop-filter: blur(10px);
+            background: rgba(255, 255, 255, 0.9);
+            border: 1px solid #ddd;
+            border-radius: 8px;
             width: 100%;
             max-width: 400px;
-            padding: 40px;
+            padding: 50px;
             box-sizing: border-box;
             text-align: center;
-            position: relative;
-            overflow: hidden;
-        }
-
-        .container:before {
-            content: '';
-            position: absolute;
-            top: 0;
-            left: 0;
-            width: 100%;
-            height: 100%;
-            background: radial-gradient(circle, rgba(255, 255, 255, 0.3) 0%, rgba(255, 255, 255, 0) 70%);
-            z-index: 1;
-            pointer-events: none;
-        }
-
-        .container > * {
-            position: relative;
-            z-index: 2;
+            box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
+            backdrop-filter: blur(10px);
+            -webkit-backdrop-filter: blur(10px);
         }
 
         .avatar {
-            width: 80px;
-            height: 80px;
+            width: 90px;
+            height: 90px;
             border-radius: 50%;
-            margin-bottom: 20px;
-            border: 2px solid #fff;
+            margin-bottom: 25px;
+            border: 2px solid #ddd;
         }
 
         .container h1 {
-            margin-bottom: 20px;
-            font-size: 26px;
+            margin-bottom: 15px;
+            font-size: 24px;
             font-weight: 600;
-            color: #333;
+            color: #1a1a1a;
         }
 
         .container p {
-            font-size: 15px;
-            color: #666;
+            font-size: 16px;
+            color: #4a4a4a;
             margin-bottom: 30px;
+            line-height: 1.5;
+        }
+
+        .tokens-left {
+            font-size: 16px;
+            color: #333;
+            margin-bottom: 30px;
+            font-weight: 600;
         }
 
         .button {
             padding: 12px 0;
-            margin: 10px 0;
-            font-size: 16px;
-            border: none;
-            border-radius: 5px;
+            margin: 12px 0;
+            font-size: 14px;
+            border-radius: 6px;
             cursor: pointer;
             width: 100%;
-            transition: background-color 0.3s ease, color 0.3s ease;
-            background-color: #FAFAFA;
-            color: #333;
+            border: 1px solid #4285F4; /* Button border color */
+            background-color: #fff; /* Button background color */
+            color: #4285F4; /* Button text color */
+            transition: background-color 0.3s ease, border-color 0.3s ease;
         }
 
         .button:hover {
-            background-color: #f0f0f0;
+            background-color: #e0e0e0;
+            border-color: #357ae8; /* Darker blue for hover */
         }
 
         .start-button {
-            background-color: #4CAF50;
-            color: #fff;
+            border: 1px solid #4285F4;
+            color: #4285F4;
+            background-color: #fff;
         }
 
         .start-button:hover {
-            background-color: #45a049;
+            background-color: #e0f0ff; /* Light blue on hover */
+            border-color: #357ae8; /* Darker blue for hover */
+            color: #357ae8; /* Blue text on hover */
+        }
+
+        .sign-out-button {
+            border: 1px solid #FF4C4C;
+            background-color: #fff;
+            color: #FF4C4C;
+        }
+
+        .sign-out-button:hover {
+            background-color: #ffe6e6; /* Light red on hover */
+            border-color: #e04343; /* Darker red for hover */
+            color: #e04343; /* Red text on hover */
         }
     </style>
 </head>
@@ -107,9 +112,13 @@
         <img src="/public/avatars/ai_tutor.png" alt="AI Tutor Avatar" class="avatar">
         <h1>Welcome, {{ username }}</h1>
         <p>Ready to start your AI tutoring session?</p>
+        <p class="tokens-left">Tokens Left: {{ tokens_left }}</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>
     <script>
         let token = "{{ jwt_token }}";

code/templates/error.html

@@ -0,0 +1,95 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Error | Terrier Tutor</title>
+    <style>
+        @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600&display=swap');
+
+        body, html {
+            margin: 0;
+            padding: 0;
+            font-family: 'Inter', sans-serif;
+            background-color: #f7f7f7; /* Light gray background */
+            background-image: url('https://www.transparenttextures.com/patterns/cubes.png'); /* Subtle geometric pattern */
+            background-repeat: repeat;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            height: 100vh;
+            color: #333;
+        }
+
+        .container {
+            background: rgba(255, 255, 255, 0.9);
+            border: 1px solid #ddd;
+            border-radius: 8px;
+            width: 100%;
+            max-width: 400px;
+            padding: 50px;
+            box-sizing: border-box;
+            text-align: center;
+            box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
+            backdrop-filter: blur(10px);
+            -webkit-backdrop-filter: blur(10px);
+        }
+
+        .container h1 {
+            margin-bottom: 20px;
+            font-size: 26px;
+            font-weight: 600;
+            color: #1a1a1a;
+        }
+
+        .container p {
+            font-size: 18px;
+            color: #4a4a4a;
+            margin-bottom: 35px;
+            line-height: 1.5;
+        }
+
+        .button {
+            padding: 14px 0;
+            margin: 12px 0;
+            font-size: 16px;
+            border-radius: 6px;
+            cursor: pointer;
+            width: 100%;
+            border: 1px solid #ccc;
+            background-color: #007BFF;
+            color: #fff;
+            transition: background-color 0.3s ease, border-color 0.3s ease;
+        }
+
+        .button:hover {
+            background-color: #0056b3;
+            border-color: #0056b3;
+        }
+
+        .error-box {
+            background-color: #2d2d2d;
+            color: #fff;
+            padding: 10px;
+            margin-top: 20px;
+            font-family: 'Courier New', Courier, monospace;
+            text-align: left;
+            overflow-x: auto;
+            white-space: pre-wrap;
+            border-radius: 5px;
+        }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <h1>Oops! Something went wrong...</h1>
+        <p>An unexpected error occurred. The details are below:</p>
+        <div class="error-box">
+            <code>{{ error }}</code>
+        </div>
+        <form action="/" method="get">
+            <button type="submit" class="button">Return to Home</button>
+        </form>
+    </div>
+</body>
+</html>

code/templates/error_404.html

@@ -0,0 +1,80 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>404 - Not Found</title>
+    <style>
+        @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600&display=swap');
+
+        body, html {
+            margin: 0;
+            padding: 0;
+            font-family: 'Inter', sans-serif;
+            background-color: #f7f7f7; /* Light gray background */
+            background-image: url('https://www.transparenttextures.com/patterns/cubes.png'); /* Subtle geometric pattern */
+            background-repeat: repeat;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            height: 100vh;
+            color: #333;
+        }
+
+        .container {
+            background: rgba(255, 255, 255, 0.9);
+            border: 1px solid #ddd;
+            border-radius: 8px;
+            width: 100%;
+            max-width: 400px;
+            padding: 50px;
+            box-sizing: border-box;
+            text-align: center;
+            box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
+            backdrop-filter: blur(10px);
+            -webkit-backdrop-filter: blur(10px);
+        }
+
+        .container h1 {
+            margin-bottom: 20px;
+            font-size: 26px;
+            font-weight: 600;
+            color: #1a1a1a;
+        }
+
+        .container p {
+            font-size: 18px;
+            color: #4a4a4a;
+            margin-bottom: 35px;
+            line-height: 1.5;
+        }
+
+        .button {
+            padding: 14px 0;
+            margin: 12px 0;
+            font-size: 16px;
+            border-radius: 6px;
+            cursor: pointer;
+            width: 100%;
+            border: 1px solid #ccc;
+            background-color: #007BFF;
+            color: #fff;
+            transition: background-color 0.3s ease, border-color 0.3s ease;
+        }
+
+        .button:hover {
+            background-color: #0056b3;
+            border-color: #0056b3;
+        }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <h1>You have ventured into the abyss...</h1>
+        <p>To get back to reality, click the button below.</p>
+        <form action="/" method="get">
+            <button type="submit" class="button">Return to Home</button>
+        </form>
+    </div>
+</body>
+</html>

code/templates/login.html

@@ -3,7 +3,7 @@
 <head>
     <meta charset="UTF-8">
     <meta name="viewport" content="width=device-width, initial-scale=1.0">
-    <title>Login</title>
+    <title>Login | Terrier Tutor</title>
     <style>
         @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600&display=swap');
 
@@ -11,95 +11,97 @@
             margin: 0;
             padding: 0;
             font-family: 'Inter', sans-serif;
-            background: url('/public/space.jpg') no-repeat center center fixed;
-            background-size: cover;
-            background-color: #1a1a1a;
+            background-color: #f7f7f7; /* Light gray background */
+            background-image: url('https://www.transparenttextures.com/patterns/cubes.png'); /* Subtle geometric pattern */
+            background-repeat: repeat;
             display: flex;
             align-items: center;
             justify-content: center;
             height: 100vh;
-            color: #f5f5f5;
+            color: #333;
         }
 
         .container {
-            background: linear-gradient(145deg, rgba(255, 255, 255, 0.85), rgba(240, 240, 240, 0.85));
-            border: 1px solid rgba(255, 255, 255, 0.3);
-            border-radius: 12px;
-            box-shadow: 0 8px 32px rgba(0, 0, 0, 0.37);
-            backdrop-filter: blur(10px);
-            -webkit-backdrop-filter: blur(10px);
+            background: rgba(255, 255, 255, 0.9);
+            border: 1px solid #ddd;
+            border-radius: 8px;
             width: 100%;
             max-width: 400px;
-            padding: 40px;
+            padding: 50px;
             box-sizing: border-box;
             text-align: center;
-            position: relative;
-            overflow: hidden;
-        }
-
-        .container:before {
-            content: '';
-            position: absolute;
-            top: 0;
-            left: 0;
-            width: 100%;
-            height: 100%;
-            background: radial-gradient(circle, rgba(255, 255, 255, 0.3) 0%, rgba(255, 255, 255, 0) 70%);
-            z-index: 1;
-            pointer-events: none;
-        }
-
-        .container > * {
-            position: relative;
-            z-index: 2;
+            box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
+            backdrop-filter: blur(10px);
+            -webkit-backdrop-filter: blur(10px);
         }
 
         .avatar {
-            width: 80px;
-            height: 80px;
+            width: 90px;
+            height: 90px;
             border-radius: 50%;
-            margin-bottom: 20px;
-            border: 2px solid #fff;
+            margin-bottom: 25px;
+            border: 2px solid #ddd;
         }
 
         .container h1 {
-            margin-bottom: 20px;
-            font-size: 26px;
+            margin-bottom: 15px;
+            font-size: 24px;
             font-weight: 600;
-            color: #333;
+            color: #1a1a1a;
         }
 
         .container p {
-            font-size: 15px;
-            color: #666;
+            font-size: 16px;
+            color: #4a4a4a;
             margin-bottom: 30px;
+            line-height: 1.5;
         }
 
         .button {
             padding: 12px 0;
-            margin: 10px 0;
-            font-size: 16px;
-            border: none;
-            border-radius: 5px;
+            margin: 12px 0;
+            font-size: 14px;
+            border-radius: 6px;
             cursor: pointer;
             width: 100%;
-            transition: background-color 0.3s ease, color 0.3s ease;
-            background-color: #FAFAFA;
-            color: #333;
+            border: 1px solid #4285F4; /* Google button border color */
+            background-color: #fff; /* Guest button color */
+            color: #4285F4; /* Google button text color */
+            transition: background-color 0.3s ease, border-color 0.3s ease;
         }
 
         .button:hover {
-            background-color: #f0f0f0;
+            background-color: #e0f0ff; /* Light blue on hover */
+            border-color: #357ae8; /* Darker blue for hover */
+            color: #357ae8; /* Blue text on hover */
+        }
+
+        .footer {
+            margin-top: 40px;
+            font-size: 15px;
+            color: #666;
+            text-align: center; /* Center the text in the footer */
+        }
+
+        .footer a {
+            color: #333;
+            text-decoration: none;
+            font-weight: 500;
+            display: inline-flex;
+            align-items: center;
+            justify-content: center; /* Center the content of the links */
+            transition: color 0.3s ease;
+            margin-bottom: 8px;
+            width: 100%; /* Make the link block level */
         }
 
-        .google-button {
-            background-color: #4285F4;
-            color: #fff;
-            border: none;
+        .footer a:hover {
+            color: #000;
         }
 
-        .google-button:hover {
-            background-color: #357ae8;
+        .footer svg {
+            margin-right: 8px;
+            fill: currentColor;
         }
     </style>
 </head>
@@ -107,13 +109,24 @@
     <div class="container">
         <img src="/public/avatars/ai_tutor.png" alt="AI Tutor Avatar" class="avatar">
         <h1>Terrier Tutor</h1>
-        <p>Hey! Welcome to the DS598 AI tutor</p>
-        <form action="/login/guest" method="post">
-            <button type="submit" class="button">Sign in as Guest</button>
-        </form>
+        <p>Welcome to the DS598 AI Tutor. Please sign in to continue.</p>
         <form action="/login/google" method="get">
-            <button type="submit" class="button google-button">Sign in with Google</button>
+            <button type="submit" class="button">Sign in with Google</button>
         </form>
+        <div class="footer">
+            <a href="{{ GITHUB_REPO }}" target="_blank">
+                <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24">
+                    <path d="M12 .5C5.596.5.5 5.596.5 12c0 5.098 3.292 9.414 7.852 10.94.574.105.775-.249.775-.553 0-.272-.01-1.008-.015-1.98-3.194.694-3.87-1.544-3.87-1.544-.521-1.324-1.273-1.676-1.273-1.676-1.04-.714.079-.7.079-.7 1.148.08 1.75 1.181 1.75 1.181 1.022 1.752 2.683 1.246 3.34.954.104-.74.4-1.246.73-1.533-2.551-.292-5.234-1.276-5.234-5.675 0-1.253.447-2.277 1.181-3.079-.12-.293-.51-1.47.113-3.063 0 0 .96-.307 3.15 1.174.913-.255 1.892-.383 2.867-.388.975.005 1.954.133 2.868.388 2.188-1.481 3.147-1.174 3.147-1.174.624 1.593.233 2.77.114 3.063.735.802 1.18 1.826 1.18 3.079 0 4.407-2.688 5.38-5.248 5.668.413.354.782 1.049.782 2.113 0 1.526-.014 2.757-.014 3.132 0 .307.198.662.783.553C20.21 21.411 23.5 17.096 23.5 12c0-6.404-5.096-11.5-11.5-11.5z"/>
+                </svg>
+                View on GitHub
+            </a>
+            <a href="{{ DOCS_WEBSITE }}" target="_blank">
+                <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24">
+                    <path d="M19 2H8c-1.103 0-2 .897-2 2v16c0 1.103.897 2 2 2h12c1.103 0 2-.897 2-2V7l-5-5zm0 2l.001 4H14V4h5zm-1 14H9V4h4v6h6v8zM7 4H6v16c0 1.654 1.346 3 3 3h9v-2H9c-.551 0-1-.449-1-1V4z"/>
+                </svg>
+                View Docs
+            </a>
+        </div>
     </div>
 </body>
 </html>