Update app.py

app.py

@@ -55,18 +55,31 @@ tab1, tab2 = st.tabs(["📑 Contract", "📁 Technical"])
 with tab1:
     ASSISTANT_ID = "asst_KsQRedoJUnEeStzfox1o06lO"
 
-    for key in ["messages", "thread_id", "image_url", "image_updated", "pending_prompt"]:
-        if key not in st.session_state:
-            st.session_state[key] = None if "url" in key else []
-
+    # Correct session state initialization
+    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_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
+    if "pending_prompt" not in st.session_state:
+        st.session_state.pending_prompt = None
+
+    # Sidebar tools
     with st.sidebar:
         st.header("ℹ️ Contract Tools")
         if st.button("🧹 Clear Chat"):
-            for k in ["messages", "thread_id", "image_url", "image_updated", "pending_prompt"]:
-                st.session_state[k] = None if "url" in k else []
+            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
             st.rerun()
 
         show_image = st.toggle("📑 Show Page Image", value=True)
+
         keyword = st.text_input("Search by Keyword", placeholder="e.g. defects, WHS, delay")
         if st.button("🔎 Search Keyword") and keyword:
             st.session_state.pending_prompt = f"Find clauses or references related to: {keyword}"
@@ -99,6 +112,7 @@ with tab1:
         if action != actions[0]:
             st.session_state.pending_prompt = action
 
+    # Chat + Image layout
     chat_col, image_col = st.columns([2, 1])
     with chat_col:
         st.markdown("### 🧠 Ask a Document-Specific Question")