Update app.py

app.py

@@ -7,6 +7,9 @@ from PIL import Image
 from io import BytesIO
 from openai import OpenAI
 
+# ------------------ Page Config ------------------
+st.set_page_config(page_title="AI Pathology Assistant", layout="wide", initial_sidebar_state="collapsed")
+
 # ------------------ Authentication ------------------
 VALID_USERS = {
     "[email protected]": "Pass.123",
@@ -33,10 +36,6 @@ if not st.session_state.authenticated:
     login()
     st.stop()
 
-# ------------------ App Config ------------------
-st.set_page_config(page_title="AI Pathology Assistant", layout="wide", initial_sidebar_state="collapsed")
-st.title("🧬 AI Pathology Assistant")
-
 # ------------------ Load OpenAI ------------------
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
 if not OPENAI_API_KEY:
@@ -48,30 +47,26 @@ 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
-
-# ------------------ Tabs ------------------
+# ------------------ State Initialization ------------------
+for key in ["messages", "thread_id", "image_urls", "pending_prompt"]:
+    if key not in st.session_state:
+        st.session_state[key] = [] if key.endswith("s") else None
+
+# ------------------ UI Tabs ------------------
 tab1, tab2 = st.tabs(["💬 Chat Assistant", "🖼️ Visual Reference Search"])
 
+# ------------------ Tab 1: Chat Assistant ------------------
 with tab1:
-    # Sidebar
     with st.sidebar:
         st.header("🧪 Pathology Tools")
-        
+
         if st.button("🧹 Clear Chat"):
             for k in ["messages", "thread_id", "image_urls", "pending_prompt"]:
                 st.session_state[k] = [] if k.endswith("s") else None
             st.rerun()
 
         show_image = st.toggle("📸 Show Images", value=True)
+
         keyword = st.text_input("Keyword Search", placeholder="e.g. mitosis, carcinoma")
         if st.button("🔎 Search") and keyword:
             st.session_state.pending_prompt = f"Find clauses or references related to: {keyword}"
@@ -97,14 +92,12 @@ with tab1:
         st.markdown("### 💬 Ask a Pathology-Specific Question")
         user_input = st.chat_input("Example: What are features of squamous cell carcinoma?")
 
-        # Append input or pending prompt
         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})
             st.session_state.pending_prompt = None
 
-        # Trigger assistant run
         if st.session_state.messages and st.session_state.messages[-1]["role"] == "user":
             try:
                 if not st.session_state.thread_id:
@@ -133,13 +126,10 @@ with tab1:
                     for m in reversed(responses):
                         if m.role == "assistant":
                             reply = m.content[0].text.value.strip()
-
-                            # De-dupe reply
                             if not any(reply in msg["content"] or msg["content"] in reply
                                        for msg in st.session_state.messages if msg["role"] == "assistant"):
                                 st.session_state.messages.append({"role": "assistant", "content": reply})
 
-                                # Image detection
                                 images = re.findall(
                                     r'https://raw\.githubusercontent\.com/AndrewLORTech/witspathologai/main/[^\s\n"]+\.png',
                                     reply
@@ -152,12 +142,10 @@ with tab1:
             except Exception as e:
                 st.error(f"❌ Error: {e}")
 
-        # Display conversation
         for msg in st.session_state.messages:
             with st.chat_message(msg["role"]):
                 st.markdown(msg["content"], unsafe_allow_html=True)
 
-        # Optional follow-ups
         if st.session_state.messages and st.session_state.messages[-1]["role"] == "assistant":
             last = st.session_state.messages[-1]["content"]
             if "Some Possible Questions:" in last:
@@ -179,6 +167,7 @@ with tab1:
                 except Exception:
                     st.warning(f"⚠️ Failed to load image: {url}")
 
+
 # ------------------ Tab 2: Visual Reference Search ------------------
 with tab2:
     ASSISTANT_ID = "asst_9v09zgizdcuuhNdcFQpRo9RO"