# /home/user/app/app.py import streamlit as st from pathlib import Path from sqlmodel import select # <--- IMPORT SELECT FOR SQLMODEL QUERIES from config.settings import settings from models import ( create_db_and_tables, get_session_context, # Your SQLModel session context manager User, ChatMessage, ChatSession ) 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 # --- Page Configuration --- (Assuming this is already correct) st.set_page_config( page_title=settings.APP_TITLE, page_icon="⚕️", layout="wide", initial_sidebar_state="expanded" ) # --- Database Initialization --- (Assuming this is already correct) @st.cache_resource 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) 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"): 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") 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: st.success(f"Welcome back, {username_input}!") app_logger.info(f"User {username_input} authenticated successfully.") try: with get_session_context() as db_session: # db_session is a SQLModel Session # --- SQLMODEL QUERY --- 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.") st.session_state.authenticated_user_id = None st.session_state.authenticated_username = None st.rerun() return 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})") 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 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() except Exception as e: 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}") 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}") # 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") 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: 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 ) user = create_user_in_db(user_data) if user: st.success(f"Account created for {new_username}. Please log in.") 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}") if not st.session_state.get("authenticated_user_id"): 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: with st.sidebar: 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) 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: st.sidebar.markdown(f"#### {settings.APP_TITLE}") username_for_display = st.session_state.get("authenticated_username", "User") st.sidebar.markdown(f"### Welcome, {username_for_display}!") st.sidebar.markdown("---") if st.sidebar.button("Logout", key="sidebar_logout_button"): logged_out_username = st.session_state.get("authenticated_username", "UnknownUser") app_logger.info(f"User {logged_out_username} logging out.") 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.") st.rerun() app_logger.info(f"Streamlit app '{settings.APP_TITLE}' initialized and running.")