Update app.py

app.py

@@ -1,9 +1,11 @@
 # /home/user/app/app.py
 import streamlit as st
 from pathlib import Path
-from sqlmodel import select # <--- IMPORT SELECT FOR SQLMODEL QUERIES
+from datetime import datetime # For default chat session title
+from sqlmodel import select # For SQLModel queries
 
 from config.settings import settings
+# Ensure models are imported correctly based on your __init__.py in models/
 from models import (
     create_db_and_tables,
     get_session_context, # Your SQLModel session context manager
@@ -11,12 +13,12 @@ from models import (
     ChatMessage,
     ChatSession
 )
-from models.user import UserCreate
+from models.user import UserCreate # For type hinting
 from services.auth import create_user_in_db, authenticate_user
 from services.logger import app_logger
-from assets.logo import get_logo_path
+from assets.logo import get_logo_path # For consistent logo handling
 
-# --- Page Configuration --- (Assuming this is already correct)
+# --- Page Configuration ---
 st.set_page_config(
     page_title=settings.APP_TITLE,
     page_icon="⚕️",
@@ -24,22 +26,23 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
-# --- Database Initialization --- (Assuming this is already correct)
-@st.cache_resource
+# --- Database Initialization ---
+@st.cache_resource # Ensure this runs only once
 def init_db():
     app_logger.info("Initializing database and tables...")
     create_db_and_tables() # This uses SQLModel.metadata.create_all(engine)
     app_logger.info("Database initialized.")
+
 init_db()
 
-# --- Session State Initialization --- (Assuming this is already correct)
+# --- Session State Initialization (Using primitive types for user auth) ---
 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:
+if 'chat_messages' not in st.session_state: # For the current active chat in Consult page
     st.session_state.chat_messages = []
 
 
@@ -47,79 +50,100 @@ if 'chat_messages' not in st.session_state:
 def display_login_form():
     with st.form("login_form"):
         st.subheader("Login")
-        username_input = st.text_input("Username", key="login_username_input")
-        password_input = st.text_input("Password", type="password", key="login_password_input")
+        username_input = st.text_input("Username", key="login_username_input_main_app") # Unique key
+        password_input = st.text_input("Password", type="password", key="login_password_input_main_app") # Unique key
         submit_button = st.form_submit_button("Login")
 
         if submit_button:
             user_object_from_auth = authenticate_user(username_input, password_input)
 
-            if user_object_from_auth:
+            if user_object_from_auth: # Indicates successful authentication
                 st.success(f"Welcome back, {username_input}!")
-                app_logger.info(f"User {username_input} authenticated successfully.")
+                app_logger.info(f"User '{username_input}' authenticated successfully by authenticate_user.")
 
                 try:
                     with get_session_context() as db_session: # db_session is a SQLModel Session
-                        # --- SQLMODEL QUERY ---
+                        # Re-fetch the user to ensure we have a live object for this session
+                        # and to get definitive ID/username if authenticate_user returned a detached object.
                         statement = select(User).where(User.username == username_input)
                         live_user = db_session.exec(statement).first()
-                        # --------------------
 
                         if not live_user:
-                            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.")
+                            # This is a critical inconsistency if authenticate_user said OK.
+                            st.error("Authentication inconsistency. User details not found after login. Please contact support or try again.")
+                            app_logger.error(f"CRITICAL: User '{username_input}' authenticated but then NOT FOUND in DB by username query.")
+                            # Clear potentially corrupted auth state
                             st.session_state.authenticated_user_id = None
                             st.session_state.authenticated_username = None
-                            st.rerun()
+                            st.rerun() # Force rerun to show login page
                             return
 
+                        app_logger.info(f"Live user object for '{live_user.username}' (ID: {live_user.id}) obtained in session.")
+                        # Store primitive data in st.session_state
                         st.session_state.authenticated_user_id = live_user.id
                         st.session_state.authenticated_username = live_user.username
-                        app_logger.info(f"Stored user ID {live_user.id} and username '{live_user.username}' in session state.")
-
-                        new_chat_session = ChatSession(user_id=live_user.id, title=f"Session for {live_user.username} ({Path(settings.APP_TITLE).name})")
+                        app_logger.info(f"Stored user ID {live_user.id} and username '{live_user.username}' in Streamlit session state.")
+
+                        # Create a new chat session for the user upon login
+                        app_logger.debug(f"Attempting to create new chat session for user_id: {live_user.id}")
+                        new_chat_session = ChatSession(
+                            user_id=live_user.id,
+                            title=f"Session for {live_user.username} - {datetime.now().strftime('%Y-%m-%d %H:%M')}"
+                            # start_time is default_factory=datetime.utcnow in the model
+                        )
                         db_session.add(new_chat_session)
-                        # db_session.commit() # Handled by get_session_context manager
-                        db_session.refresh(new_chat_session) # Refresh after add, before commit if ID needed
+                        app_logger.debug("New ChatSession ORM object added to SQLModel session.")
+
+                        # Flush the session to persist new_chat_session and get its ID
+                        app_logger.debug("Flushing SQLModel session to assign ID to new_chat_session...")
+                        db_session.flush()
+                        app_logger.debug(f"Session flushed. new_chat_session potential ID: {new_chat_session.id}")
+
+                        # Check if ID was assigned; if not, something is wrong with DB/model setup for auto-increment PK.
+                        if new_chat_session.id is None:
+                            app_logger.error("CRITICAL: new_chat_session.id is None after flush. Cannot proceed with session creation.")
+                            raise Exception("Failed to obtain ID for new chat session after database flush.")
+
+                        # Refresh the object to load any server-side defaults or confirm state
+                        app_logger.debug(f"Refreshing new_chat_session (ID: {new_chat_session.id}) from database...")
+                        db_session.refresh(new_chat_session)
+                        app_logger.debug("new_chat_session ORM object refreshed.")
+
+                        # Store the new chat session ID for use in the Consult page
                         st.session_state.current_chat_session_id = new_chat_session.id
-                        st.session_state.chat_messages = []
-                        app_logger.info(f"New chat session (ID: {new_chat_session.id}) created for user {live_user.username}.")
-                    # Session commit and close handled by get_session_context
-                    st.rerun()
+                        st.session_state.chat_messages = [] # Clear any old messages for the new session
+                        app_logger.info(f"New chat session (ID: {new_chat_session.id}) created and assigned for user '{live_user.username}'.")
+                        # The commit will be handled by the get_session_context manager upon successful exit of this 'with' block.
+
+                    # If all operations within the 'with' block succeed, the context manager will commit.
+                    app_logger.info(f"Post-login setup complete for '{username_input}'. Rerunning app.")
+                    st.rerun() # Rerun to reflect login state and navigate to the default page (e.g., Home)
 
                 except Exception as e:
-                    app_logger.error(f"Error during post-login session setup for user {username_input}: {e}", exc_info=True)
+                    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 potentially partial auth state to force re-login or prevent inconsistent state
                     st.session_state.authenticated_user_id = None
                     st.session_state.authenticated_username = None
+                    st.session_state.current_chat_session_id = None
+                    # No rerun here, let user see the error and try again or contact support.
             else:
                 st.error("Invalid username or password.")
                 app_logger.warning(f"Failed login attempt for username: {username_input}")
 
-# display_signup_form() remains largely the same as it calls create_user_in_db,
-# which should internally use SQLModel syntax if it interacts with the DB.
-
-# --- Main App Logic --- (Remains the same, using updated session state keys)
-# ... (rest of app.py as previously provided, ensuring it uses authenticated_user_id and authenticated_username)
-# For brevity, I'm omitting the parts of app.py that don't change due to SQLModel query syntax.
-# The signup form and the main structure checking st.session_state.get("authenticated_user_id")
-# and the sidebar display are correct from the previous version.
-# Ensure any other direct DB queries in app.py (if any were added) are also updated.
-
-# (Make sure the rest of your app.py from the previous good version is here)
 def display_signup_form():
     with st.form("signup_form"):
         st.subheader("Sign Up")
-        new_username = st.text_input("Choose a Username", key="signup_username_input")
-        new_email = st.text_input("Email (Optional)", key="signup_email_input")
-        new_password = st.text_input("Choose a Password", type="password", key="signup_password_input")
-        confirm_password = st.text_input("Confirm Password", type="password", key="signup_confirm_password_input")
+        new_username = st.text_input("Choose a Username", key="signup_username_input_main_app") # Unique key
+        new_email = st.text_input("Email (Optional)", key="signup_email_input_main_app") # Unique key
+        new_password = st.text_input("Choose a Password", type="password", key="signup_password_input_main_app") # Unique key
+        confirm_password = st.text_input("Confirm Password", type="password", key="signup_confirm_password_input_main_app") # Unique key
         submit_button = st.form_submit_button("Sign Up")
 
         if submit_button:
             if not new_username or not new_password:
                 st.error("Username and password are required.")
-            elif len(new_password) < 6:
+            elif len(new_password) < 6: # Example: Minimum password length
                  st.error("Password must be at least 6 characters long.")
             elif new_password != confirm_password:
                 st.error("Passwords do not match.")
@@ -128,45 +152,84 @@ def display_signup_form():
                     username=new_username,
                     password=new_password,
                     email=new_email if new_email else None
+                    # full_name can be added if part of UserCreate and signup form
                 )
+                # create_user_in_db should handle its own session and return a User object or None.
+                # It's responsible for hashing password and committing the new user.
                 user = create_user_in_db(user_data)
                 if user:
+                    # Use 'new_username' (from form input) for the message
                     st.success(f"Account created for {new_username}. Please log in.")
-                    app_logger.info(f"Account created for {new_username}.")
+                    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}")
+                    # create_user_in_db returns None if user exists or another error occurred.
+                    # The function itself should log the specific reason.
+                    st.error("Username or Email might already be taken, or another error occurred during signup. Please check logs or try different credentials.")
+                    app_logger.warning(f"Signup failed for username: '{new_username}'. create_user_in_db returned None.")
 
