Update app.py

app.py

@@ -216,20 +216,21 @@ Please format the output with:
                 return None
             time.sleep(2)
 
-# ระบบ login
 # Main Streamlit interface
+# Initialize session state variables
+if 'login_step' not in st.session_state:
+    st.session_state.login_step = 'username'
 if 'username' not in st.session_state:
+    st.session_state.username = None
+
+# Login system
+if st.session_state.username is None:
     st.title("Login")
     username_input = st.text_input("Enter your username:")
-    
-    # Step 1: Initial login button
-    if st.button("Login"):
+
+    if st.session_state.login_step == 'username' and st.button("Login"):
         if username_input:
             if check_user_in_sheet(username_input):
-                # Store username in session but don't redirect yet
-                st.session_state['temp_username'] = username_input
-                
-                # Show welcome message and stats
                 stats = get_user_stats(username_input)
                 if stats:
                     st.success(f"Welcome, {username_input}! 👋")
@@ -238,40 +239,37 @@ if 'username' not in st.session_state:
                     - Today's Usage: {stats['daily_count']}/{DAILY_API_LIMIT} generations
                     - Total All-Time Usage: {stats['total_count']} generations
                     """)
-                    
-                    # Step 2: Enter application button
-                    if st.button("🎯 Enter GeneXam Application"):
-                        st.session_state['username'] = username_input
-                        del st.session_state['temp_username']
-                        st.rerun()
+                    st.session_state.login_step = 'enter_app'
             else:
                 st.warning("Username not found. Please try again.")
         else:
             st.warning("Please enter a valid username.")
-            
-    # Add some helpful information about the two-step process
-    if 'temp_username' in st.session_state:
-        st.info("👆 Click 'Enter GeneXam Application' to start generating questions!")
-    else:
+
+    if st.session_state.login_step == 'enter_app':
+        if st.button("🎯 Enter GeneXam Application"):
+            st.session_state.username = username_input
+            st.rerun()
+
+    # Show instructions
+    if st.session_state.login_step == 'username':
         st.markdown("""
         ### How to Login:
         1. Enter your username and click 'Login' to verify your account
-        2. Then click 'Enter GeneXam Application' to start using the system
+        2. After verification, click 'Enter GeneXam Application' to start using the system
         """)
-
 else:
-    # Main application interface remains the same
-    st.title(f"Welcome to GeneXam, {st.session_state['username']}! 🎓")
+    # Main application code (ส่วนที่เหลือเหมือนเดิม)
+    st.title(f"Welcome to GeneXam, {st.session_state.username}! 🎓")
     
     # Show current usage stats
-    stats = get_user_stats(st.session_state['username'])
+    stats = get_user_stats(st.session_state.username)
     if stats:
         remaining = DAILY_API_LIMIT - stats['daily_count']
         st.info(f"""
         📊 Usage Statistics:
         - Daily Generations Remaining: {remaining}/{DAILY_API_LIMIT}
         - Total All-Time Generations: {stats['total_count']}
-        """)    
+        """)
     
     # Create tabs for input methods
     tab1, tab2 = st.tabs(["Text Input", "PDF Upload"])