Update app.py

app.py

@@ -166,16 +166,25 @@ with tab1:
 
 
 # ------------------ 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)")
 
+with tab2:
+    ASSISTANT_ID = "asst_9v09zgizdcuuhNdcFQpRo9RO"  # Your pathology assistant
     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
+    if "image_results" not in st.session_state:
+        st.session_state.image_results = []
+    if "image_lightbox" not in st.session_state:
+        st.session_state.image_lightbox = None
+
+    # User input for assistant
+    image_input = st.chat_input("Ask for histology visual references (e.g. ovary histology, mitosis)")
+    if image_input:
+        st.session_state.image_response = None
+        st.session_state.image_results = []
+        st.session_state.image_lightbox = None
 
-    if st.button("Ask Assistant") and user_query:
         try:
             if st.session_state.image_thread_id is None:
                 thread = client.beta.threads.create()
@@ -184,7 +193,7 @@ with tab2:
             client.beta.threads.messages.create(
                 thread_id=st.session_state.image_thread_id,
                 role="user",
-                content=user_query
+                content=image_input
             )
 
             run = client.beta.threads.runs.create(
@@ -192,7 +201,7 @@ with tab2:
                 assistant_id=ASSISTANT_ID
             )
 
-            with st.spinner("🔬 Searching for visual references..."):
+            with st.spinner("🔬 Searching for histology references..."):
                 while True:
                     run_status = client.beta.threads.runs.retrieve(
                         thread_id=st.session_state.image_thread_id,
@@ -204,36 +213,48 @@ with tab2:
 
             if run_status.status == "completed":
                 messages = client.beta.threads.messages.list(thread_id=st.session_state.image_thread_id)
-                for m in reversed(messages.data):
-                    if m.role == "assistant":
-                        st.session_state.image_response = m.content[0].text.value
+                for msg in reversed(messages.data):
+                    if msg.role == "assistant":
+                        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]
                         break
-            else:
-                st.error("❌ Assistant failed to return an image match.")
         except Exception as e:
-            st.error(f"Error: {e}")
+            st.error(f"❌ Visual Assistant Error: {e}")
 
-    # 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)
+    # Display split UI (image results right, response text left)
+    text_col, image_col = st.columns([2, 1])
 
-        # 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", " ")
-        )
+    with text_col:
+        if st.session_state.image_response:
+            st.markdown("### 🧠 Assistant Response")
+            st.markdown(st.session_state.image_response, unsafe_allow_html=True)
 
-        if url_matches:
+    with image_col:
+        if st.session_state.image_results:
             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"⚠️ Could not load image: {url}")
-                    st.error(f"🛑 Error: {str(e)}")
+            for i, item in enumerate(st.session_state.image_results):
+                image_url = item.get("image") or item.get("Image URL") or item.get("ImageURL")
+                if image_url:
+                    try:
+                        st.image(image_url, caption=image_url.split("/")[-1], use_container_width=True)
+                        if st.button("🖼️ View Full Image", key=f"img_{i}"):
+                            st.session_state.image_lightbox = image_url
+                    except Exception as e:
+                        st.warning(f"⚠️ Could not load: {image_url}")
+                        st.error(f"{e}")
         else:
-            st.info("ℹ️ No valid image URLs found in the assistant's response.")
+            st.info("ℹ️ No image references found yet.")
+
+    if st.session_state.image_lightbox:
+        st.image(st.session_state.image_lightbox, caption="🔍 Full Image View", use_container_width=True)
+        if st.button("❌ Close Preview"):
+            st.session_state.image_lightbox = None
+            st.rerun()