Merge pull request #63 from DL4DS/fastapi

Dockerfile

@@ -35,4 +35,4 @@ RUN --mount=type=secret,id=LITERAL_API_KEY_LOGGING,mode=0444,required=true
 RUN --mount=type=secret,id=CHAINLIT_AUTH_SECRET,mode=0444,required=true 
 
 # Default command to run the application
-CMD ["sh", "-c", "python -m modules.vectorstore.store_manager && chainlit run main.py --host 0.0.0.0 --port 7860"]
+CMD ["sh", "-c", "python -m modules.vectorstore.store_manager && uvicorn app:app --host 0.0.0.0 --port 7860"]

code/app.py

@@ -0,0 +1,236 @@
+from fastapi import FastAPI, Request, Response
+from fastapi.responses import HTMLResponse, RedirectResponse
+from fastapi.templating import Jinja2Templates
+from google.oauth2 import id_token
+from google.auth.transport import requests as google_requests
+from google_auth_oauthlib.flow import Flow
+from chainlit.utils import mount_chainlit
+import secrets
+import json
+import base64
+from modules.config.constants import (
+    OAUTH_GOOGLE_CLIENT_ID,
+    OAUTH_GOOGLE_CLIENT_SECRET,
+    CHAINLIT_URL,
+)
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+
+GOOGLE_CLIENT_ID = OAUTH_GOOGLE_CLIENT_ID
+GOOGLE_CLIENT_SECRET = OAUTH_GOOGLE_CLIENT_SECRET
+GOOGLE_REDIRECT_URI = f"{CHAINLIT_URL}/auth/oauth/google/callback"
+
+app = FastAPI()
+app.mount("/public", StaticFiles(directory="public"), name="public")
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Update with appropriate origins
+    allow_methods=["*"],
+    allow_headers=["*"],  # or specify the headers you want to allow
+    expose_headers=["X-User-Info"],  # Expose the custom header
+)
+
+templates = Jinja2Templates(directory="templates")
+session_store = {}
+CHAINLIT_PATH = "/chainlit_tutor"
+
+USER_ROLES = {
+    "[email protected]": ["instructor", "bu"],
+    "[email protected]": ["instructor", "bu"],
+    "[email protected]": ["instructor", "bu"],
+    "[email protected]": ["guest"],
+    # Add more users and roles as needed
+}
+
+# Create a Google OAuth flow
+flow = Flow.from_client_config(
+    {
+        "web": {
+            "client_id": GOOGLE_CLIENT_ID,
+            "client_secret": GOOGLE_CLIENT_SECRET,
+            "auth_uri": "https://accounts.google.com/o/oauth2/auth",
+            "token_uri": "https://oauth2.googleapis.com/token",
+            "redirect_uris": [GOOGLE_REDIRECT_URI],
+            "scopes": [
+                "openid",
+                # "https://www.googleapis.com/auth/userinfo.email",
+                # "https://www.googleapis.com/auth/userinfo.profile",
+            ],
+        }
+    },
+    scopes=[
+        "openid",
+        "https://www.googleapis.com/auth/userinfo.email",
+        "https://www.googleapis.com/auth/userinfo.profile",
+    ],
+    redirect_uri=GOOGLE_REDIRECT_URI,
+)
+
+
+def get_user_role(username: str):
+    return USER_ROLES.get(username, ["student"])  # Default to "student" role
+
+
+def get_user_info_from_cookie(request: Request):
+    user_info_encoded = request.cookies.get("X-User-Info")
+    if user_info_encoded:
+        try:
+            user_info_json = base64.b64decode(user_info_encoded).decode()
+            return json.loads(user_info_json)
+        except Exception as e:
+            print(f"Error decoding user info: {e}")
+            return None
+    return None
+
+
+def get_user_info(request: Request):
+    session_token = request.cookies.get("session_token")
+    if session_token and session_token in session_store:
+        return session_store[session_token]
+    return None
+
+
+@app.get("/", response_class=HTMLResponse)
+async def login_page(request: Request):
+    user_info = 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
+
+
+@app.get("/login/google")
+async def login_google(request: Request):
+    # Clear any existing session cookies to avoid conflicts with guest sessions
+    response = RedirectResponse(url="/post-signin")
+    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}")
+    # Check if user is already signed in using Google
+    if user_info and user_info.get("google_signed_in"):
+        return RedirectResponse("/post-signin")
+    else:
+        authorization_url, _ = flow.authorization_url(prompt="consent")
+        return RedirectResponse(authorization_url, headers=response.headers)
+
+
+@app.get("/auth/oauth/google/callback")
+async def auth_google(request: Request):
+    try:
+        flow.fetch_token(code=request.query_params.get("code"))
+        credentials = flow.credentials
+        user_info = id_token.verify_oauth2_token(
+            credentials.id_token, google_requests.Request(), GOOGLE_CLIENT_ID
+        )
+
+        email = user_info["email"]
+        name = user_info.get("name", "")
+        profile_image = user_info.get("picture", "")
+
+        session_token = secrets.token_hex(16)
+        session_store[session_token] = {
+            "email": email,
+            "name": name,
+            "profile_image": profile_image,
+            "google_signed_in": True,  # Set this flag to True for Google-signed users
+        }
+
+        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
+    except Exception as e:
+        print(f"Error during Google OAuth callback: {e}")
+        return RedirectResponse(url="/", status_code=302)
+
+
+@app.get("/post-signin", response_class=HTMLResponse)
+async def post_signin(request: Request):
+    user_info = 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"):
+    if user_info:
+        username = user_info["email"]
+        role = get_user_role(username)
+        jwt_token = request.cookies.get("X-User-Info")
+        return templates.TemplateResponse(
+            "dashboard.html",
+            {
+                "request": request,
+                "username": username,
+                "role": role,
+                "jwt_token": jwt_token,
+            },
+        )
+    return RedirectResponse("/")
+
+
+@app.post("/start-tutor")
+async def start_tutor(request: Request):
+    user_info = 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()
+
+        response = RedirectResponse(CHAINLIT_PATH, status_code=303)
+        response.set_cookie(key="X-User-Info", value=user_info_encoded, httponly=True)
+        return response
+
+    return RedirectResponse(url="/")
+
+
+@app.exception_handler(Exception)
+async def exception_handler(request: Request, exc: Exception):
+    return templates.TemplateResponse(
+        "error.html", {"request": request, "error": str(exc)}, status_code=500
+    )
+
+
+@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)
+
+
+mount_chainlit(app=app, target="main.py", path=CHAINLIT_PATH)
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="127.0.0.1", port=7860)

