Update app.py

app.py

@@ -49,24 +49,16 @@ if st.sidebar.button("🧹 Clear Chat"):
 
 show_image = st.sidebar.toggle("📑 Show Page Image", value=True)
 
-# ------------------ Center Chat UI ------------------
-st.markdown("### 🧠 Ask a Document-Specific Question")
-prompt = st.chat_input("Example: What is the defects liability period?")
-
-# -- Chat Output
-for msg in st.session_state.messages:
-    with st.chat_message(msg["role"]):
-        st.markdown(msg["content"], unsafe_allow_html=True)
-
-# -- 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")
-    with col2:
-        if st.button("🔎 Search Keyword") and keyword:
-            prompt = f"Find clauses or references related to: {keyword}"
-            st.session_state.search_history.append(keyword)
+# ------------------ Horizontal Layout ------------------
+left, center = st.columns([1, 2])
+
+# ------------------ Expandable Left Section ------------------
+with left.expander("📂 Navigate or Query Document", expanded=False):
+    keyword = st.text_input("Jump to clause keyword or term", placeholder="e.g. defects, payment, WHS")
+    if st.button("🔎 Search Keyword") and keyword:
+        st.session_state.search_history.append(keyword)
+        if not st.session_state.run_in_progress:
+            st.session_state.messages.append({"role": "user", "content": f"Find clauses or references related to: {keyword}"})
 
     section_options = [
         "Select a section...",
@@ -82,7 +74,7 @@ with st.expander("🔍 Navigate or Query Document"):
     ]
     selected_section = st.selectbox("📄 Browse Sections", section_options)
     if selected_section and selected_section != section_options[0]:
-        prompt = f"Summarize or list key points from section: {selected_section}"
+        st.session_state.messages.append({"role": "user", "content": f"Summarize or list key points from section: {selected_section}"})
 
     actions = [
         "Select an action...",
@@ -94,7 +86,18 @@ with st.expander("🔍 Navigate or Query Document"):
     ]
     selected_action = st.selectbox("⚙️ Common Actions", actions)
     if selected_action and selected_action != actions[0]:
-        prompt = selected_action
+        st.session_state.messages.append({"role": "user", "content": selected_action})
+
+    if show_image and st.session_state.image_url:
+        st.markdown("---")
+        st.markdown("### 📄 Page 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)
+        except Exception as e:
+            st.error(f"❗ Failed to load image: {e}")
 
     if st.session_state.search_history:
         st.markdown("---")
@@ -102,63 +105,57 @@ with st.expander("🔍 Navigate or Query Document"):
         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}")
-
+# ------------------ Chat Section ------------------
+with center:
+    st.markdown("### 🧠 Ask a Document-Specific Question")
+    prompt = st.chat_input("Example: What is the defects liability period?")
+
+    for msg in st.session_state.messages:
+        with st.chat_message(msg["role"]):
+            st.markdown(msg["content"], unsafe_allow_html=True)
+
+    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}")