-if not st.session_state.get("authenticated_user_id"):
+# --- Main App Logic (Checks for authenticated_user_id) ---
+if not st.session_state.get("authenticated_user_id"): # Check if user_id is set (i.e., user is logged in)
     st.title(f"Welcome to {settings.APP_TITLE}")
     st.markdown("Your AI-powered partner for advanced healthcare insights.")
+
     login_tab, signup_tab = st.tabs(["Login", "Sign Up"])
     with login_tab:
         display_login_form()
     with signup_tab:
         display_signup_form()
 else:
+    # --- User is Authenticated - Display Sidebar and Page Content ---
     with st.sidebar:
-        logo_path_str_sidebar = get_logo_path()
+        # Display Logo using get_logo_path for consistency
+        logo_path_str_sidebar = get_logo_path() # Assumes assets/logo.py is set up
         if logo_path_str_sidebar:
             logo_file_sidebar = Path(logo_path_str_sidebar)
             if logo_file_sidebar.exists():
                 try:
-                    st.image(str(logo_file_sidebar), width=100)
+                    st.image(str(logo_file_sidebar), width=100) # Adjust width as needed
                 except Exception as e:
                     app_logger.warning(f"Could not display logo in sidebar from path '{logo_path_str_sidebar}': {e}")
-        elif settings.APP_TITLE:
+            # else: # get_logo_path should log if file doesn't exist
+            #    app_logger.warning(f"Sidebar logo path from get_logo_path() does not exist: {logo_path_str_sidebar}")
+        elif settings.APP_TITLE: # Fallback to title if no logo path or logo load fails
              st.sidebar.markdown(f"#### {settings.APP_TITLE}")
