Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import gradio as gr
 from src.auth.auth import handle_login
 from src.auth.db import initialize_db
@@ -8,41 +9,57 @@ from src.interface import Interface
 initialize_db()
 load_dotenv()
 
-bot = None  # Initially, no bot is created until the user logs in or registers
-
+bot = None
+bot_ready = False
 
 def start_bot(userid, password, api_key_input):
-    global bot
+    global bot_ready
     login_status = handle_login(userid, password, api_key_input)
-    
-    if "successful" in login_status:  # Check for successful login
-        bot = Interface()  # Initialize after login success
+
+    if "✅" in login_status:
+        bot_ready = False
         return (
-            login_status,
+            "Login successful. Setting up your bot, please wait...",
             gr.update(visible=False),  # Hide login/registration section
-            gr.update(visible=True)    # Show chat section
+            gr.update(visible=False)   # Still hide chat section for now
         )
     else:
         return (
             login_status,
-            gr.update(visible=True),   # Keep login/registration section visible
-            gr.update(visible=False)   # Keep chat section hidden
+            gr.update(visible=True),    # Keep login/register open
+            gr.update(visible=False)    # Hide chat section
         )
 
+def setup_bot():
+    global bot, bot_ready
+    try:
+        bot = Interface()  # Interface will pick up API key from environment
+        bot_ready = True
+        return (
+            gr.update(visible=False),  # Hide login/registration
+            gr.update(visible=True)    # Show chat section
+        )
+    except Exception as e:
+        return (
+            gr.update(visible=True),
+            gr.update(visible=False)
+        )
 
 def answer(message, history):
-    answer_md, tables_display, images_display, retrieved_display = bot.get_answer(message)
-    
-    # Combine all parts into a single response string for chat
-    combined_response = f"{answer_md}\n\n{tables_display}"
+    if bot is None or not bot_ready:
+        return "Please log in and wait for your bot to be ready."
     
-    # Add images as markdown
-    if images_display:
-        combined_response += "\n\n" + "\n\n".join(images_display)
-    
-    return combined_response
+    try:
+        answer_md, tables_display, images_display, retrieved_display = bot.get_answer(message)
+
+        combined_response = f"{answer_md}\n\n{tables_display}"
 
+        if images_display:
+            combined_response += "\n\n" + "\n\n".join(images_display)
 
+        return combined_response
+    except Exception as e:
+        return f"An error occurred: {e}"
 
 with gr.Blocks(fill_height=True, fill_width=True) as app:
 
@@ -57,8 +74,9 @@ with gr.Blocks(fill_height=True, fill_width=True) as app:
 
             with gr.TabItem("Register"):
                 gr.Markdown("## 🔐 Enter Your Groq Cloud API Key")
-                gr.Markdown("You can create an API key at [Groq Cloud Console]"
-                "(https://console.groq.com/keys)")
+                gr.Markdown(
+                    "You can create an API key at [Groq Cloud Console](https://console.groq.com/keys)"
+                )
                 userid_register = gr.Textbox(label="UserID")
                 password_register = gr.Textbox(label="Password", type="password")
                 api_key_register = gr.Textbox(
@@ -67,31 +85,38 @@ with gr.Blocks(fill_height=True, fill_width=True) as app:
                     placeholder="sk-... (required)"
                 )
                 register_btn = gr.Button("Register")
-                register_output = gr.Textbox(label="Registration Status", 
-                                             interactive=False)
+                register_output = gr.Textbox(label="Registration Status", interactive=False)
 
     # Chat Section (Initially hidden)
     with gr.Column(visible=False) as chat_section:
         gr.ChatInterface(
-                answer,
-                title="🩺 Medico-Bot",
-                examples=["briefly explain me about cancer", "types of skin diseases?"],
-                flagging_options = ['Like', 'Dislike']
-                )
-
+            answer,
+            title="🩺 Medico-Bot",
+            examples=["briefly explain me about cancer", "types of skin diseases?"],
+            flagging_options=['Like', 'Dislike']
+        )
 
-    # Function connections
+    # Login button connection
     login_btn.click(
         start_bot,
-        inputs=[userid_login, password_login],
+        inputs=[userid_login, password_login, api_key_register],
         outputs=[login_output, login_register_section, chat_section]
+    ).then(
+        setup_bot,
+        inputs=[],
+        outputs=[login_register_section, chat_section]
     )
 
+    # Register button connection
     register_btn.click(
         start_bot,
         inputs=[userid_register, password_register, api_key_register],
         outputs=[register_output, login_register_section, chat_section]
+    ).then(
+        setup_bot,
+        inputs=[],
+        outputs=[login_register_section, chat_section]
     )
 
-app.queue()  
+app.queue(concurrency_count=1, max_size=5)
 app.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)