Update app.py

app.py

@@ -1,91 +1,127 @@
 import streamlit as st
-
-# ──────────────────────────────────────────────────────────────────────────────
-# 1) Initialize DB (imports models under the hood, then create_all)
-from models.db import init_db
-init_db()
-
-# ──────────────────────────────────────────────────────────────────────────────
-# 2) First-run admin signup (before any queries to the user table)
-from repositories.user_repo import UserRepo
 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 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
 
-repo = UserRepo(settings.database_url)
-if not repo.get_all_users():
-    st.title("🚀 Welcome to Quantum Healthcare AI")
-    st.warning("No users exist yet. Create the first admin account below:")
-    new_user = st.text_input("Username")
-    new_name = st.text_input("Full name")
-    new_pw   = st.text_input("Password", type="password")
-    if st.button("Create Admin User"):
-        if new_user and new_name and new_pw:
-            repo.add_user(new_user, new_name, new_pw)
-            st.success(f"✅ Admin `{new_user}` created! Refresh the page to log in.")
-        else:
-            st.error("All fields are required.")
-    st.stop()
+# --- Page Configuration ---
+st.set_page_config(
+    page_title=settings.APP_TITLE,
+    page_icon="⚕️", # You can use an emoji or a path to an image
+    layout="wide",
+    initial_sidebar_state="expanded"
+)
 
-# ──────────────────────────────────────────────────────────────────────────────
-# 3) Authentication
-from services.auth import authenticator, require_login
-username = require_login()
+# --- Database Initialization ---
+@st.cache_resource # Ensure this runs only once
+def init_db():
+    app_logger.info("Initializing database and tables...")
+    create_db_and_tables()
+    app_logger.info("Database initialized.")
 
-# ──────────────────────────────────────────────────────────────────────────────
-# 4) Core services & repositories
-from agent.gemini_agent import chat_with_gemini
-from clinical_nlp.umls_bioportal import lookup_umls, lookup_bioportal
-from quantum.optimizer import optimize_treatment
-from repositories.chat_repo import ChatRepo
-
-# ──────────────────────────────────────────────────────────────────────────────
-# 5) Logging & metrics
-from services.logger import logger
-from services.metrics import CHAT_COUNT, OPTIMIZE_COUNT
+init_db()
 
-# ──────────────────────────────────────────────────────────────────────────────
-# 6) UI Layout
-st.set_page_config(page_title="Quantum Healthcare AI", layout="wide")
-st.image("assets/logo.png", width=64)
-st.title(f"Hello, {username}!")
+# --- Session State Initialization ---
+if 'authenticated_user' not in st.session_state:
+    st.session_state.authenticated_user = None # Stores User object upon successful login
+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
+    st.session_state.chat_messages = []
 
-tab1, tab2 = st.tabs(["🩺 Consult", "📊 Reports"])
 
-with tab1:
-    query = st.text_area("Describe your symptoms or ask a clinical question:", height=100)
+# --- Authentication Logic ---
+def display_login_form():
+    with st.form("login_form"):
+        st.subheader("Login")
+        username = st.text_input("Username")
+        password = st.text_input("Password", type="password")
+        submit_button = st.form_submit_button("Login")
 
-    if st.button("Ask Gemini"):
-        CHAT_COUNT.inc()
-        with st.spinner("🤖 Consulting Gemini..."):
-            response = chat_with_gemini(username, query)
-            logger.info(f"[Chat] user={username} prompt={query!r}")
-        st.markdown(f"**AI Response:** {response}")
-        ChatRepo().save(user=username, prompt=query, response=response)
+        if submit_button:
+            user = authenticate_user(username, password)
+            if user:
+                st.session_state.authenticated_user = user
+                st.success(f"Welcome back, {user.username}!")
+                # Create a new chat session for the user upon login
+                with get_session_context() as db_session:
+                    new_chat_session = ChatSession(user_id=user.id, title=f"Session for {user.username}")
+                    db_session.add(new_chat_session)
+                    db_session.commit()
+                    db_session.refresh(new_chat_session)
+                    st.session_state.current_chat_session_id = new_chat_session.id
+                    st.session_state.chat_messages = [] # Clear previous messages
+                st.rerun() # Rerun to reflect login state
+            else:
+                st.error("Invalid username or password.")
 
