Update app.py

app.py

@@ -166,9 +166,10 @@ with tab1:
 
 
 # ------------------ Tab 2: Visual Reference Search ------------------
-
 with tab2:
-    ASSISTANT_ID = "asst_9v09zgizdcuuhNdcFQpRo9RO"  # Your pathology assistant
+    ASSISTANT_ID = "asst_9v09zgizdcuuhNdcFQpRo9RO"  # Your visual reference assistant ID
+
+    # Session 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:
@@ -178,7 +179,7 @@ with tab2:
     if "image_lightbox" not in st.session_state:
         st.session_state.image_lightbox = None
 
-    # User input for assistant
+    # User input
     image_input = st.chat_input("Ask for histology visual references (e.g. ovary histology, mitosis)")
     if image_input:
         st.session_state.image_response = None
@@ -218,18 +219,24 @@ with tab2:
                         response_text = msg.content[0].text.value
                         st.session_state.image_response = response_text
 
-                        # Try to extract structured data (JSON-like block) if present
-                        try:
-                            st.session_state.image_results = json.loads(response_text.strip("`json "))
-                        except:
-                            # Fallback to manually extract all GitHub image URLs
-                            urls = re.findall(r'https://raw\.githubusercontent\.com/[^\s)>\]]+\.png', response_text.replace("\n", " "))
-                            st.session_state.image_results = [{"image": url} for url in urls]
+                        # 🔍 Robustly extract "Image URL" entries
+                        lines = response_text.splitlines()
+                        image_urls = []
+                        for i, line in enumerate(lines):
+                            if "Image URL:" in line:
+                                url = line.split("Image URL:")[-1].strip()
+                                if not url.startswith("http") and i + 1 < len(lines):
+                                    next_line = lines[i + 1].strip()
+                                    if next_line.startswith("http"):
+                                        url = next_line
+                                if url.startswith("http"):
+                                    image_urls.append(url)
+                        st.session_state.image_results = [{"image": url} for url in image_urls]
                         break
         except Exception as e:
             st.error(f"❌ Visual Assistant Error: {e}")
 
-    # Display split UI (image results right, response text left)
+    # Layout split
     text_col, image_col = st.columns([2, 1])
 
     with text_col:
@@ -241,7 +248,7 @@ with tab2:
         if st.session_state.image_results:
             st.markdown("### 🖼️ Image Preview(s)")
             for i, item in enumerate(st.session_state.image_results):
-                image_url = item.get("image") or item.get("Image URL") or item.get("ImageURL")
+                image_url = item.get("image")
                 if image_url:
                     try:
                         st.image(image_url, caption=image_url.split("/")[-1], use_container_width=True)