Update app.py

app.py

@@ -164,9 +164,11 @@ with tab2:
 
     for key in ["image_thread_id", "image_response", "image_results", "image_lightbox"]:
         if key not in st.session_state:
-            st.session_state[key] = None if "id" in key or "lightbox" in key else []
+            st.session_state[key] = None if "id" in key or "response" in key else []
+
+    st.markdown("### 🔍 Search Histology Visual References")
+    image_input = st.chat_input("e.g. mitosis, seminiferous tubule, cell cycle")
 
-    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 = []
@@ -187,85 +189,68 @@ with tab2:
                 assistant_id=ASSISTANT_ID
             )
 
-            with st.spinner("🔬 Searching for histology references..."):
+            with st.spinner("🧬 Retrieving images..."):
                 while True:
-                    run_status = client.beta.threads.runs.retrieve(
+                    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"):
+                    if status.status in ("completed", "failed", "cancelled"):
                         break
                     time.sleep(1)
 
-            if run_status.status == "completed":
+            if status.status == "completed":
                 messages = client.beta.threads.messages.list(thread_id=st.session_state.image_thread_id)
                 for msg in reversed(messages.data):
                     if msg.role == "assistant":
                         response_text = msg.content[0].text.value
                         st.session_state.image_response = response_text
 
-                        # ✅ Image metadata parser
+                        # Extract URLs and details for grid card display
                         lines = response_text.splitlines()
-                        image_entries = []
-                        current = {"title": "", "description": "", "url": ""}
-
+                        block = {}
+                        all_cards = []
                         for line in lines:
-                            if line.startswith("### 🖼️"):
-                                current = {"title": line.replace("### 🖼️", "").strip(), "description": "", "url": ""}
-                            elif line.strip().startswith("- **Description:**"):
-                                current["description"] = line.split("**Description:**")[-1].strip()
+                            if line.strip().startswith("### 🖼️"):
+                                if block: all_cards.append(block)
+                                block = {"title": line.replace("### 🖼️", "").strip()}
+                            elif "Page:" in line:
+                                block["page"] = line.split("Page:")[1].strip()
+                            elif "Description:" in line:
+                                block["desc"] = line.split("Description:")[1].strip()
                             elif "Image URL:" in line:
-                                match = re.search(r'(https://[^\s]+)', line)
-                                if match:
-                                    current["url"] = match.group(1)
-                                    image_entries.append(current)
-                            elif line.strip().startswith("https://") and current["url"] == "":
-                                current["url"] = line.strip()
-                                image_entries.append(current)
-
-                        st.session_state.image_results = image_entries
-                        if image_entries and not st.session_state.image_lightbox:
-                            st.session_state.image_lightbox = image_entries[0]["url"]
+                                block["url"] = line.split("Image URL:")[1].strip()
+                        if block: all_cards.append(block)
+                        st.session_state.image_results = all_cards
                         break
-        except Exception as e:
-            st.error(f"❌ Visual Assistant Error: {e}")
-
-    text_col, image_col = st.columns([2, 1])
-    with text_col:
-        if st.session_state.image_response:
-            st.markdown("### 🧠 Assistant Response")
-            st.markdown(st.session_state.image_response, unsafe_allow_html=True)
 
-    with image_col:
-        if st.session_state.image_results:
-            st.markdown("### 🖼️ Image Preview(s)")
-            cols = st.columns(4)
-            for idx, entry in enumerate(st.session_state.image_results):
-                col = cols[idx % 4]
-                with col:
-                    try:
-                        r = requests.get(entry["url"], timeout=10)
-                        r.raise_for_status()
-                        img = Image.open(BytesIO(r.content))
-                        st.image(img, caption=entry["title"], use_container_width=True)
-                        st.caption(entry["description"][:100] + "...")
-                        if st.button("🔍 View", key=f"view_{idx}"):
-                            st.session_state.image_lightbox = entry["url"]
-                    except Exception as e:
-                        st.warning(f"⚠️ Could not load image: {entry['url']}")
-        else:
-            st.info("ℹ️ No image references found yet.")
+        except Exception as e:
+            st.error(f"❌ Assistant Error: {e}")
+
+    # 🔍 Display Image Cards
+    if st.session_state.image_results:
+        st.markdown("### 🖼️ Image Results")
+        cols = st.columns(4)
+        for idx, item in enumerate(st.session_state.image_results):
+            with cols[idx % 4]:
+                try:
+                    img = Image.open(BytesIO(requests.get(item["url"]).content))
+                    st.image(img, use_container_width=True)
+                except:
+                    st.image("https://via.placeholder.com/150?text=Error", use_container_width=True)
+
+                st.markdown(f"**{item['title']}**", unsafe_allow_html=True)
+                st.caption(item.get("desc", "")[:150] + "...")
+                if st.button("🔍 View", key=f"img_{idx}"):
+                    st.session_state.image_lightbox = item["url"]
 
     if st.session_state.image_lightbox:
         st.markdown("### 🔬 Full Image View")
         try:
-            r = requests.get(st.session_state.image_lightbox, timeout=10)
-            r.raise_for_status()
-            full_img = Image.open(BytesIO(r.content))
-            st.image(full_img, caption=st.session_state.image_lightbox.split('/')[-1], use_container_width=True)
+            full_img = Image.open(BytesIO(requests.get(st.session_state.image_lightbox).content))
+            st.image(full_img, use_container_width=True)
         except Exception as e:
-            st.warning("⚠️ Could not load full image.")
-            st.error(str(e))
-        if st.button("❌ Close Preview"):
+            st.error(f"⚠️ Could not load full image: {e}")
+        if st.button("❌ Close Image View"):
             st.session_state.image_lightbox = None
             st.rerun()