Update app.py

app.py

@@ -31,6 +31,8 @@ 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 "search_history" not in st.session_state:
+    st.session_state.search_history = []
 
 # ------------------ Sidebar ------------------
 st.sidebar.header("⚙️ Settings")
@@ -39,6 +41,7 @@ if st.sidebar.button("🧹 Clear Chat"):
     st.session_state.thread_id = None
     st.session_state.image_url = None
     st.session_state.image_updated = False
+    st.session_state.search_history = []
     st.rerun()
 
 show_image = st.sidebar.toggle("📑 Show Page Image", value=True)
@@ -64,19 +67,46 @@ 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("#### ⚡ Prompt Templates")
-    col_btn1, col_btn2, col_btn3 = st.columns(3)
-    with col_btn1:
-        if st.button("Summarize Current Page"):
-            prompt = "Summarize the current page."
-    with col_btn2:
-        if st.button("List Contractual Obligations"):
-            prompt = "List the contractual obligations from this section."
-    with col_btn3:
-        if st.button("Find Related AS Clauses"):
-            prompt = "Which clauses reference AS 4000-1997 or similar standards?"
-
+    st.markdown("#### 🔍 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)
+
+    section_options = [
+        "Select a section...",
+        "1. Formal Instrument of Contract",
+        "2. Offer and Acceptance",
+        "3. Key Personnel",
+        "4. Contract Pricing",
+        "5. Contract Specifications",
+        "6. WHS Policies",
+        "7. Penalties and Delays",
+        "8. Dispute Resolution",
+        "9. Principal Obligations"
+    ]
+    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}"
+
+    actions = [
+        "Select an action...",
+        "List all contractual obligations",
+        "Summarize payment terms",
+        "List WHS responsibilities",
+        "Find delay-related penalties",
+        "Extract dispute resolution steps"
+    ]
+    selected_action = st.selectbox("⚙️ Common Actions", actions)
+    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:
@@ -91,30 +121,40 @@ with chat_col:
             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 == "completed":
+                    if run_status.status in ("completed", "failed", "cancelled"):
                         break
                     time.sleep(1)
 
-            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
+            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)
+            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}")