Update app.py

app.py

@@ -163,6 +163,7 @@ with tab1:
                 except Exception:
                     continue
 
+# ------------------ Tab 2: Visual Search ------------------
 # ------------------ Tab 2: Visual Search ------------------
 with tab2:
     st.title("🔍 Visual Reference Search")
@@ -175,19 +176,16 @@ with tab2:
 
     if st.button("Ask Assistant") and user_query:
         try:
-            # Create thread if needed
             if st.session_state.image_thread_id is None:
                 thread = client.beta.threads.create()
                 st.session_state.image_thread_id = thread.id
 
-            # Send user message
             client.beta.threads.messages.create(
                 thread_id=st.session_state.image_thread_id,
                 role="user",
                 content=user_query
             )
 
-            # Run assistant
             run = client.beta.threads.runs.create(
                 thread_id=st.session_state.image_thread_id,
                 assistant_id="asst_9v09zgizdcuuhNdcFQpRo9RO"
@@ -195,7 +193,10 @@ with tab2:
 
             with st.spinner("🔬 Searching for visual references..."):
                 while True:
-                    run_status = client.beta.threads.runs.retrieve(thread_id=st.session_state.image_thread_id, run_id=run.id)
+                    run_status = client.beta.threads.runs.retrieve(
+                        thread_id=st.session_state.image_thread_id,
+                        run_id=run.id
+                    )
                     if run_status.status in ("completed", "failed", "cancelled"):
                         break
                     time.sleep(1)
@@ -211,19 +212,28 @@ with tab2:
         except Exception as e:
             st.error(f"Error: {e}")
 
-    # Render assistant's structured image reply
+    # Render assistant response + build GitHub image URLs from filenames
     if st.session_state.image_response:
         st.markdown("### 🧠 Assistant Response")
         st.markdown(st.session_state.image_response, unsafe_allow_html=True)
 
-        # Extract and render image(s)
-        image_links = re.findall(r'https://raw\\.githubusercontent\\.com/[^\s"\')]+\\.png', st.session_state.image_response)
-        for img_url in image_links:
+        # Extract image filenames from assistant response
+        filename_matches = re.findall(r'Image Filename:\s*(.*?\\.png)', st.session_state.image_response)
+
+        for fname in filename_matches:
+            if "Autopsy" in fname:
+                folder = "Handbook Autopsy Practice Waters"
+            else:
+                folder = "JunqueirasBasicHistologyTextAtlas14thed.AnthonyMescher (2)"
+
+            encoded_folder = folder.replace(" ", "%20")
+            encoded_fname = fname.replace(" ", "%20")
+            url = f"https://raw.githubusercontent.com/AndrewLORTech/witspathologaihistotech/main/{encoded_folder}/{encoded_fname}"
+
             try:
-                r = requests.get(img_url, timeout=5)
+                r = requests.get(url, timeout=5)
                 r.raise_for_status()
                 img = Image.open(BytesIO(r.content))
-                st.image(img, caption=img_url.split('/')[-1], use_container_width=True)
-            except:
-                st.warning(f"⚠️ Failed to load: {img_url}")
-
+                st.image(img, caption=fname, use_container_width=True)
+            except Exception:
+                st.warning(f"⚠️ Failed to load: {url}")