Update app.py

app.py

@@ -1,22 +1,21 @@
 # /home/user/app/app.py
 import streamlit as st
 from pathlib import Path
-from datetime import datetime # For default chat session title
-from sqlmodel import select # For SQLModel queries
+from datetime import datetime
+from sqlmodel import select
 
 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
+    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
-from assets.logo import get_logo_path # For consistent logo handling
+from assets.logo import get_logo_path
 
 # --- Page Configuration ---
 st.set_page_config(
@@ -27,209 +26,166 @@ 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() # This uses SQLModel.metadata.create_all(engine)
+    create_db_and_tables()
     app_logger.info("Database initialized.")
-
 init_db()
 
-# --- Session State Initialization (Using primitive types for user auth) ---
+# --- Session State Initialization ---
 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: # For the current active chat in Consult page
+if 'chat_messages' not in st.session_state:
     st.session_state.chat_messages = []
+if 'user_understands_disclaimer' not in st.session_state:
+    st.session_state.user_understands_disclaimer = False
 
+# --- Disclaimer Function ---
+def display_disclaimer():
+    st.warning(f"**Important Disclaimer:** {settings.MAIN_DISCLAIMER_LONG}")
+    st.info(settings.SIMULATION_DISCLAIMER)
+    if st.button("I Understand and Agree to Proceed", key="disclaimer_agree_btn"):
+        st.session_state.user_understands_disclaimer = True
+        st.rerun()
 
 # --- Authentication Logic ---
 def display_login_form():
+    # (This function remains the same as the last full app.py provided, which had the fix for
+    #  `Instance '<ChatSession ...>' is not persistent within this Session`.
+    #  Ensure that version is used here for robustness.)
     with st.form("login_form"):
         st.subheader("Login")
-        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
+        username_input = st.text_input("Username", key="login_username_input_main_app")
+        password_input = st.text_input("Password", type="password", key="login_password_input_main_app")
         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: # Indicates successful authentication
+            if user_object_from_auth:
                 st.success(f"Welcome back, {username_input}!")
-                app_logger.info(f"User '{username_input}' authenticated successfully by authenticate_user.")
-
+                app_logger.info(f"User '{username_input}' authenticated successfully.")
                 try:
-                    with get_session_context() as db_session: # db_session is a SQLModel Session
-                        # 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.
+                    with get_session_context() as db_session:
                         statement = select(User).where(User.username == username_input)
                         live_user = db_session.exec(statement).first()
-
                         if not live_user:
-                            # 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.error("Auth inconsistency. User not found post-login.")
+                            app_logger.error(f"CRITICAL: User '{username_input}' auth OK but NOT FOUND in DB.")
                             st.session_state.authenticated_user_id = None
                             st.session_state.authenticated_username = None
-                            st.rerun() # Force rerun to show login page
+                            st.rerun()
                             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 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}")
+                        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} - {datetime.now().strftime('%Y-%m-%d %H:%M')}"
-                            # start_time is default_factory=datetime.utcnow in the model
                         )
                         db_session.add(new_chat_session)
-                        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...")
+                        app_logger.debug("New ChatSession added to SQLModel 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.
+                        app_logger.debug(f"Session flushed. new_chat_session ID: {new_chat_session.id}")
                         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...")
+                            app_logger.error("CRITICAL: new_chat_session.id is None after flush.")
+                            raise Exception("Failed to obtain ID for new chat session after flush.")
                         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
+                        app_logger.debug("new_chat_session refreshed.")
                         st.session_state.current_chat_session_id = new_chat_session.id
-                        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.
+                        st.session_state.chat_messages = []
+                        app_logger.info(f"New chat session (ID: {new_chat_session.id}) created for user '{live_user.username}'.")
                     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)
-
+                    st.rerun()
                 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 '{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}")
 
 def display_signup_form():
+    # (This function remains largely the same as the last full app.py provided.
+    #  Ensure it correctly calls your `create_user_in_db` which handles hashing and DB commit.)
     with st.form("signup_form"):
         st.subheader("Sign Up")
