added auth

app.py

@@ -7,13 +7,15 @@ import re
 from typing import List, Literal, TypedDict, Tuple
 from transformers import AutoTokenizer
 from gradio_client import Client
-import constants as C
-import utils as U
-
 from openai import OpenAI
 import anthropic
 from groq import Groq
 
+import constants as C
+import utils as U
+from auth import authenticateFunc
+from sidebar import showSidebar
+
 from dotenv import load_dotenv
 load_dotenv()
 
@@ -226,7 +228,7 @@ def __logLlmRequest(messagesFormatted: list):
     # U.pprint(f"{messagesFormatted=}")
 
 
-def predict():
+def __predict():
     messagesFormatted = []
 
     try:
@@ -352,112 +354,118 @@ U.pprint("\n")
 U.applyCommonStyles()
 st.title("Kommuneity Story Creator 🪄")
 
-if "startMsg" not in st.session_state:
-    __setStartMsg("")
-    st.button(C.START_MSG, on_click=lambda: __setStartMsg(C.START_MSG))
-
-for chat in st.session_state.chatHistory:
-    role = chat["role"]
-    content = chat["content"]
-    imagePath = chat.get("image")
-    avatar = C.AI_ICON if role == "assistant" else C.USER_ICON
-    with st.chat_message(role, avatar=avatar):
-        st.markdown(content)
-        if imagePath:
-            st.image(imagePath)
-
-# U.pprint(f"{st.session_state.buttonValue=}")
-# U.pprint(f"{st.session_state.selectedStory=}")
-# U.pprint(f"{st.session_state.startMsg=}")
-
-if prompt := (
-    st.chat_input()
-    or st.session_state["buttonValue"]
-    or st.session_state["selectedStoryTitle"]
-    or st.session_state["startMsg"]
-):
-    __resetButtonState()
-    __setStartMsg("")
-    if st.session_state["selectedStoryTitle"] != prompt:
-        __resetSelectedStory()
-    st.session_state.selectedStoryTitle = ""
 
-    with st.chat_message("user", avatar=C.USER_ICON):
-        st.markdown(prompt)
-    U.pprint(f"{prompt=}")
-    st.session_state.chatHistory.append({"role": "user", "content": prompt })
-    st.session_state.messages.append({"role": "user", "content": prompt})
+def mainApp():
+    if "startMsg" not in st.session_state:
+        __setStartMsg("")
+        st.button(C.START_MSG, on_click=lambda: __setStartMsg(C.START_MSG))
+
+    for chat in st.session_state.chatHistory:
+        role = chat["role"]
+        content = chat["content"]
+        imagePath = chat.get("image")
+        avatar = C.AI_ICON if role == "assistant" else C.USER_ICON
+        with st.chat_message(role, avatar=avatar):
+            st.markdown(content)
+            if imagePath:
+                st.image(imagePath)
+
+    # U.pprint(f"{st.session_state.buttonValue=}")
+    # U.pprint(f"{st.session_state.selectedStoryTitle=}")
+    # U.pprint(f"{st.session_state.startMsg=}")
+
+    if prompt := (
+        st.chat_input()
+        or st.session_state["buttonValue"]
+        or st.session_state["selectedStoryTitle"]
+        or st.session_state["startMsg"]
+    ):
+        __resetButtonState()
+        __setStartMsg("")
+        if st.session_state["selectedStoryTitle"] != prompt:
+            __resetSelectedStory()
+        st.session_state.selectedStoryTitle = ""
 
-    with st.chat_message("assistant", avatar=C.AI_ICON):
-        responseContainer = st.empty()
+        with st.chat_message("user", avatar=C.USER_ICON):
+            st.markdown(prompt)
+        U.pprint(f"{prompt=}")
+        st.session_state.chatHistory.append({"role": "user", "content": prompt })
+        st.session_state.messages.append({"role": "user", "content": prompt})
 
-        def __printAndGetResponse():
-            response = ""
-            responseContainer.image(C.TEXT_LOADER)
-            responseGenerator = predict()
+        with st.chat_message("assistant", avatar=C.AI_ICON):
+            responseContainer = st.empty()
 
