Update app.py

app.py

@@ -1,16 +1,12 @@
 import streamlit as st
-from pathlib import Path # Ensure Path is imported if used for settings.LOGO_PATH
-# Make sure sqlalchemy.orm.exc is available if you want to catch DetachedInstanceError specifically
-from sqlalchemy.orm.exc import DetachedInstanceError as SQLAlchemyDetachedInstanceError
-
+from pathlib import Path
+from sqlalchemy.orm.exc import DetachedInstanceError as SQLAlchemyDetachedInstanceError # For specific error catching if needed
 
 from config.settings import settings
 from models import create_db_and_tables, get_session_context, User, ChatMessage, ChatSession
-from models.user import UserCreate # For type hinting
+from models.user import UserCreate
 from services.auth import create_user_in_db, authenticate_user
 from services.logger import app_logger
-# Agent will be initialized and used in specific pages like Consult.py
-# from agent import get_agent_executor # Import in pages where needed
 
 # --- Page Configuration ---
 st.set_page_config(
@@ -21,7 +17,7 @@ st.set_page_config(
 )
 
 # --- Database Initialization ---
-@st.cache_resource # Ensure this runs only once
+@st.cache_resource
 def init_db():
     app_logger.info("Initializing database and tables...")
     create_db_and_tables()
@@ -29,15 +25,16 @@ def init_db():
 
 init_db()
 
-# --- Session State Initialization ---
-if 'authenticated_user' not in st.session_state:
-    st.session_state.authenticated_user = None
+# --- Session State Initialization (Using new keys) ---
+if 'authenticated_user_id' not in st.session_state:
+    st.session_state.authenticated_user_id = None
+if 'authenticated_username' not in st.session_state:
+    st.session_state.authenticated_username = None
 if 'current_chat_session_id' not in st.session_state:
     st.session_state.current_chat_session_id = None
 if 'chat_messages' not in st.session_state:
     st.session_state.chat_messages = []
 
-
 # --- Authentication Logic ---
 def display_login_form():
     with st.form("login_form"):
@@ -47,67 +44,57 @@ def display_login_form():
         submit_button = st.form_submit_button("Login")
 
         if submit_button:
-            user_from_auth = authenticate_user(username_input, password_input)
+            # authenticate_user should ideally return a User object with at least ID loaded,
+            # or be robust enough that its session handling doesn't cause immediate detachment issues
+            # if we were to access its attributes directly (though we now avoid that for long-term storage).
+            user_object_from_auth = authenticate_user(username_input, password_input)
 
-            if user_from_auth:
-                # Success message uses the username from form input, which is safe.
-                st.success(f"Welcome back, {username_input}!")
-                app_logger.info(f"User {username_input} logged in successfully.")
+            if user_object_from_auth:
+                st.success(f"Welcome back, {username_input}!") # Use form input for immediate message
+                app_logger.info(f"User {username_input} authenticated successfully.")
 
                 try:
                     with get_session_context() as db_session:
-                        # WORKAROUND for DetachedInstanceError from user_from_auth.id:
-                        # Re-fetch the user by username within this new session.
-                        # This ensures 'live_user' is attached and its attributes are loaded.
-                        # The PREFERRED FIX is in 'services.auth.authenticate_user' to return
-                        # a User object with its 'id' (and other essential attributes) already loaded/refreshed.
-                        app_logger.info(f"Attempting to fetch live user details for '{username_input}' in new session.")
+                        # Fetch the "live" user to get definite ID and current username,
+                        # and to ensure it's attached to *this* session for creating the chat session.
+                        # The primary fix for 'authenticate_user' is still for it to return a usable object,
+                        # but re-fetching here ensures data for st.session_state is from a controlled point.
                         live_user = db_session.query(User).filter(User.username == username_input).first()
 
                         if not live_user:
-                            # This case means authentication (somehow) passed,
-                            # but the user is not findable by username immediately after.
-                            # This indicates a more serious data consistency issue or flaw in authenticate_user.
-                            st.error("Authentication inconsistency. User details not found after successful login. Please contact support.")
-                            app_logger.error(f"CRITICAL: User '{username_input}' authenticated but then not found by username in new session context.")
-                            st.session_state.authenticated_user = None # Clear any partial auth state
-                            st.rerun() # Rerun to show login page again
+                            st.error("Authentication inconsistency. User details not found after login. Please contact support.")
+                            app_logger.error(f"CRITICAL: User '{username_input}' authenticated but then not found in DB by username.")
+                            # Clear any potentially partially set auth state from previous attempts
+                            st.session_state.authenticated_user_id = None
+                            st.session_state.authenticated_username = None
+                            st.rerun()
                             return
 
-                        # Store the "live" user object (attached to db_session and attributes loaded)
-                        st.session_state.authenticated_user = live_user
-                        app_logger.info(f"Live user '{live_user.username}' (ID: {live_user.id}) fetched and stored in session_state.")
+                        # IMPORTANT: Store primitive data in st.session_state
+                        # Do this *before* the commit for new_chat_session, as commit might expire 'live_user' attributes.
+                        st.session_state.authenticated_user_id = live_user.id
+                        st.session_state.authenticated_username = live_user.username # Username from DB
+                        app_logger.info(f"Stored user ID {live_user.id} and username '{live_user.username}' in session state.")
 
-                        # Now use 'live_user' for creating the chat session
+                        # Now create the chat session using the live_user's ID
                         new_chat_session = ChatSession(user_id=live_user.id, title=f"Session for {live_user.username}")
                         db_session.add(new_chat_session)
-                        db_session.commit()
-                        db_session.refresh(new_chat_session) # Refresh to get DB-generated values for new_chat_session
+                        db_session.commit() # This commit is for new_chat_session
+                        # After commit, live_user attributes might be expired if expire_on_commit=True.
+                        # But we've already stored the primitives we need.
+                        db_session.refresh(new_chat_session) # Get ID for new_chat_session
                         st.session_state.current_chat_session_id = new_chat_session.id
-                        st.session_state.chat_messages = [] # Clear previous messages
+                        st.session_state.chat_messages = []
                         app_logger.info(f"New chat session (ID: {new_chat_session.id}) created for user {live_user.username}.")
+
                     st.rerun() # Rerun to reflect login state and navigate
 
-                except SQLAlchemyDetachedInstanceError as detached_error:
-                    # This specific catch might be less likely if the above workaround (re-fetching live_user) is effective.
-                    # However, it's good for diagnostics if the issue persists due to other detached instances.
-                    app_logger.error(
-                        f"DetachedInstanceError occurred during session setup for user '{username_input}'. "
-                        f"This typically means an object from 'authenticate_user' or 'create_user_in_db' "
-                        f"had expired attributes. Error: {detached_error}",
-                        exc_info=True
-                    )
-                    st.error(
-                        "Could not start a new session due to an object state issue. "
-                        "Please ensure user services return objects with loaded attributes."
-                    )
-                    st.session_state.authenticated_user = None
                 except Exception as e:
-                    app_logger.error(f"Error creating chat session for user {username_input}: {e}", exc_info=True)
-                    st.error(f"Could not start a new session: {e}")
-                    # Optionally, clear authenticated_user if session creation is critical and failed
-                    # st.session_state.authenticated_user = None
-                    # st.rerun()
+                    app_logger.error(f"Error during post-login session setup for user {username_input}: {e}", exc_info=True)
+                    st.error(f"Could not complete login process: {e}")
+                    # Clear auth state on error
+                    st.session_state.authenticated_user_id = None
+                    st.session_state.authenticated_username = None
             else:
                 st.error("Invalid username or password.")
                 app_logger.warning(f"Failed login attempt for username: {username_input}")
@@ -132,18 +119,16 @@ def display_signup_form():
                     password=new_password,
                     email=new_email if new_email else None
                 )
