Update app.py

app.py

@@ -163,18 +163,23 @@ with tab1:
                 except Exception:
                     continue
                     
-from urllib.parse import quote  # ✅ Make sure this is at the top of your file
+from PIL import Image
+import requests
+import re
+from io import BytesIO
 
 # ------------------ Tab 2: Visual Reference Search ------------------
 with tab2:
     st.title("🔍 Visual Reference Search")
     user_query = st.text_input("Enter keyword to search images (e.g. ovary, thyroid, mitosis)")
 
+    # Initialize thread and response states
     if "image_thread_id" not in st.session_state:
         st.session_state.image_thread_id = None
     if "image_response" not in st.session_state:
         st.session_state.image_response = None
 
+    # Submit user query to assistant
     if st.button("Ask Assistant") and user_query:
         try:
             if st.session_state.image_thread_id is None:
@@ -189,7 +194,7 @@ with tab2:
 
             run = client.beta.threads.runs.create(
                 thread_id=st.session_state.image_thread_id,
-                assistant_id="asst_9v09zgizdcuuhNdcFQpRo9RO"
+                assistant_id=ASSISTANT_ID  # Reuse your assistant ID
             )
 
             with st.spinner("🔬 Searching for visual references..."):
@@ -213,17 +218,24 @@ with tab2:
         except Exception as e:
             st.error(f"Error: {e}")
 
-    # Render assistant response and load each image using the actual URL
+    # Display assistant text response
     if st.session_state.image_response:
         st.markdown("### 🧠 Assistant Response")
         st.markdown(st.session_state.image_response, unsafe_allow_html=True)
 
-        # Extract valid GitHub URLs from assistant response
+        # Extract GitHub image URLs
         url_matches = re.findall(r'Image URL:\s*(https?://[^\s]+)', st.session_state.image_response)
 
-        for url in url_matches:
-            try:
-                st.image(url, caption=url.split("/")[-1], use_container_width=True)
-            except Exception as e:
-                st.warning(f"⚠️ Could not load image: {url}")
-                st.error(f"🛑 Error: {str(e)}")
+        if url_matches:
+            st.markdown("### 🖼️ Image Preview(s)")
+            for url in url_matches:
+                try:
+                    response = requests.get(url, timeout=5)
+                    response.raise_for_status()
+                    img = Image.open(BytesIO(response.content))
+                    st.image(img, caption=url.split("/")[-1], use_container_width=True)
+                except Exception as e:
+                    st.warning(f"⚠️ Failed to load image: {url}")
+                    st.error(f"🛑 Error: {e}")
+        else:
+            st.info("ℹ️ No valid image URLs found in the assistant's response.")