Update app.py

app.py

@@ -778,7 +778,41 @@ def generate_response_stream(prompt, temperature=0.7, backend=None, max_tokens=2
     final_tokens_per_sec = tokens_generated / elapsed if elapsed > 0 else 0
     yield "", False, final_tokens_per_sec, final_tokens_per_sec, False
 
-# PART 3 - Continue from Part 2 - FIXED VERSION
+# PART 3 - Continue from Part 2 - SESSION CODE VERSION
+
+import secrets
+
+# Global session codes storage
+active_sessions = {}  # {session_code: user_data}
+session_lock = threading.Lock()
+
+def generate_session_code():
+    """Generate a unique 4-digit session code."""
+    with session_lock:
+        while True:
+            code = ''.join([str(secrets.randbelow(10)) for _ in range(4)])
+            if code not in active_sessions:
+                return code
+
+def create_session(user_data):
+    """Create a new session and return the code."""
+    code = generate_session_code()
+    with session_lock:
+        active_sessions[code] = user_data
+    return code
+
+def validate_session(code):
+    """Validate a session code and return user data."""
+    with session_lock:
+        return active_sessions.get(code, None)
+
+def invalidate_session(code):
+    """Remove a session code."""
+    with session_lock:
+        if code in active_sessions:
+            del active_sessions[code]
+            return True
+        return False
 
 if __name__ == "__main__":
     import gradio as gr
@@ -804,6 +838,8 @@ if __name__ == "__main__":
     .send-btn { background: linear-gradient(135deg, #10a37f 0%, #0d8c6c 100%) !important; border: none !important; }
     .stop-btn { background: linear-gradient(135deg, #ef4444 0%, #dc2626 100%) !important; border: none !important; }
     .announcement-banner { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 20px 28px; border-radius: 12px; margin-bottom: 20px; box-shadow: 0 4px 6px rgba(0,0,0,0.1); text-align: center; font-size: 16px; font-weight: 500; line-height: 1.6; }
+    .session-code-box { background: linear-gradient(135deg, #10a37f 0%, #0d8c6c 100%); color: white; padding: 20px; border-radius: 12px; text-align: center; margin: 20px 0; box-shadow: 0 4px 12px rgba(0,0,0,0.2); }
+    .session-code-display { font-size: 32px; font-weight: 700; letter-spacing: 8px; margin: 10px 0; font-family: monospace; }
     """
 
     def format_message_html(role, content, show_thinking=True):
@@ -845,10 +881,10 @@ if __name__ == "__main__":
             html += format_message_html(msg["role"], msg["content"], show_thinking)
         return html
 
-    def render_limits_panel(user_session):
-        if not user_session or 'user_id' not in user_session:
+    def render_limits_panel(user_data):
+        if not user_data or 'user_id' not in user_data:
             return ""
-        info = get_user_limits_info(user_session['user_id'])
+        info = get_user_limits_info(user_data['user_id'])
         if not info:
             return ""
         plan_badge_class = f"plan-{info['plan']}"
@@ -870,31 +906,34 @@ if __name__ == "__main__":
         return html
 
     with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="slate")) as demo:
-        # Session state
-        user_session = gr.State({})
         
-        gr.HTML('<div class="announcement-banner">🔐 <strong>SAM-X-1 V3.0 - SECURE ACCESS</strong> 🔐<br>✨ Sign in to unlock AI-powered conversations!<br>🆓 FREE: Nano & Mini unlimited, Fast 10/3h, Large 8/3h<br>⭐ PLUS: Nano/Mini/Fast unlimited, Large 20/3h<br>💎 PRO: Everything unlimited!</div>')
+        gr.HTML('<div class="announcement-banner">🔐 <strong>SAM-X-1 V3.0 - SESSION CODE ACCESS</strong> 🔐<br>✨ Sign in to get your 4-digit session code!<br>🆓 FREE: Nano & Mini unlimited, Fast 10/3h, Large 8/3h<br>⭐ PLUS: Nano/Mini/Fast unlimited, Large 20/3h<br>💎 PRO: Everything unlimited!</div>')
         
         with gr.Tabs() as main_tabs:
             with gr.Tab("🔐 Sign In"):
                 with gr.Column():
                     login_username = gr.Textbox(label="Username", placeholder="Enter username")
                     login_password = gr.Textbox(label="Password", type="password", placeholder="Enter password")
-                    login_btn = gr.Button("Sign In", variant="primary")
+                    login_btn = gr.Button("Sign In", variant="primary", size="lg")
                     login_msg = gr.Markdown("")
+                    session_code_display = gr.HTML("")
             
             with gr.Tab("📝 Sign Up"):
                 with gr.Column():
                     signup_username = gr.Textbox(label="Username", placeholder="Choose a username")
                     signup_email = gr.Textbox(label="Email (optional)", placeholder="your@email.com")
                     signup_password = gr.Textbox(label="Password", type="password", placeholder="Choose a password")
-                    signup_btn = gr.Button("Create Account", variant="primary")
+                    signup_btn = gr.Button("Create Account", variant="primary", size="lg")
                     signup_msg = gr.Markdown("")
             
             with gr.Tab("💬 Chat") as chat_tab:
+                with gr.Row():
+                    chat_session_code = gr.Textbox(label="🔑 Enter Your 4-Digit Session Code", placeholder="0000", max_lines=1, scale=3)
+                    verify_session_btn = gr.Button("✅ Verify", variant="primary", scale=1)
+                
                 with gr.Row():
                     with gr.Column(scale=4):
-                        user_info = gr.Markdown("Please sign in first")
+                        user_info = gr.Markdown("❌ Not authenticated - Enter your session code above")
                     with gr.Column(scale=1):
                         logout_btn = gr.Button("🚪 Logout", size="sm")
                 
@@ -912,10 +951,10 @@ if __name__ == "__main__":
                 chat_html = gr.HTML(value="", elem_classes=["chat-container"])
                 
                 with gr.Row():
-                    msg_input = gr.Textbox(placeholder="Ask me anything...", show_label=False, scale=8)
+                    msg_input = gr.Textbox(placeholder="Enter session code first to chat...", show_label=False, scale=8, interactive=False)
                     with gr.Column(scale=1, min_width=120):
                         with gr.Row():
-                            send_btn = gr.Button("▶", variant="primary", elem_classes=["circular-btn", "send-btn"])
+                            send_btn = gr.Button("▶", variant="primary", elem_classes=["circular-btn", "send-btn"], interactive=False)
                             stop_btn = gr.Button("⏹", variant="stop", elem_classes=["circular-btn", "stop-btn"], interactive=False)
                 
                 with gr.Row():
@@ -946,24 +985,27 @@ if __name__ == "__main__":
                             request_msg = gr.Markdown("")
                 
                 with gr.Accordion("⭐ Request Plan Upgrade", visible=False, open=False) as upgrade_panel:
+                    upgrade_session_code = gr.Textbox(label="Your Session Code", placeholder="0000")
                     upgrade_plan_choice = gr.Radio(choices=["plus", "pro"], label="Select Plan", value="plus")
                     upgrade_reason = gr.Textbox(label="Reason for Upgrade", placeholder="Why do you need this upgrade?", lines=3)
                     submit_upgrade_btn = gr.Button("Submit Request", variant="primary")
                     upgrade_msg = gr.Markdown("")
+            
+            with gr.Tab("👨‍💼 Admin Access") as admin_tab:
+                admin_session_code = gr.Textbox(label="🔑 Enter Your Admin Session Code", placeholder="0000", max_lines=1)
+                verify_admin_btn = gr.Button("✅ Verify Admin", variant="primary", size="lg")
+                admin_verify_msg = gr.Markdown("")
+                admin_logout_btn = gr.Button("🚪 Logout", size="sm")
         
         # Event handlers
         def handle_login(username, password):
             success, user_data = authenticate_user(username, password)
             if success:
-                # Return the session data and success message
-                return (
-                    user_data, 
-                    f"✅ Welcome back, {username}! Redirecting to chat...", 
-                    "",
-                    gr.update(selected=2)  # Switch to Chat tab (index 2)
-                )
-            return {}, "❌ Invalid credentials!", "", gr.update()
-                
+                session_code = create_session(user_data)
+                code_html = f'<div class="session-code-box"><div style="font-size: 18px; margin-bottom: 10px;">✅ Login Successful!</div><div style="font-size: 16px; margin-bottom: 5px;">Your Session Code:</div><div class="session-code-display">{session_code}</div><div style="font-size: 14px; margin-top: 10px;">💡 Use this code in the Chat or Admin tab</div><div style="font-size: 13px; margin-top: 5px; opacity: 0.9;">⚠️ Keep this code private!</div></div>'
+                return f"✅ Welcome back, **{username}**! Use your session code above to access chat.", code_html
+            return "❌ Invalid credentials!", ""
+        
         def handle_signup(username, email, password):
             if len(username) < 3:
                 return "❌ Username must be at least 3 characters!"
@@ -971,36 +1013,24 @@ if __name__ == "__main__":
                 return "❌ Password must be at least 6 characters!"
             success, message = create_user(username, password, email)
             if success:
-                return f"✅ {message} Now sign in with your credentials!"
+                return f"✅ {message} Now sign in to get your session code!"
             return f"❌ {message}"
         
-        def load_user_interface(session):
-            # Add debug logging
-            print(f"DEBUG - load_user_interface called with session: {session}")
+        def verify_session_code(code):
+            if not code or len(code) != 4 or not code.isdigit():
+                return "❌ Invalid code format", "", gr.update(visible=False), gr.update(), gr.update(), gr.update(interactive=False, placeholder="Enter valid session code first..."), gr.update(interactive=False)
             
-            if not session or 'user_id' not in session:
-                # Return updates but don't show error message
-                return (
-                    "⚠️ Please sign in to use chat", 
-                    "", 
-                    gr.update(visible=False), 
-                    gr.update(choices=["🤖 Auto (Recommended)"], value="🤖 Auto (Recommended)"), 
-                    gr.update()
-                )
+            user_data = validate_session(code)
+            if not user_data:
+                return "❌ Invalid or expired session code", "", gr.update(visible=False), gr.update(), gr.update(), gr.update(interactive=False, placeholder="Enter valid session code first..."), gr.update(interactive=False)
             
-            info = get_user_limits_info(session['user_id'])
+            info = get_user_limits_info(user_data['user_id'])
             if not info:
-                return (
-                    "⚠️ Could not load user info", 
-                    "", 
-                    gr.update(visible=False), 
-                    gr.update(choices=["🤖 Auto (Recommended)"], value="🤖 Auto (Recommended)"), 
-                    gr.update()
-                )
+                return "❌ Could not load user info", "", gr.update(visible=False), gr.update(), gr.update(), gr.update(interactive=False), gr.update(interactive=False)
             
             plan_badge = f'<span class="plan-badge plan-{info["plan"]}">{info["plan"].upper()}</span>'
-            user_info_text = f"👤 **{session['username']}** {plan_badge}"
-            limits_html = render_limits_panel(session)
+            user_info_text = f"✅ **Authenticated as: {user_data['username']}** {plan_badge}"
+            limits_html = render_limits_panel(user_data)
             
             if info['can_choose_model']:
                 available_model_names = list(available_models.keys())
@@ -1008,33 +1038,37 @@ if __name__ == "__main__":
             else:
                 choices = ["🤖 Auto (Recommended)"]
             
-            # Show admin panel only for admins
-            is_admin = session.get('is_admin', False)
+            is_admin = user_data.get('is_admin', False)
             
             return (
                 user_info_text, 
                 limits_html, 
                 gr.update(visible=is_admin), 
                 gr.update(choices=choices, value="🤖 Auto (Recommended)"), 
-                gr.update(maximum=info['max_tokens'], value=min(256, info['max_tokens']))
+                gr.update(maximum=info['max_tokens'], value=min(256, info['max_tokens'])),
+                gr.update(interactive=True, placeholder="Ask me anything..."),
+                gr.update(interactive=True)
             )
-
         
-        def send_message_handler(message, show_thinking, temperature, model_choice, max_tokens, session):
+        def send_message_handler(message, show_thinking, temperature, model_choice, max_tokens, session_code):
             global stop_generation
             stop_generation.clear()
             
-            if not session or 'user_id' not in session:
-                return "", "", "❌ Not logged in", gr.update(), gr.update()
+            if not session_code or len(session_code) != 4:
+                return "", "", "❌ Invalid session code", gr.update(), gr.update()
+            
+            user_data = validate_session(session_code)
+            if not user_data:
+                return "", "", "❌ Session expired - please re-enter your code", gr.update(), gr.update()
             
             if not message.strip():
                 return "", "", "⚡ Ready", gr.update(interactive=True), gr.update(interactive=False)
             
-            info = get_user_limits_info(session['user_id'])
+            info = get_user_limits_info(user_data['user_id'])
             
             # Auto or manual model selection
             if model_choice == "🤖 Auto (Recommended)" or not info['can_choose_model']:
-                user_available = get_available_models_for_user(session['user_id'])
+                user_available = get_available_models_for_user(user_data['user_id'])
                 if not user_available:
                     return "", "", "❌ No models available (limits reached)", gr.update(interactive=True), gr.update(interactive=False)
                 backend = select_model_auto(message, available_models, user_available)
@@ -1043,18 +1077,18 @@ if __name__ == "__main__":
                 model_name = backend.get_name()
             else:
                 model_name = model_choice
-                can_use, msg = can_use_model(session['user_id'], model_name)
+                can_use, msg = can_use_model(user_data['user_id'], model_name)
                 if not can_use:
                     return "", "", f"❌ {msg}", gr.update(interactive=True), gr.update(interactive=False)
                 backend = available_models[model_name]
             
             # Final check
-            can_use, msg = can_use_model(session['user_id'], model_name)
+            can_use, msg = can_use_model(user_data['user_id'], model_name)
             if not can_use:
                 return "", "", f"❌ {msg}", gr.update(interactive=True), gr.update(interactive=False)
             
             # Increment usage
-            increment_model_usage(session['user_id'], model_name)
+            increment_model_usage(user_data['user_id'], model_name)
             
             yield "", "", f"⚡ Using {model_name}...", gr.update(interactive=False), gr.update(interactive=True)
             
@@ -1097,14 +1131,52 @@ if __name__ == "__main__":
         def show_upgrade_panel():
             return gr.update(visible=True, open=True)
         
-        def submit_upgrade_request(plan, reason, session):
-            if not session or 'user_id' not in session:
-                return "❌ Not logged in"
+        def submit_upgrade_request(session_code, plan, reason):
+            if not session_code or len(session_code) != 4:
+                return "❌ Invalid session code"
+            
+            user_data = validate_session(session_code)
+            if not user_data:
+                return "❌ Session expired"
+            
             if not reason.strip():
                 return "❌ Please provide a reason"
-            success, msg = request_upgrade(session['user_id'], plan, reason)
+            
+            success, msg = request_upgrade(user_data['user_id'], plan, reason)
             return f"{'✅' if success else '❌'} {msg}"
         
+        def handle_logout(session_code):
+            if session_code and len(session_code) == 4:
+                invalidate_session(session_code)
+            return (
+                "", 
+                "❌ Logged out - Session code invalidated", 
+                "", 
+                gr.update(visible=False), 
+                gr.update(choices=["🤖 Auto (Recommended)"], value="🤖 Auto (Recommended)"),
+                gr.update(value=256),
+                gr.update(interactive=False, placeholder="Enter session code first..."),
+                gr.update(interactive=False)
+            )
+        
+        def verify_admin_session(code):
+            if not code or len(code) != 4 or not code.isdigit():
+                return "❌ Invalid code format", gr.update(visible=False)
+            
+            user_data = validate_session(code)
+            if not user_data:
+                return "❌ Invalid or expired session code", gr.update(visible=False)
+            
+            if not user_data.get('is_admin', False):
+                return "❌ Access denied - Admin privileges required", gr.update(visible=False)
+            
+            return f"✅ Admin access granted for **{user_data['username']}**", gr.update(visible=True, open=True)
+        
+        def admin_logout_handler(code):
+            if code and len(code) == 4:
+                invalidate_session(code)
+            return "", "❌ Logged out", gr.update(visible=False)
+        
         def load_all_users():
             users = get_all_users()
             formatted = []
@@ -1137,27 +1209,31 @@ if __name__ == "__main__":
             success, msg = deny_request(int(request_id))
             return f"{'✅' if success else '❌'} {msg}"
         
-        def handle_logout():
-            return {}, "❌ Logged out - Please sign in", ""
-        
         # Wire up events
-        login_btn.click(handle_login, [login_username, login_password], [user_session, login_msg, chat_html])
+        login_btn.click(handle_login, [login_username, login_password], [login_msg, session_code_display])
         signup_btn.click(handle_signup, [signup_username, signup_email, signup_password], [signup_msg])
         
-        # Load user info when chat tab is selected
-        chat_tab.select(load_user_interface, [user_session], [user_info, limits_display, admin_panel, model_selector, max_tokens_slider])
+        # Session verification
+        verify_outputs = [user_info, limits_display, admin_panel, model_selector, max_tokens_slider, msg_input, send_btn]
+        verify_session_btn.click(verify_session_code, [chat_session_code], verify_outputs)
         
+        # Chat functionality
         send_outputs = [msg_input, chat_html, speed_display, send_btn, stop_btn]
-        send_btn.click(send_message_handler, [msg_input, show_thinking_checkbox, temperature_slider, model_selector, max_tokens_slider, user_session], send_outputs)
-        msg_input.submit(send_message_handler, [msg_input, show_thinking_checkbox, temperature_slider, model_selector, max_tokens_slider, user_session], send_outputs)
+        send_btn.click(send_message_handler, [msg_input, show_thinking_checkbox, temperature_slider, model_selector, max_tokens_slider, chat_session_code], send_outputs)
+        msg_input.submit(send_message_handler, [msg_input, show_thinking_checkbox, temperature_slider, model_selector, max_tokens_slider, chat_session_code], send_outputs)
         
         stop_btn.click(stop_generation_handler, outputs=[speed_display, send_btn, stop_btn])
         clear_btn.click(clear_chat, outputs=[chat_html, speed_display, send_btn, stop_btn])
         upgrade_btn.click(show_upgrade_panel, outputs=[upgrade_panel])
-        submit_upgrade_btn.click(submit_upgrade_request, [upgrade_plan_choice, upgrade_reason, user_session], [upgrade_msg])
-        login_btn.click(handle_login, [login_username, login_password], [user_session, login_msg, chat_html, main_tabs])
+        submit_upgrade_btn.click(submit_upgrade_request, [upgrade_session_code, upgrade_plan_choice, upgrade_reason], [upgrade_msg])
+        
+        logout_outputs = [chat_session_code, user_info, limits_display, admin_panel, model_selector, max_tokens_slider, msg_input, send_btn]
+        logout_btn.click(handle_logout, [chat_session_code], logout_outputs)
         
         # Admin panel events
+        verify_admin_btn.click(verify_admin_session, [admin_session_code], [admin_verify_msg, admin_panel])
+        admin_logout_btn.click(admin_logout_handler, [admin_session_code], [admin_session_code, admin_verify_msg, admin_panel])
+        
         refresh_users_btn.click(load_all_users, outputs=[users_table])
         refresh_requests_btn.click(load_pending_requests, outputs=[requests_table])
         update_plan_btn.click(admin_update_plan, [admin_username, admin_new_plan], [admin_msg])
@@ -1166,6 +1242,13 @@ if __name__ == "__main__":
         
         gr.Markdown("""
         ---
+        ### 🔑 How Session Codes Work
+        
+        1. **Sign In** on the "Sign In" tab to get your unique 4-digit code
+        2. **Copy** your session code (displayed after login)
+        3. **Enter** the code in the Chat or Admin tab to access features
+        4. **Logout** invalidates your code (you'll need to sign in again)
+        
         ### 📊 Plan Comparison
         
         | Feature | FREE | PLUS ⭐ | PRO 💎 |
@@ -1178,7 +1261,7 @@ if __name__ == "__main__":
         | **Max Tokens** | 256 | 384 | 512 |
         
         ### 🆓 Sign up for FREE account - Nano & Mini unlimited!
-        ### 👨‍💼 Admins: Sign in to access the Admin Panel in the Chat tab
+        ### 👨‍💼 Admins: Use your session code in the Admin Access tab
         """)
     
     demo.launch(debug=True, share=False, server_name="0.0.0.0", server_port=7860)