Update app.py

app.py

@@ -1,5 +1,8 @@
 import gradio as gr
 import os
+import secrets
+from functools import partial
+
 api_token = os.getenv("HF_TOKEN")
 
 from langchain_community.vectorstores import FAISS
@@ -13,9 +16,14 @@ from langchain_community.llms import HuggingFaceEndpoint
 list_llm = ["meta-llama/Meta-Llama-3-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.2"]  
 list_llm_simple = [os.path.basename(llm) for llm in list_llm]
 
-# Simulated user credentials (replace with a secure method in production)
-VALID_USERNAME = "admin"
-VALID_PASSWORD = "password123"
+# Simulated user database (replace with a real database in production)
+USER_DB = {
+    "admin": {"password": "securepass123", "email": "[email protected]"},
+    "user1": {"password": "userpass456", "email": "[email protected]"}
+}
+
+# Session storage (in-memory for simplicity)
+SESSIONS = {}
 
 # Load and split PDF document
 def load_doc(list_file_path):
@@ -121,31 +129,50 @@ def conversation(qa_chain, message, history, language):
     return qa_chain, gr.update(value=""), new_history, response_source1, response_source1_page, response_source2, response_source2_page, response_source3, response_source3_page
 
 # Login function
-def check_login(username, password):
-    if username == VALID_USERNAME and password == VALID_PASSWORD:
-        return True, "Login successful! Access the chatbot below."
+def login(username, password):
+    if username in USER_DB and USER_DB[username]["password"] == password:
+        session_token = secrets.token_hex(16)
+        SESSIONS[session_token] = username
+        return True, session_token, f"Welcome, {username}! You are now logged in."
     else:
-        return False, "Invalid username or password. Please try again."
+        return False, None, "Invalid username or password. Please try again."
 
-# Main demo with login
+# Logout function
+def logout(session_token):
+    if session_token in SESSIONS:
+        del SESSIONS[session_token]
+    return False, None, "You have been logged out."
+
+# Main demo with modern login
 def demo():
-    with gr.Blocks(theme=gr.themes.Default(primary_hue="red", secondary_hue="pink", neutral_hue="sky")) as demo:
+    with gr.Blocks(
+        theme=gr.themes.Soft(primary_hue="blue", secondary_hue="gray", neutral_hue="slate"),
+        css="""
+            .login-box { max-width: 400px; margin: 50px auto; padding: 20px; border-radius: 10px; box-shadow: 0 4px 8px rgba(0,0,0,0.1); }
+            .title { text-align: center; font-size: 2em; margin-bottom: 20px; }
+            .button { background-color: #007bff; color: white; border-radius: 5px; }
+            .button:hover { background-color: #0056b3; }
+        """
+    ) as demo:
         # State variables
         vector_db = gr.State()
         qa_chain = gr.State()
         logged_in = gr.State(value=False)
+        session_token = gr.State(value=None)
 
         # Login interface
-        with gr.Column(visible=True) as login_col:
-            gr.HTML("<center><h1>RAG PDF Chatbot - Login</h1></center>")
-            username = gr.Textbox(label="Username", placeholder="Enter username")
-            password = gr.Textbox(label="Password", type="password", placeholder="Enter password")
-            login_btn = gr.Button("Login")
-            login_message = gr.Textbox(value="Please log in to access the chatbot.", show_label=False)
+        with gr.Column(elem_classes="login-box", visible=True) as login_col:
+            gr.HTML("<h1 class='title'>RAG PDF Chatbot Login</h1>")
+            username = gr.Textbox(label="Username", placeholder="Enter your username", lines=1)
+            password = gr.Textbox(label="Password", type="password", placeholder="Enter your password", lines=1)
+            login_btn = gr.Button("Login", elem_classes="button")
+            login_message = gr.Textbox(value="Please log in to access the chatbot.", show_label=False, interactive=False)
 
         # Chatbot interface (hidden until login)
         with gr.Column(visible=False) as chatbot_col:
