Update app.py

app.py

@@ -164,7 +164,7 @@ 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 "response" in key else []
+            st.session_state[key] = None if key in ["image_thread_id", "image_response", "image_lightbox"] else []
 
     st.markdown("### 🔍 Search Histology Visual References")
     image_input = st.chat_input("e.g. mitosis, seminiferous tubule, cell cycle")
@@ -189,7 +189,7 @@ with tab2:
                 assistant_id=ASSISTANT_ID
             )
 
-            with st.spinner("🧬 Retrieving images..."):
+            with st.spinner("🧬 Retrieving histology references..."):
                 while True:
                     status = client.beta.threads.runs.retrieve(
                         thread_id=st.session_state.image_thread_id,
@@ -203,46 +203,55 @@ with tab2:
                 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
+                        response_text = msg.content[0].text.value.strip()
                         st.session_state.image_response = response_text
 
-                        # Extract URLs and details for grid card display
-                        lines = response_text.splitlines()
-                        block = {}
-                        all_cards = []
-                        for line in lines:
-                            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:
-                                block["url"] = line.split("Image URL:")[1].strip()
-                        if block: all_cards.append(block)
-                        st.session_state.image_results = all_cards
+                        # ✅ Robust markdown block parser
+                        blocks = response_text.split("### 🖼️")[1:]
+                        results = []
+                        for b in blocks:
+                            lines = b.strip().splitlines()
+                            card = {
+                                "title": lines[0].strip() if lines else "Untitled",
+                                "desc": "",
+                                "url": ""
+                            }
+                            for line in lines:
+                                if line.startswith("- **Description:**"):
+                                    card["desc"] = line.replace("- **Description:**", "").strip()
+                                elif line.startswith("- **Image URL:**"):
+                                    card["url"] = line.replace("- **Image URL:**", "").strip()
+                                elif line.startswith("http") and not card["url"]:
+                                    card["url"] = line.strip()
+                            if card["url"]:
+                                results.append(card)
+
+                        st.session_state.image_results = results
                         break
 
         except Exception as e:
             st.error(f"❌ Assistant Error: {e}")
 
-    # 🔍 Display Image Cards
+    # 🖼️ Display Image Cards in 4-column layout
     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"]
+        for i in range(0, len(st.session_state.image_results), 4):
+            row = st.columns(4)
+            for j, col in enumerate(row):
+                idx = i + j
+                if idx < len(st.session_state.image_results):
+                    item = st.session_state.image_results[idx]
+                    with col:
+                        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=Image+Error", use_container_width=True)
+
+                        st.markdown(f"**{item.get('title', 'No Title')}**", unsafe_allow_html=True)
+                        st.caption(item.get("desc", "")[:150] + "..." if item.get("desc") else "")
+                        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")