-        with st.expander("🔎 UMLS Concept Lookup"):
-            umls_results = lookup_umls(query)
-            st.write(umls_results or "No concepts found in UMLS.")
+def display_signup_form():
+    with st.form("signup_form"):
+        st.subheader("Sign Up")
+        new_username = st.text_input("Choose a Username")
+        new_email = st.text_input("Email (Optional)")
+        new_password = st.text_input("Choose a Password", type="password")
+        confirm_password = st.text_input("Confirm Password", type="password")
+        submit_button = st.form_submit_button("Sign Up")
 
-        with st.expander("🔬 BioPortal Concept Lookup"):
-            bio_results = lookup_bioportal(query)
-            st.write(bio_results or "No matches found in BioPortal.")
+        if submit_button:
+            if not new_username or not new_password:
+                st.error("Username and password are required.")
+            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 {user.username}. Please log in.")
+                    # Optionally log them in directly:
+                    # st.session_state.authenticated_user = user
+                    # st.rerun()
+                else:
+                    st.error("Username might already be taken or an error occurred.")
 
-    if st.button("🧬 Quantum Optimize Care Plan"):
-        OPTIMIZE_COUNT.inc()
-        with st.spinner("⚛️ Running quantum-inspired optimizer..."):
-            plan = optimize_treatment(query)
-            logger.info(f"[Optimize] user={username} plan={plan}")
-        st.markdown("### Optimized Treatment Plan")
-        st.json(plan)
+# --- Main App Logic ---
+if not st.session_state.authenticated_user:
+    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:
+    # If authenticated, Streamlit automatically navigates to pages in the `pages/` directory.
+    # The content of `app.py` typically acts as the "Home" page if no `1_Home.py` exists,
+    # or it can be used for global elements like a custom sidebar if not using Streamlit's default page navigation.
 
-with tab2:
-    st.header("Generate PDF Report of Recent Chats")
-    if st.button("Download Last 5 Chats"):
-        recent = ChatRepo().get_recent(user=username, limit=5)
-        from services.pdf_report import generate_pdf
-        pdf_path = generate_pdf({"Recent Chats": recent})
-        with open(pdf_path, "rb") as f:
-            st.download_button("Download PDF", f, file_name=pdf_path)
+    # Custom Sidebar for logged-in users (Streamlit handles page navigation automatically)
+    with st.sidebar:
+        st.markdown(f"### Welcome, {st.session_state.authenticated_user.username}!")
+        st.image("assets/logo.png", width=100) # Display logo if available
+        st.markdown("---")
+        
+        if st.button("Logout"):
+            st.session_state.authenticated_user = None
+            st.session_state.current_chat_session_id = None
+            st.session_state.chat_messages = []
+            st.success("You have been logged out.")
+            st.rerun()
+    
+    # This content will show if no other page is selected, or if you don't have a 1_Home.py
+    # If you have 1_Home.py, Streamlit will show that by default after login.
+    # So, this part might be redundant if 1_Home.py exists and is the intended landing.
+    st.sidebar.success("Select a page above to get started.")
+    st.markdown(f"# {settings.APP_TITLE}")
+    st.markdown("Navigate using the sidebar to consult with the AI or view your reports.")
+    st.markdown("---")
+    st.info("This is the main application area. If you see this, ensure you have a `pages/1_Home.py` or that this `app.py` is your intended landing page after login.")
 
-# ──────────────────────────────────────────────────────────────────────────────
-st.markdown("---")
-st.caption("Powered by Gemini LLM • UMLS/BioPortal • Quantum-inspired optimization")
+app_logger.info("Streamlit app initialized and running.")