Update app.py

app.py

@@ -140,28 +140,23 @@ with tab1:
 # ------------------ Tab 2: Contract Queries ------------------
 with tab2:
     ASSISTANT_ID = "asst_Qcl6k4bYxK2UiOfBfqsC6LZv"
-    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
 
+    # State initialization
+    for key in ["messages", "thread_id", "image_url", "image_updated", "pending_prompt"]:
+        if key not in st.session_state:
+            st.session_state[key] = [] if key == "messages" else None if "id" in key or key.endswith("_url") else False
+
+    # Sidebar
     with st.sidebar:
         st.header("ℹ️ Contract Tools")
+
         if st.button("🧹 Clear Chat"):
-            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
+            for key in ["messages", "thread_id", "image_url", "image_updated", "pending_prompt"]:
+                st.session_state[key] = [] if key == "messages" else None if "id" in key or key.endswith("_url") else False
             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}"
@@ -195,6 +190,7 @@ with tab2:
             st.rerun()
 
     chat_col, image_col = st.columns([2, 1])
+
     with chat_col:
         st.markdown("### 🧠 Ask a Document-Specific Question")
         user_input = st.chat_input("Example: What is the defects liability period?")
@@ -206,7 +202,7 @@ with tab2:
 
         if st.session_state.messages and st.session_state.messages[-1]["role"] == "user":
             try:
-                if st.session_state.thread_id is None:
+                if not st.session_state.thread_id:
                     thread = client.beta.threads.create()
                     st.session_state.thread_id = thread.id
 
@@ -223,7 +219,10 @@ with tab2:
 
                 with st.spinner("🤖 Thinking..."):
                     while True:
-                        status = client.beta.threads.runs.retrieve(thread_id=st.session_state.thread_id, run_id=run.id)
+                        status = client.beta.threads.runs.retrieve(
+                            thread_id=st.session_state.thread_id,
+                            run_id=run.id
+                        )
                         if status.status in ("completed", "failed", "cancelled"):
                             break
                         time.sleep(1)
@@ -232,16 +231,26 @@ with tab2:
                     messages = client.beta.threads.messages.list(thread_id=st.session_state.thread_id)
                     for m in reversed(messages.data):
                         if m.role == "assistant":
-                            reply = m.content[0].text.value.strip()
+                            full_reply = m.content[0].text.value.strip()
 
-                            # 🧼 Remove citation artifacts like [8:1:source]
-                            import re
-                            reply = re.sub(r"\[\d+:\d+:source\]", "", reply)
+                            # ✂️ Strip "Some Possible Questions"
+                            match = re.search(r"Some Possible Questions:\s*(.*?)(?=\n{2,}|Other References:|\Z)", full_reply, re.DOTALL)
+                            question_block = None
+                            if match:
+                                question_block = match.group(1).strip()
+                                full_reply = full_reply.replace(f"Some Possible Questions:\n{question_block}", "").strip()
+
+                            # 🧼 Remove all source tags like [8:0†source], [source], [8:1:source]
+                            cleaned_reply = re.sub(r"\[\d+:\d+†[^\]]+\]", "", full_reply)
+                            cleaned_reply = re.sub(r"\[\d+:\d+:source\]", "", cleaned_reply)
+                            cleaned_reply = re.sub(r"\[source\]", "", cleaned_reply)
 
-                            if not any(reply in msg["content"] for msg in st.session_state.messages if msg["role"] == "assistant"):
-                                st.session_state.messages.append({"role": "assistant", "content": reply})
+                            # Save reply and extract image if unique
+                            if not any(cleaned_reply in msg["content"] for msg in st.session_state.messages if msg["role"] == "assistant"):
+                                st.session_state.messages.append({"role": "assistant", "content": cleaned_reply})
 
-                                match = re.search(r'📄\s*(.*?)\.txt\s*🔢.*?Page\s*(\d+)', reply)
+                                # 🔍 Extract image reference
+                                match = re.search(r'📄\s*(.*?)\.txt\s*🔢.*?Page\s*(\d+)', cleaned_reply)
                                 if match:
                                     doc, page = match.group(1).strip(), int(match.group(2))
                                     folder = quote(doc)
@@ -259,14 +268,12 @@ with tab2:
             with st.chat_message(msg["role"]):
                 st.markdown(msg["content"], unsafe_allow_html=True)
 
-        if st.session_state.messages and st.session_state.messages[-1]["role"] == "assistant":
-            latest_reply = st.session_state.messages[-1]["content"]
-            match = re.search(r"###\s*\*\*Some Possible Questions:\*\*(.*?)###", latest_reply + "###", re.DOTALL)
-            questions = []
-            if match:
-                block = match.group(1)
-                questions = [line.strip(" -•") for line in block.splitlines() if line.strip().startswith(("-", "•"))]
-
+        # ✅ Follow-up questions
+        if 'question_block' in locals() and question_block:
+            questions = [
+                line.strip(" -•") for line in question_block.splitlines()
+                if line.strip().startswith(("-", "•")) and line.strip().endswith("?")
+            ]
             if questions:
                 st.markdown("#### 💡 Follow-Up Suggestions")
                 for q in questions:
@@ -284,6 +291,7 @@ with tab2:
             except Exception as e:
                 st.error(f"🖼️ Failed to load image: {e}")
 
+
 # ------------------ Technical Tab ------------------
 with tab3:
     ASSISTANT_ID = "asst_DjvuWBc7tCvMbAhY7n1em4BZ"