implement signin signout

llm_demo/app.py

@@ -14,7 +14,7 @@
 
 import os
 from contextlib import asynccontextmanager
-from typing import Optional
+from typing import Any, Optional
 
 import uvicorn
 from fastapi import APIRouter, Body, FastAPI, HTTPException, Request
@@ -48,14 +48,24 @@ async def index(request: Request):
     # User session setup
     orchestrator = request.app.state.orchestrator
     session = request.session
+
     if "uuid" not in session or not orchestrator.user_session_exist(session["uuid"]):
         await orchestrator.user_session_create(session)
+
+    # recheck if token is still valid
+    user_id_token = orchestrator.get_user_id_token(session["uuid"])
+    if user_id_token:
+        check_user_info_validity(session, user_id_token, request.app.state.client_id)
+
     return templates.TemplateResponse(
         "index.html",
         {
             "request": request,
             "messages": request.session["history"],
             "client_id": request.app.state.client_id,
+            "user_img": request.session["user_info"]["user_img"]
+            if "user_info" in request.session
+            else None,
         },
     )
 
@@ -72,14 +82,17 @@ async def login_google(
     client_id = request.app.state.client_id
     if not client_id:
         raise HTTPException(status_code=400, detail="Client id not found")
-    user_name = get_user_name(str(user_id_token), client_id)
+
+    session = request.session
+    user_info = get_user_info(str(user_id_token), client_id)
+    session["user_info"] = user_info
 
     # create new request session
     orchestrator = request.app.state.orchestrator
     orchestrator.set_user_session_header(request.session["uuid"], str(user_id_token))
     print("Logged in to Google.")
 
-    welcome_text = f"Welcome to Cymbal Air, {user_name}! How may I assist you?"
+    welcome_text = f"Welcome to Cymbal Air, {request.session['user_info']['name']}! How may I assist you?"
     if len(request.session["history"]) == 1:
         request.session["history"][0] = {
             "type": "ai",
@@ -131,11 +144,22 @@ async def reset(request: Request):
     request.session.clear()
 
 
-def get_user_name(user_token_id: str, client_id: str) -> str:
-    id_info = id_token.verify_oauth2_token(
-        user_token_id, requests.Request(), audience=client_id
-    )
-    return id_info["name"]
+def get_user_info(user_id_token: str, client_id: str) -> dict[str, str]:
+    try:
+        id_info = id_token.verify_oauth2_token(
+            user_id_token, requests.Request(), audience=client_id
+        )
+        return {
+            "user_img": id_info["picture"],
+            "name": id_info["name"],
+        }
+    except ValueError as err:
+        return {}
+
+
+def check_user_info_validity(session: dict[str, Any], user_id_token: str, client_id: str):
+    if not get_user_info(user_id_token, client_id):
+        del session["user_info"]
 
 
 def init_app(

llm_demo/orchestrator/orchestrator.py

@@ -14,7 +14,7 @@
 
 import asyncio
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, Optional
 
 
 class classproperty:
@@ -61,6 +61,15 @@ def set_user_session_header(self, uuid: str, user_id_token: str):
         user_session = self.get_user_session(uuid)
         user_session.client.headers["User-Id-Token"] = f"Bearer {user_id_token}"
 
+    def get_user_id_token(self, uuid: str) -> Optional[str]:
+        user_session = self.get_user_session(uuid)
+        if user_session.client and "User-Id-Token" in user_session.client.headers:
+            token = user_session.client.headers["User-Id-Token"]
+            parts = str(token).split(" ")
+            if len(parts) == 2:
+                return parts[1]
+        return None
+
 
 def createOrchestrator(orchestration_type: str) -> "BaseOrchestrator":
     for cls in BaseOrchestrator.__subclasses__():

llm_demo/static/index.css

@@ -51,8 +51,23 @@ body {
 #g_id_onload,
 .g_id_signin {
     position: absolute;
-    top: 8px;
+    top: 6px;
     right: 10px;
+    .chat-user-state {
+        top: 6px;
+        position: relative;
+        .chat-user-image {
+            border-radius: 50%;
+            height: 36px;
+            width: 36px;
+        }
+        .chat-signout-btn {
+            border-radius: 5px;
+            margin-left: 10px;
+            height: 32px;
+            font-family: "Roboto";
+        }
+    }
 }
 
 #menuButton {
@@ -191,4 +206,4 @@ div.chat-wrapper div.chat-content div .sender-icon img {
 .send-button {
     margin-top: 12px;
     margin-right: 12px;
-}
+}

llm_demo/static/index.js

@@ -31,6 +31,11 @@ $('#resetButton').click(async (e) => {
     await reset();
 });
 
+// Sign out current user via click
+$('#signoutButton').click(async (e) => {
+    await reset();
+});
+
 async function submitMessage() {
     let msg = $('.chat-bar input').val();
     // Add message to UI
@@ -68,6 +73,7 @@ async function askQuestion(prompt) {
 }
 
 async function reset() {
+    print("here")
     await fetch('reset', {
         method: 'POST',
     }).then(()=>{

llm_demo/templates/index.html

@@ -33,26 +33,32 @@
 </head>
 
 <body>
-
     <div class="container">
         <div class="chat-header">
             <span class="material-symbols-outlined" id="menuButton">menu</span>
             <span>Cymbal Air Customer Service Assistant</span>
             <span class="material-symbols-outlined" id="resetButton">refresh</span>
         </div>
-        <div id="g_id_onload"
+        <div class="chat-user-wrapper"
+            id="g_id_onload"
             data-context="signin"
             data-ux_mode="popup"
             data-auto_prompt="false">
-        </div>
-        <div class="g_id_signin"
-            data-type="standard"
-            data-shape="rectangular"
-            data-theme="outline"
-            data-text="signin_with"
-            data-size="large"
-            data-logo_alignment="left"
-            data-onsuccess="onSignIn">
+        {% if user_img %}
+            <div class="chat-user-state">
+                <img class="chat-user-image" src="{{ user_img }}" alt="signed in user image">
+                <button type="button" class="btn btn-default chat-signout-btn" id="signoutButton" onclick="reset()">Sign out</button>
+            </div>
+        {% else %}
+            <div class="g_id_signin"
+                data-type="standard"
+                data-shape="rectangular"
+                data-theme="outline"
+                data-text="signin"
+                data-size="medium"
+                data-logo_alignment="left">
+            </div>
+        {% endif %}
         </div>
         <div class="chat-wrapper">
             <div class="chat-content">