code/main.py

@@ -5,7 +5,6 @@ from modules.config.constants import (
     LITERAL_API_URL,
 )
 from modules.chat_processor.literal_ai import CustomLiteralDataLayer
-
 import json
 import yaml
 from typing import Any, Dict, no_type_check
@@ -21,6 +20,7 @@ import copy
 from typing import Optional
 from chainlit.types import ThreadDict
 import time
+import base64
 
 USER_TIMEOUT = 60_000
 SYSTEM = "System"
@@ -445,14 +445,41 @@ 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.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]:
+
+        print("\n\n\nI am here\n\n\n")
+        # try: # TODO: Add try-except block after testing
+        # TODO: Implement to get the user information from the headers (not the cookie)
+        cookie = headers.get("cookie")  # gets back a str
+        # Create a dictionary from the pairs
+        cookie_dict = {}
+        for pair in cookie.split("; "):
+            key, value = pair.split("=", 1)
+            # Strip surrounding quotes if present
+            cookie_dict[key] = value.strip('"')
+
+        decoded_user_info = base64.b64decode(
+            cookie_dict.get("X-User-Info", "")
+        ).decode()
+        decoded_user_info = json.loads(decoded_user_info)
+
+        return cl.User(
+            identifier=decoded_user_info["email"],
+            metadata={
+                "name": decoded_user_info["name"],
+                "avatar": decoded_user_info["profile_image"],
+            },
+        )
 
     async def on_follow_up(self, action: cl.Action):
         message = await cl.Message(
@@ -482,4 +509,8 @@ async def start_app():
     cl.action_callback("follow up question")(chatbot.on_follow_up)
 
 
-asyncio.run(start_app())
+loop = asyncio.get_event_loop()
+if loop.is_running():
+    asyncio.ensure_future(start_app())
+else:
+    asyncio.run(start_app())

code/modules/config/constants.py

@@ -12,6 +12,7 @@ LLAMA_CLOUD_API_KEY = os.getenv("LLAMA_CLOUD_API_KEY")
 HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 LITERAL_API_KEY_LOGGING = os.getenv("LITERAL_API_KEY_LOGGING")
 LITERAL_API_URL = os.getenv("LITERAL_API_URL")
+CHAINLIT_URL = os.getenv("CHAINLIT_URL")
 
 OAUTH_GOOGLE_CLIENT_ID = os.getenv("OAUTH_GOOGLE_CLIENT_ID")
 OAUTH_GOOGLE_CLIENT_SECRET = os.getenv("OAUTH_GOOGLE_CLIENT_SECRET")

code/templates/dashboard.html

@@ -0,0 +1,120 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Dashboard</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: url('./public/space.jpg') no-repeat center center fixed;
+            background-size: cover;
+            background-color: #1a1a1a;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            height: 100vh;
+            color: #f5f5f5;
+        }
+
+        .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);
+            width: 100%;
+            max-width: 400px;
+            padding: 40px;
+            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;
+        }
+
+        .avatar {
+            width: 80px;
+            height: 80px;
+            border-radius: 50%;
+            margin-bottom: 20px;
+            border: 2px solid #fff;
+        }
+
+        .container h1 {
+            margin-bottom: 20px;
+            font-size: 26px;
+            font-weight: 600;
+            color: #333;
+        }
+
+        .container p {
+            font-size: 15px;
+            color: #666;
+            margin-bottom: 30px;
+        }
+
+        .button {
+            padding: 12px 0;
+            margin: 10px 0;
+            font-size: 16px;
+            border: none;
+            border-radius: 5px;
+            cursor: pointer;
+            width: 100%;
+            transition: background-color 0.3s ease, color 0.3s ease;
+            background-color: #FAFAFA;
+            color: #333;
+        }
+
+        .button:hover {
+            background-color: #f0f0f0;
+        }
+
+        .start-button {
+            background-color: #4CAF50;
+            color: #fff;
+        }
+
+        .start-button:hover {
+            background-color: #45a049;
+        }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <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>
+        <form action="/start-tutor" method="post">
+            <button type="submit" class="button start-button">Start AI Tutor</button>
+        </form>
+    </div>
+    <script>
+        let token = "{{ jwt_token }}";
+        console.log("Token: ", token);
+        localStorage.setItem('token', token);
+    </script>
+</body>
+</html>

