Update app.py

app.py

@@ -54,7 +54,7 @@ show_image = st.sidebar.checkbox("📖 Show Document Image", value=True)
 
 # ------------------ Assistant Query Function ------------------
 def query_assistant(prompt):
-    st.session_state.messages.insert(0, {"role": "user", "content": prompt})  # Insert at top
+    st.session_state.messages.insert(0, {"role": "user", "content": prompt})
 
     try:
         if st.session_state.thread_id is None:
@@ -84,21 +84,22 @@ def query_assistant(prompt):
                     break
                 time.sleep(1)
 
-        messages = client.beta.threads.messages.list(thread_id=thread_id)
-        for message in reversed(messages.data):
-            if message.role == "assistant":
-                assistant_message = message.content[0].text.value
-                st.session_state.messages.insert(0, {"role": "assistant", "content": assistant_message})
-
-                # Extract GitHub image URL if available
-                image_match = re.search(
-                    r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png',
-                    assistant_message
-                )
-                if image_match:
-                    st.session_state.image_url = image_match.group(0)
-                    st.session_state.image_updated = True
-                return assistant_message
+        # ✅ Only fetch the latest assistant message to avoid infinite loop
+        messages = client.beta.threads.messages.list(thread_id=thread_id, limit=1)
+        latest_message = messages.data[0]
+        if latest_message.role == "assistant":
+            assistant_message = latest_message.content[0].text.value
+            st.session_state.messages.insert(0, {"role": "assistant", "content": assistant_message})
+
+            # Extract GitHub image URL if available
+            image_match = re.search(
+                r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png',
+                assistant_message
+            )
+            if image_match:
+                st.session_state.image_url = image_match.group(0)
+                st.session_state.image_updated = True
+            return assistant_message
 
     except Exception as e:
         st.error(f"❌ Error: {str(e)}")
@@ -116,7 +117,7 @@ with center:
         prompt = st.text_input("💡 Ask a question about the document:", key="chat_input")
         if prompt:
             query_assistant(prompt)
-            st.rerun()  # ✅ Updated to current Streamlit standard
+            st.rerun()  # Rerun to refresh with new message
 
     # Show messages: latest at top
     for message in st.session_state.messages: