Update app.py

app.py

@@ -163,23 +163,18 @@ with tab1:
                 except Exception:
                     continue
                     
-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:
@@ -194,7 +189,7 @@ with tab2:
 
             run = client.beta.threads.runs.create(
                 thread_id=st.session_state.image_thread_id,
-                assistant_id=ASSISTANT_ID  # Reuse your assistant ID
+                assistant_id=ASSISTANT_ID
             )
 
             with st.spinner("🔬 Searching for visual references..."):
@@ -218,13 +213,16 @@ with tab2:
         except Exception as e:
             st.error(f"Error: {e}")
 
-    # Display assistant text response
+    # Display assistant response and render images
     if st.session_state.image_response:
         st.markdown("### 🧠 Assistant Response")
         st.markdown(st.session_state.image_response, unsafe_allow_html=True)
 
-        # Extract GitHub image URLs
-        url_matches = re.findall(r'Image URL:\s*(https?://[^\s]+)', st.session_state.image_response)
+        # More robust regex to extract URLs even with encoded characters and line breaks
+        url_matches = re.findall(
+            r'https://raw\.githubusercontent\.com/[^\s)>\]]+\.png',
+            st.session_state.image_response.replace("\n", " ")
+        )
 
         if url_matches:
             st.markdown("### 🖼️ Image Preview(s)")
@@ -235,7 +233,7 @@ with tab2:
                     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}")
+                    st.warning(f"⚠️ Could not load image: {url}")
+                    st.error(f"🛑 Error: {str(e)}")
         else:
             st.info("ℹ️ No valid image URLs found in the assistant's response.")