-        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
+        new_username = st.text_input("Choose a Username", key="signup_username_input_main_app")
+        new_email = st.text_input("Email (Optional)", key="signup_email_input_main_app")
+        new_password = st.text_input("Choose a Password", type="password", key="signup_password_input_main_app")
+        confirm_password = st.text_input("Confirm Password", type="password", key="signup_confirm_password_input_main_app")
         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: # 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.")
+            if not new_username or not new_password: st.error("Username and password are required.")
+            elif len(new_password) < 6: st.error("Password must be at least 6 characters long.")
+            elif new_password != confirm_password: st.error("Passwords do not match.")
             else:
-                user_data = UserCreate(
-                    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_data = UserCreate(username=new_username, password=new_password, email=new_email if new_email else None)
                 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}'.")
                 else:
-                    # 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.")
-
-# --- 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()
+                    st.error("Username/Email might be taken, or another error occurred. Check logs or try different credentials.")
+                    app_logger.warning(f"Signup failed for '{new_username}'. create_user_in_db returned None.")
+
+# --- Main App Logic ---
+if not st.session_state.get("authenticated_user_id"): # Not logged in
+    if not st.session_state.user_understands_disclaimer:
+        st.title(f"Welcome to {settings.APP_TITLE}")
+        st.markdown("Your AI-powered partner for advanced healthcare insights.")
+        st.markdown("---")
+        display_disclaimer()
+        st.stop() # Stop further execution until disclaimer is agreed to
+    else:
+        # Disclaimer agreed, show login/signup
+        st.title(f"Welcome to {settings.APP_TITLE}")
+        st.caption(settings.MAIN_DISCLAIMER_SHORT) # Short disclaimer always visible on login page
+        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 ---
+    # --- User is Authenticated ---
     with st.sidebar:
-        # Display Logo using get_logo_path for consistency
-        logo_path_str_sidebar = get_logo_path() # Assumes assets/logo.py is set up
+        logo_path_str_sidebar = get_logo_path()
         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) # 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}")
-            # 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}")
+                try: st.image(str(logo_file_sidebar), width=100)
+                except Exception as e: app_logger.warning(f"Could not display sidebar logo: {e}")
+        elif settings.APP_TITLE: st.sidebar.markdown(f"#### {settings.APP_TITLE}")
 
-        # Use the stored primitive username
-        username_for_display = st.session_state.get("authenticated_username", "User") # Fallback to "User"
+        username_for_display = st.session_state.get("authenticated_username", "User")
         st.sidebar.markdown(f"### Welcome, {username_for_display}!")
-        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.
+        st.sidebar.info(settings.MAIN_DISCLAIMER_SHORT) # Constant reminder
+        st.sidebar.markdown("---")
 
-        if st.sidebar.button("Logout", key="sidebar_logout_button_main_app"): # Unique key
+        if st.sidebar.button("Logout", key="sidebar_logout_button_main_app"):
             logged_out_username = st.session_state.get("authenticated_username", "UnknownUser")
             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]
+            # Clear all session state, or selectively
+            for key in list(st.session_state.keys()): # Clear most session state keys on logout
+                if key not in ['query_params']: # Keep essential Streamlit keys if needed
+                    del st.session_state[key]
+            st.session_state.user_understands_disclaimer = False # Require disclaimer agreement again
             st.success("You have been logged out.")
             st.rerun()
 
-    # 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.
+    # Main content area: Streamlit MPA handles this by showing the selected page
+    # from the `pages/` directory.
+    # If on app.py itself, give a pointer.
+    current_page = st.session_state.get('page_script_hash', None)
+    if not current_page or current_page == st.elements.utils.calc_page_script_hash('app.py'):
+         st.sidebar.success("Select a page from the navigation.")
+         st.info("Please select a page from the sidebar (e.g., Home, Consult, Reports) to begin.")
 
-    # 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.")