-            gr.HTML("<center><h1>RAG PDF Chatbot</h1></center>")
+            with gr.Row():
+                gr.HTML("<h1 class='title'>RAG PDF Chatbot</h1>")
+                logout_btn = gr.Button("Logout", elem_classes="button", scale=0)
             gr.Markdown("""<b>Query your PDF documents!</b> This AI agent is designed to perform retrieval augmented generation (RAG) on PDF documents. \
             <b>Please do not upload confidential documents.</b>""")
             
@@ -153,7 +180,7 @@ def demo():
                 with gr.Column(scale=86):
                     gr.Markdown("<b>Step 1 - Upload PDF documents and Initialize RAG pipeline</b>")
                     document = gr.Files(height=300, file_count="multiple", file_types=["pdf"], interactive=True, label="Upload PDF documents")
-                    db_btn = gr.Button("Create vector database")
+                    db_btn = gr.Button("Create vector database", elem_classes="button")
                     db_progress = gr.Textbox(value="Not initialized", show_label=False)
                     gr.Markdown("<b>Select Large Language Model (LLM) and input parameters</b>")
                     llm_btn = gr.Radio(list_llm_simple, label="Available LLMs", value=list_llm_simple[0], type="index")
@@ -161,7 +188,7 @@ def demo():
                         slider_temperature = gr.Slider(minimum=0.01, maximum=1.0, value=0.5, step=0.1, label="Temperature", interactive=True)
                         slider_maxtokens = gr.Slider(minimum=128, maximum=9192, value=4096, step=128, label="Max New Tokens", interactive=True)
                         slider_topk = gr.Slider(minimum=1, maximum=10, value=3, step=1, label="top-k", interactive=True)
-                    qachain_btn = gr.Button("Initialize Question Answering Chatbot")
+                    qachain_btn = gr.Button("Initialize Question Answering Chatbot", elem_classes="button")
                     llm_progress = gr.Textbox(value="Not initialized", show_label=False)
 
                 with gr.Column(scale=200):
@@ -176,14 +203,14 @@ def demo():
                         doc_source3 = gr.Textbox(label="Reference 3", lines=2, container=True, scale=20)
                         source3_page = gr.Number(label="Page", scale=1)
                     msg = gr.Textbox(placeholder="Ask a question", container=True)
-                    submit_btn = gr.Button("Submit")
+                    submit_btn = gr.Button("Submit", elem_classes="button")
                     clear_btn = gr.ClearButton([msg, chatbot], value="Clear")
 
         # Login event
         login_btn.click(
-            fn=check_login,
+            fn=login,
             inputs=[username, password],
-            outputs=[logged_in, login_message]
+            outputs=[logged_in, session_token, login_message]
         ).then(
             fn=lambda logged: (gr.update(visible=not logged), gr.update(visible=logged)),
             inputs=[logged_in],
@@ -191,7 +218,24 @@ def demo():
             queue=False
         )
 
-        # Preprocessing events (only accessible after login)
+        # Logout event
+        logout_btn.click(
+            fn=logout,
+            inputs=[session_token],
+            outputs=[logged_in, session_token, login_message]
+        ).then(
+            fn=lambda logged: (gr.update(visible=not logged), gr.update(visible=logged)),
+            inputs=[logged_in],
+            outputs=[login_col, chatbot_col],
+            queue=False
+        ).then(
+            fn=lambda: gr.update(value="Please log in to access the chatbot."),
+            inputs=None,
+            outputs=[login_message],
+            queue=False
+        )
+
+        # Preprocessing events
         db_btn.click(initialize_database, inputs=[document], outputs=[vector_db, db_progress])
         qachain_btn.click(initialize_LLM, inputs=[llm_btn, slider_temperature, slider_maxtokens, slider_topk, vector_db], outputs=[qa_chain, llm_progress]).then(
             lambda: [None, "", 0, "", 0, "", 0], inputs=None, outputs=[chatbot, doc_source1, source1_page, doc_source2, source2_page, doc_source3, source3_page], queue=False