Update app.py

app.py

@@ -33,6 +33,8 @@ if "image_updated" not in st.session_state:
     st.session_state.image_updated = False
 if "search_history" not in st.session_state:
     st.session_state.search_history = []
+if "run_in_progress" not in st.session_state:
+    st.session_state.run_in_progress = False
 
 # ------------------ Sidebar ------------------
 st.sidebar.header("⚙️ Settings")
@@ -42,34 +44,22 @@ if st.sidebar.button("🧹 Clear Chat"):
     st.session_state.image_url = None
     st.session_state.image_updated = False
     st.session_state.search_history = []
+    st.session_state.run_in_progress = False
     st.rerun()
 
 show_image = st.sidebar.toggle("📑 Show Page Image", value=True)
 
-# ------------------ Layout ------------------
-image_col, chat_col = st.columns([1, 2])
+# ------------------ Center Chat UI ------------------
+st.markdown("### 🧠 Ask a Document-Specific Question")
+prompt = st.chat_input("Example: What is the defects liability period?")
 
-# ------------------ Image Viewer ------------------
-with image_col:
-    if show_image and st.session_state.image_url:
-        with st.spinner("Loading document preview..."):
-            try:
-                response = requests.get(st.session_state.image_url)
-                response.raise_for_status()
-                img = Image.open(BytesIO(response.content))
-                st.image(img, caption="📄 OCR Page Image", use_container_width=True)
-                st.session_state.image_updated = False
-            except Exception as e:
-                st.error(f"❗ Failed to load image: {e}")
+# -- Chat Output
+for msg in st.session_state.messages:
+    with st.chat_message(msg["role"]):
+        st.markdown(msg["content"], unsafe_allow_html=True)
 
-# ------------------ Chat Interface ------------------
-with chat_col:
-    st.markdown("### 🧠 Ask a Document-Specific Question")
-    prompt = st.chat_input("Example: What is the defects liability period?")
-
-    # -- Smart Query Controls
-    st.markdown("<hr style='margin:0.5em 0;'>", unsafe_allow_html=True)
-    st.markdown("#### 🔍 Navigate or Query Document")
+# -- Expandable Tools Below Chat
+with st.expander("🔍 Navigate or Query Document"):
     col1, col2 = st.columns([1.2, 1])
     with col1:
         keyword = st.text_input("Jump to clause keyword or term", placeholder="e.g. defects, payment, WHS")
@@ -106,55 +96,69 @@ with chat_col:
     if selected_action and selected_action != actions[0]:
         prompt = selected_action
 
-    # -- Send prompt to assistant
-    if prompt:
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        try:
-            if st.session_state.thread_id is None:
-                thread = client.beta.threads.create()
-                st.session_state.thread_id = thread.id
-
-            client.beta.threads.messages.create(thread_id=st.session_state.thread_id, role="user", content=prompt)
-
-            run = client.beta.threads.runs.create(thread_id=st.session_state.thread_id, assistant_id=ASSISTANT_ID)
-
-            with st.spinner("🤖 Parsing and responding with referenced content..."):
-                while True:
-                    run_status = client.beta.threads.runs.retrieve(thread_id=st.session_state.thread_id, run_id=run.id)
-                    if run_status.status in ("completed", "failed", "cancelled"):
-                        break
-                    time.sleep(1)
-
-            if run_status.status != "completed":
-                st.error(f"⚠️ Assistant failed: {run_status.status}")
-            else:
-                messages = client.beta.threads.messages.list(thread_id=st.session_state.thread_id)
-                for message in reversed(messages.data):
-                    if message.role == "assistant":
-                        assistant_message = message.content[0].text.value
-                        st.session_state.messages.append({"role": "assistant", "content": assistant_message})
-                        break
-
-                match = re.search(r'Document Reference:\s+(.+?),\s+Page\s+(\d+)', assistant_message)
-                if match:
-                    doc_name = match.group(1).strip()
-                    page = int(match.group(2))
-                    page_str = f"{page:04d}"
-                    image_url = f"https://raw.githubusercontent.com/AndrewLORTech/c2ozschlaegerforrestdale/main/{doc_name}/{doc_name}_page_{page_str}.png"
-                    st.session_state.image_url = image_url
-                    st.session_state.image_updated = True
-
-            st.rerun()
-        except Exception as e:
-            st.error(f"❌ Error: {e}")
-
-    # -- Display Chat History
-    for msg in st.session_state.messages:
-        with st.chat_message(msg["role"]):
-            st.markdown(msg["content"], unsafe_allow_html=True)
-
     if st.session_state.search_history:
         st.markdown("---")
         st.markdown("#### 🔁 Recent Searches")
         for term in reversed(st.session_state.search_history[-5:]):
             st.markdown(f"- {term}")
+
+# ------------------ Image Viewer ------------------
+with st.expander("📄 Page Image (Toggle)"):
+    if show_image and st.session_state.image_url:
+        with st.spinner("Loading document preview..."):
+            try:
+                response = requests.get(st.session_state.image_url)
+                response.raise_for_status()
+                img = Image.open(BytesIO(response.content))
+                st.image(img, caption="📄 OCR Page Image", use_container_width=True)
+                st.session_state.image_updated = False
+            except Exception as e:
+                st.error(f"❗ Failed to load image: {e}")
+
+# ------------------ Prompt Handling ------------------
+if prompt and not st.session_state.run_in_progress:
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.session_state.run_in_progress = True
+
+    try:
+        if st.session_state.thread_id is None:
+            thread = client.beta.threads.create()
+            st.session_state.thread_id = thread.id
+
+        client.beta.threads.messages.create(thread_id=st.session_state.thread_id, role="user", content=prompt)
+
+        run = client.beta.threads.runs.create(thread_id=st.session_state.thread_id, assistant_id=ASSISTANT_ID)
+
+        with st.spinner("🤖 Parsing and responding with referenced content..."):
+            while True:
+                run_status = client.beta.threads.runs.retrieve(thread_id=st.session_state.thread_id, run_id=run.id)
+                if run_status.status in ("completed", "failed", "cancelled"):
+                    break
+                time.sleep(1)
+
+        if run_status.status != "completed":
+            st.error(f"⚠️ Assistant failed: {run_status.status}")
+        else:
+            messages = client.beta.threads.messages.list(thread_id=st.session_state.thread_id)
+            for message in reversed(messages.data):
+                if message.role == "assistant":
+                    assistant_message = message.content[0].text.value
+                    st.session_state.messages.append({"role": "assistant", "content": assistant_message})
+                    break
+
+            match = re.search(r'Document Reference:\s+(.+?),\s+Page\s+(\d+)', assistant_message)
+            if match:
+                doc_name = match.group(1).strip()
+                page = int(match.group(2))
+                page_str = f"{page:04d}"
+                image_url = f"https://raw.githubusercontent.com/AndrewLORTech/c2ozschlaegerforrestdale/main/{doc_name}/{doc_name}_page_{page_str}.png"
+                st.session_state.image_url = image_url
+                st.session_state.image_updated = True
+
+        st.session_state.run_in_progress = False
+        st.rerun()
+
+    except Exception as e:
+        st.session_state.run_in_progress = False
+        st.error(f"❌ Error: {e}")
+