Update app.py

app.py

@@ -1,4 +1,6 @@
 import streamlit as st
+st.set_page_config(page_title="AI Pathology Assistant", layout="wide", initial_sidebar_state="collapsed")
+
 import os
 import time
 import re
@@ -22,20 +24,17 @@ def login():
     password = st.text_input("Password", type="password")
     if st.button("Login"):
         if VALID_USERS.get(email) == password:
-            st.session_state.authenticated = True
+            st.session_state["authenticated"] = True
+            st.experimental_set_query_params(logged_in="1")
             st.rerun()
         else:
             st.error("❌ Incorrect email or password.")
 
-if "authenticated" not in st.session_state:
-    st.session_state.authenticated = False
-
-if not st.session_state.authenticated:
+if not st.session_state.get("authenticated", False):
     login()
     st.stop()
 
-# ------------------ App Config ------------------
-st.set_page_config(page_title="AI Pathology Assistant", layout="wide", initial_sidebar_state="collapsed")
+# ------------------ App Title ------------------
 st.title("🧬 AI Pathology Assistant")
 
 # ------------------ Load OpenAI ------------------
@@ -49,36 +48,25 @@ client = OpenAI(api_key=OPENAI_API_KEY)
 # ------------------ Assistant Setup ------------------
 ASSISTANT_ID = "asst_jXDSjCG8LI4HEaFEcjFVq8KB"
 
-# ------------------ State ------------------
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-if "thread_id" not in st.session_state:
-    st.session_state.thread_id = None
-if "image_urls" not in st.session_state:
-    st.session_state.image_urls = []
-if "pending_prompt" not in st.session_state:
-    st.session_state.pending_prompt = None
+# ------------------ Session State Initialization ------------------
+for key in ["messages", "thread_id", "image_urls", "pending_prompt", "image_url", "image_updated"]:
+    if key not in st.session_state:
+        st.session_state[key] = [] if key.endswith("s") else None if "url" in key else False
 
 # ------------------ Tabs ------------------
 tab1, tab2 = st.tabs(["💬 Chat Assistant", "🖼️ Visual Reference Search"])
 
-# ------------------ Tab 1 ------------------
-
-# ✅ Ensure required keys exist
-if "image_url" not in st.session_state:
-    st.session_state.image_url = None
-if "image_updated" not in st.session_state:
-    st.session_state.image_updated = False
-
+# ------------------ Tab 1: Chat Assistant ------------------
 with tab1:
     with st.sidebar:
         st.header("🔧 Tools")
         if st.button("🧹 Clear Chat"):
-            st.session_state.messages = []
-            st.session_state.thread_id = None
-            st.session_state.image_url = None
-            st.session_state.image_updated = False
-            st.session_state.pending_prompt = None
+            for key in ["messages", "thread_id", "image_url", "image_updated", "pending_prompt"]:
+                st.session_state[key] = [] if key == "messages" else None if "url" in key else False
+            st.rerun()
+
+        if st.button("🚪 Logout"):
+            st.session_state.clear()
             st.rerun()
 
         show_image = st.toggle("📸 Show Images", value=True)
@@ -106,17 +94,17 @@ with tab1:
     with chat_col:
         st.markdown("### 💬 Ask a Pathology-Specific Question")
         user_input = st.chat_input("Example: What are features of squamous cell carcinoma?")
-        if user_input:
-            st.session_state.messages.append({"role": "user", "content": user_input})
-        elif st.session_state.pending_prompt:
-            st.session_state.messages.append({"role": "user", "content": st.session_state.pending_prompt})
+        if user_input or st.session_state.pending_prompt:
+            st.session_state.messages.append({
+                "role": "user",
+                "content": user_input or st.session_state.pending_prompt
+            })
             st.session_state.pending_prompt = None
 
         if st.session_state.messages and st.session_state.messages[-1]["role"] == "user":
             try:
-                if st.session_state.thread_id is None:
-                    thread = client.beta.threads.create()
-                    st.session_state.thread_id = thread.id
+                if not st.session_state.thread_id:
+                    st.session_state.thread_id = client.beta.threads.create().id
 
                 client.beta.threads.messages.create(
                     thread_id=st.session_state.thread_id,
@@ -169,6 +157,7 @@ with tab1:
             except Exception as e:
                 st.error(f"🖼️ Failed to load image: {e}")
 
+
 # ------------------ Tab 2: Visual Reference Search ------------------
 with tab2:
     ASSISTANT_ID = "asst_9v09zgizdcuuhNdcFQpRo9RO"