-                # 'create_user_in_db' should return a User object with attributes loaded
-                # (e.g., by calling session.refresh(user) before its session closes if a commit occurred).
-                user = create_user_in_db(user_data)
+                user = create_user_in_db(user_data) # Should handle its own session and return User or None
                 if user:
-                    st.success(f"Account created for {new_username}. Please log in.")
+                    st.success(f"Account created for {new_username}. Please log in.") # Use form input
                     app_logger.info(f"Account created for {new_username}.")
                 else:
                     st.error("Username might already be taken or another error occurred during signup.")
                     app_logger.warning(f"Signup failed for username: {new_username}")
 
-# --- Main App Logic ---
-if not st.session_state.authenticated_user:
+# --- Main App Logic (Checks for authenticated_user_id) ---
+if not st.session_state.get("authenticated_user_id"): # Check if user_id is set
     st.title(f"Welcome to {settings.APP_TITLE}")
     st.markdown("Your AI-powered partner for advanced healthcare insights.")
 
@@ -153,14 +138,11 @@ if not st.session_state.authenticated_user:
     with signup_tab:
         display_signup_form()
 else:
+    # User is authenticated (authenticated_user_id is present)
     with st.sidebar:
-        try:
-            # This relies on st.session_state.authenticated_user being a "live" object
-            # with 'username' attribute accessible. The login logic now aims to ensure this.
-            st.markdown(f"### Welcome, {st.session_state.authenticated_user.username}!")
-        except AttributeError as e:
-             st.markdown(f"### Welcome!")
-             app_logger.error(f"Could not access username for authenticated_user in sidebar: {e}. User object: {st.session_state.authenticated_user}", exc_info=True)
+        # Use the stored primitive username
+        username_for_display = st.session_state.get("authenticated_username", "User")
+        st.markdown(f"### Welcome, {username_for_display}!")
 
         logo_path_str = getattr(settings, "LOGO_PATH", None)
         if logo_path_str:
@@ -178,14 +160,11 @@ else:
         st.markdown("---")
 
         if st.button("Logout"):
-            username_on_logout = "UnknownUser"
-            try:
-                username_on_logout = st.session_state.authenticated_user.username
-            except Exception:
-                app_logger.warning("Could not get username during logout from session state object.")
-            app_logger.info(f"User {username_on_logout} logging out.")
-
-            st.session_state.authenticated_user = None
+            logged_out_username = st.session_state.get("authenticated_username", "UnknownUser")
+            app_logger.info(f"User {logged_out_username} logging out.")
+            # Clear authentication state
+            st.session_state.authenticated_user_id = None
+            st.session_state.authenticated_username = None
             st.session_state.current_chat_session_id = None
             st.session_state.chat_messages = []
             st.success("You have been logged out.")