code/templates/login.html

@@ -0,0 +1,119 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Login</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: url('/public/space.jpg') no-repeat center center fixed;
+            background-size: cover;
+            background-color: #1a1a1a;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            height: 100vh;
+            color: #f5f5f5;
+        }
+
+        .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);
+            width: 100%;
+            max-width: 400px;
+            padding: 40px;
+            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;
+        }
+
+        .avatar {
+            width: 80px;
+            height: 80px;
+            border-radius: 50%;
+            margin-bottom: 20px;
+            border: 2px solid #fff;
+        }
+
+        .container h1 {
+            margin-bottom: 20px;
+            font-size: 26px;
+            font-weight: 600;
+            color: #333;
+        }
+
+        .container p {
+            font-size: 15px;
+            color: #666;
+            margin-bottom: 30px;
+        }
+
+        .button {
+            padding: 12px 0;
+            margin: 10px 0;
+            font-size: 16px;
+            border: none;
+            border-radius: 5px;
+            cursor: pointer;
+            width: 100%;
+            transition: background-color 0.3s ease, color 0.3s ease;
+            background-color: #FAFAFA;
+            color: #333;
+        }
+
+        .button:hover {
+            background-color: #f0f0f0;
+        }
+
+        .google-button {
+            background-color: #4285F4;
+            color: #fff;
+            border: none;
+        }
+
+        .google-button:hover {
+            background-color: #357ae8;
+        }
+    </style>
+</head>
+<body>
+    <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>
+        <form action="/login/google" method="get">
+            <button type="submit" class="button google-button">Sign in with Google</button>
+        </form>
+    </div>
+</body>
+</html>

code/templates/logout.html

@@ -0,0 +1,21 @@
+<!DOCTYPE html>
+<html>
+<head>
+    <title>Logout</title>
+    <script>
+        window.onload = function() {
+            fetch('/chainlit_tutor/logout', {
+                method: 'POST',
+                credentials: 'include'  // Ensure cookies are sent
+            }).then(() => {
+                window.location.href = '/';
+            }).catch(error => {
+                console.error('Logout failed:', error);
+            });
+        };
+    </script>
+</head>
+<body>
+    <p>Logging out... If you are not redirected, <a href="/">click here</a>.</p>
+</body>
+</html>