Update app.py

app.py

@@ -27,6 +27,8 @@ 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
 
 # ------------------ Sidebar Controls ------------------
 st.sidebar.header("🔧 Settings")
@@ -34,6 +36,7 @@ if st.sidebar.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.rerun()
 
 show_image = st.sidebar.checkbox("📖 Show Document Image", value=True)
@@ -45,6 +48,7 @@ col1, col2 = st.columns([1, 2])  # Adjust ratio as needed
 with col1:
     if show_image and st.session_state.image_url:
         st.image(st.session_state.image_url, caption="📑 Extracted Page", use_container_width=True)
+        st.session_state.image_updated = False  # Reset flag after rendering
 
 # ------------------ Chat Panel (Right) ------------------
 with col2:
@@ -106,6 +110,8 @@ with col2:
             )
             if image_match:
                 st.session_state.image_url = image_match.group(0)
+                st.session_state.image_updated = True
+                st.rerun()  # Trigger rerun to refresh image display
 
         except Exception as e:
             st.error(f"❌ Error: {str(e)}")