Update app.py

app.py

@@ -162,14 +162,9 @@ with tab1:
 with tab2:
     ASSISTANT_ID = "asst_9v09zgizdcuuhNdcFQpRo9RO"
 
-    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
+    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 []
 
     image_input = st.chat_input("Ask for histology visual references (e.g. ovary histology, mitosis)")
     if image_input:
@@ -178,9 +173,8 @@ with tab2:
         st.session_state.image_lightbox = None
 
         try:
-            if st.session_state.image_thread_id is None:
-                thread = client.beta.threads.create()
-                st.session_state.image_thread_id = thread.id
+            if not st.session_state.image_thread_id:
+                st.session_state.image_thread_id = client.beta.threads.create().id
 
             client.beta.threads.messages.create(
                 thread_id=st.session_state.image_thread_id,
@@ -210,34 +204,33 @@ with tab2:
                         response_text = msg.content[0].text.value
                         st.session_state.image_response = response_text
 
-                        # ✅ Extract Image URLs from assistant markdown
+                        # ✅ Image metadata parser
                         lines = response_text.splitlines()
-                        image_urls = []
-                        expecting_url = False
+                        image_entries = []
+                        current = {"title": "", "description": "", "url": ""}
 
                         for line in lines:
-                            line_clean = line.strip().replace("**", "")
-                            if "Image URL:" in line_clean:
-                                parts = line_clean.split("Image URL:")
-                                if len(parts) > 1 and parts[1].strip().startswith("http"):
-                                    image_urls.append(parts[1].strip())
-                                else:
-                                    expecting_url = True
-                            elif expecting_url:
-                                if line_clean.startswith("http"):
-                                    image_urls.append(line_clean)
-                                expecting_url = False
-
-                        st.session_state.image_results = [{"image": url} for url in image_urls]
-
-                        if image_urls and not st.session_state.image_lightbox:
-                            st.session_state.image_lightbox = image_urls[0]
+                            if line.startswith("### 🖼️"):
+                                current = {"title": line.replace("### 🖼️", "").strip(), "description": "", "url": ""}
+                            elif line.strip().startswith("- **Description:**"):
+                                current["description"] = 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"]
                         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")
@@ -246,30 +239,30 @@ with tab2:
     with image_col:
         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")
-                if image_url:
+            cols = st.columns(4)
+            for idx, entry in enumerate(st.session_state.image_results):
+                col = cols[idx % 4]
+                with col:
                     try:
-                        r = requests.get(image_url, timeout=10)  # ✅ Use as-is
+                        r = requests.get(entry["url"], timeout=10)
                         r.raise_for_status()
                         img = Image.open(BytesIO(r.content))
-                        st.image(img, caption=image_url.split("/")[-1], use_container_width=True)
-                        if st.button("🔍 View Full Image", key=f"full_img_{i}"):
-                            st.session_state.image_lightbox = image_url
+                        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_url}")
-                        st.error(f"{e}")
+                        st.warning(f"⚠️ Could not load image: {entry['url']}")
         else:
             st.info("ℹ️ No image references found yet.")
 
     if st.session_state.image_lightbox:
         st.markdown("### 🔬 Full Image View")
         try:
-            img_url = st.session_state.image_lightbox
-            r = requests.get(img_url, timeout=10)
+            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=img_url.split("/")[-1], use_container_width=True)
+            st.image(full_img, caption=st.session_state.image_lightbox.split('/')[-1], use_container_width=True)
         except Exception as e:
             st.warning("⚠️ Could not load full image.")
             st.error(str(e))