-            for chunk in responseGenerator:
-                response += chunk
-                if __isInvalidResponse(response):
-                    U.pprint(f"InvalidResponse={response}")
-                    return
+            def __printAndGetResponse():
+                response = ""
+                responseContainer.image(C.TEXT_LOADER)
+                responseGenerator = __predict()
 
-                if C.JSON_SEPARATOR not in response:
-                    responseContainer.markdown(response)
+                for chunk in responseGenerator:
+                    response += chunk
+                    if __isInvalidResponse(response):
+                        U.pprint(f"InvalidResponse={response}")
+                        return
 
-            return response
+                    if C.JSON_SEPARATOR not in response:
+                        responseContainer.markdown(response)
 
-        response = __printAndGetResponse()
-        while not response:
-            U.pprint("Empty response. Retrying..")
-            time.sleep(0.7)
-            response = __printAndGetResponse()
+                return response
 
-        U.pprint(f"{response=}")
-
-        def selectButton(optionLabel):
-            st.session_state["buttonValue"] = optionLabel
-            U.pprint(f"Selected: {optionLabel}")
-
-        rawResponse = response
-        responseParts = response.split(C.JSON_SEPARATOR)
-
-        jsonStr = None
-        if len(responseParts) > 1:
-            [response, jsonStr] = responseParts
-
-        imageContainer = st.empty()
-        imagePath = __paintImageIfApplicable(imageContainer, prompt, response)
-
-        st.session_state.chatHistory.append({
-            "role": "assistant",
-            "content": response,
-            "image": imagePath,
-        })
-        st.session_state.messages.append({
-            "role": "assistant",
-            "content": rawResponse,
-        })
-
-        if jsonStr:
-            try:
-                json.loads(jsonStr)
-                jsonObj = json.loads(jsonStr)
-                options = jsonObj.get("options")
-                action = jsonObj.get("action")
-
-                if options:
-                    for option in options:
-                        st.button(
-                            option["label"],
-                            key=option["id"],
-                            on_click=lambda label=option["label"]: selectButton(label)
-                        )
-                elif action:
-                    U.pprint(f"{action=}")
-                    if action == "SHOW_STORY_DATABASE":
-                        time.sleep(0.5)
-                        st.switch_page("pages/popular-stories.py")
-                    # st.code(jsonStr, language="json")
-            except Exception as e:
-                U.pprint(e)
+            response = __printAndGetResponse()
+            while not response:
+                U.pprint("Empty response. Retrying..")
+                time.sleep(0.7)
+                response = __printAndGetResponse()
+
+            U.pprint(f"{response=}")
+
+            def selectButton(optionLabel):
+                st.session_state["buttonValue"] = optionLabel
+                U.pprint(f"Selected: {optionLabel}")
+
+            rawResponse = response
+            responseParts = response.split(C.JSON_SEPARATOR)
+
+            jsonStr = None
+            if len(responseParts) > 1:
+                [response, jsonStr] = responseParts
+
+            imageContainer = st.empty()
+            imagePath = __paintImageIfApplicable(imageContainer, prompt, response)
+
+            st.session_state.chatHistory.append({
+                "role": "assistant",
+                "content": response,
+                "image": imagePath,
+            })
+            st.session_state.messages.append({
+                "role": "assistant",
+                "content": rawResponse,
+            })
+
+            if jsonStr:
+                try:
+                    json.loads(jsonStr)
+                    jsonObj = json.loads(jsonStr)
+                    options = jsonObj.get("options")
+                    action = jsonObj.get("action")
+
+                    if options:
+                        for option in options:
+                            st.button(
+                                option["label"],
+                                key=option["id"],
+                                on_click=lambda label=option["label"]: selectButton(label)
+                            )
+                    elif action:
+                        U.pprint(f"{action=}")
+                        if action == "SHOW_STORY_DATABASE":
+                            time.sleep(0.5)
+                            st.switch_page("pages/popular-stories.py")
+                        # st.code(jsonStr, language="json")
+                except Exception as e:
+                    U.pprint(e)
+
+
+authenticateFunc(mainApp)
+showSidebar()

auth.py