-        username_for_display = st.session_state.get("authenticated_username", "User")
+
+        # Use the stored primitive username
+        username_for_display = st.session_state.get("authenticated_username", "User") # Fallback to "User"
         st.sidebar.markdown(f"### Welcome, {username_for_display}!")
-        st.sidebar.markdown("---")
-        if st.sidebar.button("Logout", key="sidebar_logout_button"):
+        st.sidebar.markdown("---") # Separator
+
+        # Navigation is handled by Streamlit's multi-page app feature (pages/ directory)
+        # This app.py acts as the entry point.
+        # Streamlit automatically lists pages from the `pages/` directory here in the sidebar.
+
+        if st.sidebar.button("Logout", key="sidebar_logout_button_main_app"): # Unique key
             logged_out_username = st.session_state.get("authenticated_username", "UnknownUser")
-            app_logger.info(f"User {logged_out_username} logging out.")
+            app_logger.info(f"User '{logged_out_username}' logging out.")
+            # Clear all relevant session state variables
             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 = []
+            # Potentially clear other page-specific states if needed
+            # for key in list(st.session_state.keys()):
+            #     if key not in ['rerun_requested', ...]: # preserve essential Streamlit keys
+            #         del st.session_state[key]
             st.success("You have been logged out.")
             st.rerun()
-app_logger.info(f"Streamlit app '{settings.APP_TITLE}' initialized and running.")
+
+    # Main content area for app.py when logged in.
+    # If you have a `pages/1_Home.py`, Streamlit will render that page's content here by default
+    # when the user navigates to "Home" via the sidebar (or if it's the first page alphabetically).
+    # The content below in app.py might only be visible briefly or if no page is explicitly selected,
+    # or if `app.py` is configured as a page itself.
+
+    # For a typical MPA setup, you might not need much content directly in app.py's main area
+    # once the user is logged in, as page files will take over.
+    # A simple instruction or pointer could be useful.
+    st.sidebar.success("Select a page from the navigation.")
+    # Optionally, add a default landing message in the main area if no page is loaded yet,
+    # though Streamlit usually picks the first page in `pages/`.
+    # st.info("Main application area. Select a page from the sidebar to begin.")
+
+app_logger.info(f"Streamlit app '{settings.APP_TITLE}' (app.py) processed.")