Update app.py

app.py

@@ -55,13 +55,12 @@ for key in ["messages", "thread_id", "image_urls", "pending_prompt", "image_url"
         st.session_state[key] = [] if key.endswith("s") else None if "url" in key else False
 
 # ------------------ Tabs ------------------
-show_tab2 = False  # Toggle this to True to activate Tab 2
-
+show_tab2 = True  # Set to True to activate visual tab
 if show_tab2:
     tab1, tab2 = st.tabs(["💬 Chat Assistant", "📷 Visual Reference Search"])
 else:
-    tab1, = st.tabs(["💬 Chat Assistant"])
-
+    tab1 = st.tabs(["💬 Chat Assistant"])[0]
+    tab2 = None
 
 # ------------------ Tab 1: Chat Assistant ------------------
 with tab1:
@@ -166,121 +165,139 @@ with tab1:
 
 
 # ------------------ Tab 2: Visual Reference Search ------------------
-if show_tab2:
-    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
+import urllib.parse
+import requests
+from PIL import Image
+from io import BytesIO
 
-        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 = []
+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
+
+    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
+
+        try:
+            if st.session_state.image_thread_id is None:
+                thread = client.beta.threads.create()
+                st.session_state.image_thread_id = thread.id
+
+            client.beta.threads.messages.create(
+                thread_id=st.session_state.image_thread_id,
+                role="user",
+                content=image_input
+            )
+
+            run = client.beta.threads.runs.create(
+                thread_id=st.session_state.image_thread_id,
+                assistant_id=ASSISTANT_ID
+            )
+
+            with st.spinner("🔬 Searching for histology references..."):
+                while True:
+                    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)
+
+            if run_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
+
+                        # Extract and decode image URLs
+                        lines = response_text.splitlines()
+                        image_urls = []
+                        expecting_url = False
+
+                        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(urllib.parse.unquote(parts[1].strip()))
+                                else:
+                                    expecting_url = True
+                            elif expecting_url:
+                                if line_clean.startswith("http"):
+                                    image_urls.append(urllib.parse.unquote(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]
+                        break
+        except Exception as e:
+            st.error(f"❌ Visual Assistant Error: {e}")
+
+    if st.session_state.image_results and st.session_state.image_response:
+        st.subheader("🖼️ Image Preview(s)")
+
+        # Split the assistant response into metadata cards
+        cards = []
+        blocks = st.session_state.image_response.split("### 🖼️ ")
+        for block in blocks:
+            if not block.strip():
+                continue
+            lines = block.strip().splitlines()
+            title = lines[0].strip()
+            meta = "\n".join(lines[1:])
+            cards.append((title, meta))
+
+        cols = st.columns(4)
+        for i, ((title, meta), item) in enumerate(zip(cards, st.session_state.image_results)):
+            image_url = item.get("image")
+            with cols[i % 4]:
+                with st.container():
+                    # Remove Image Filename and Image URL lines from display
+                    meta_clean = "\n".join(
+                        line for line in meta.splitlines()
+                        if all(x not in line.lower() for x in ["image filename", "image url"])
+                    )
+
+                    st.markdown(f"**🔬 {title}**")
+                    st.caption(meta_clean)
+                    try:
+                        r = requests.get(image_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("🔍 Zoom", key=f"zoom_{i}"):
+                            st.session_state.image_lightbox = image_url
+                    except Exception as e:
+                        st.warning("⚠️ Could not load image.")
+                        st.error(str(e))
+    else:
+        st.info("ℹ️ No image references found yet.")
+
+    if st.session_state.image_lightbox:
+        st.markdown("### 🔬 Full Image View")
+        try:
+            img_url = urllib.parse.unquote(st.session_state.image_lightbox)
+            r = requests.get(img_url, 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)
+        except Exception as e:
+            st.warning("⚠️ Could not load full image.")
+            st.error(str(e))
+        if st.button("❌ Close Preview"):
             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
-
-                client.beta.threads.messages.create(
-                    thread_id=st.session_state.image_thread_id,
-                    role="user",
-                    content=image_input
-                )
-
-                run = client.beta.threads.runs.create(
-                    thread_id=st.session_state.image_thread_id,
-                    assistant_id=ASSISTANT_ID
-                )
-
-                with st.spinner("🔬 Searching for histology references..."):
-                    while True:
-                        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)
-
-                if run_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
-
-                            # ✅ Extract Image URLs from assistant markdown
-                            lines = response_text.splitlines()
-                            image_urls = []
-                            expecting_url = False
-
-                            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]
-                            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)")
-                for i, item in enumerate(st.session_state.image_results):
-                    image_url = item.get("image")
-                    if image_url:
-                        try:
-                            r = requests.get(image_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
-                        except Exception as e:
-                            st.warning(f"⚠️ Could not load: {image_url}")
-                            st.error(f"{e}")
-            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.raise_for_status()
-                full_img = Image.open(BytesIO(r.content))
-                st.image(full_img, caption=img_url.split("/")[-1], 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.session_state.image_lightbox = None
-                st.rerun()
+            st.rerun()