@@ -0,0 +1,56 @@
+import os
+from typing import Callable, Any
+import streamlit as st
+from descope.descope_client import DescopeClient
+from descope.exceptions import AuthException
+
+DESCOPE_PROJECT_ID = os.environ.get("DESCOPE_PROJECT_ID")
+descopeClient = DescopeClient(project_id=DESCOPE_PROJECT_ID)
+
+
+def authenticateFunc(func: Callable[[], Any]):
+    if "token" not in st.session_state:
+        if "code" in st.query_params:
+            code = st.query_params["code"]
+            st.query_params.clear()
+            try:
+                with st.spinner("Authenticating ..."):
+                    jwtResponse = descopeClient.sso.exchange_token(code)
+                st.session_state["token"] = jwtResponse["sessionToken"].get("jwt")
+                st.session_state["refreshToken"] = jwtResponse["refreshSessionToken"].get(
+                    "jwt"
+                )
+                st.session_state["user"] = jwtResponse["user"]
+                st.rerun()
+            except AuthException:
+                st.error("Login failed!")
+
+        st.warning("You're not logged in. Please login to continue")
+        with st.container(border=False):
+            if st.button(
+                "Sign in with Google",
+                type="primary",
+                use_container_width=True
+            ):
+                oauthResponse = descopeClient.oauth.start(
+                    provider="google", return_url=st.context.headers["Origin"]
+                )
+                url = oauthResponse["url"]
+                # Redirect to Google
+                st.markdown(
+                    f'<meta http-equiv="refresh" content="0; url={url}">',
+                    unsafe_allow_html=True,
+                )
+    else:
+        try:
+            with st.spinner("Authenticating ..."):
+                jwtResponse = descopeClient.validate_and_refresh_session(
+                    st.session_state.token, st.session_state.refreshToken
+                )
+                st.session_state["token"] = jwtResponse["sessionToken"].get("jwt")
+
+            func()
+        except AuthException:
+            # Log out user
+            del st.session_state.token
+            st.rerun()

requirements.txt

@@ -5,3 +5,4 @@ transformers
 gradio_client
 anthropic
 supabase
+descope

sidebar.py

@@ -0,0 +1,57 @@
+import streamlit as st
+
+
+def showSidebar():
+    with st.sidebar:
+        if "user" in st.session_state and "token" in st.session_state:
+            st.markdown("""
+                <style>
+                .user-info {
+                    display: flex;
+                    align-items: center;
+                    padding: 10px;
+                    background-color: rgba(255, 255, 255, 0.1);
+                    border-radius: 5px;
+                    margin-bottom: 10px;
+                    margin-top: -2rem;
+                }
+                .user-avatar {
+                    width: 40px;
+                    height: 40px;
+                    border-radius: 50%;
+                    margin-right: 10px;
+                }
+                .user-details {
+                    flex-grow: 1;
+                }
+                .user-name {
+                    font-weight: bold;
+                    margin: 0;
+                }
+                .user-email {
+                    font-size: 0.8em;
+                    color: #888;
+                    margin: 0;
+                }
+                </style>
+                """, unsafe_allow_html=True)
+
+            user_avatar = st.session_state["user"].get("picture", "https://example.com/default-avatar.png")
+            user_name = st.session_state["user"].get("name", "User")
+            user_email = st.session_state["user"].get("email", "")
+
+            st.markdown(f"""
+                <div class="user-info">
+                    <img src="{user_avatar}" class="user-avatar">
+                    <div class="user-details">
+                        <p class="user-name">{user_name}</p>
+                        <p class="user-email">{user_email}</p>
+                    </div>
+                </div>
+                """, unsafe_allow_html=True)
+
+            if st.button("Logout", key="logout_button", type="secondary", use_container_width=True):
+                for key in ["token", "user"]:
+                    if key in st.session_state:
+                        del st.session_state[key]
+                st.rerun()

utils.py

@@ -100,6 +100,21 @@ def applyCommonStyles():
             height: 620px;
         }
 
+        div.stSpinner {
+            margin-left: auto;
+            margin-right: auto;
+            margin-top: 1rem;
+            width: fit-content;
+        }
+
+        div.stAlert {
+            margin-top: 1rem;
+        }
+
+        # section[data-testid="stSidebar"] {
+        #     width: 20vw !important;
+        # }
+
         </style>
         """,
         unsafe_allow_html=True