Update app.py

app.py

@@ -162,21 +162,25 @@ with tab1:
 with tab2:
     ASSISTANT_ID = "asst_9v09zgizdcuuhNdcFQpRo9RO"
 
-    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 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")
+    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 not st.session_state.image_thread_id:
-                st.session_state.image_thread_id = client.beta.threads.create().id
+            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,
@@ -189,72 +193,86 @@ with tab2:
                 assistant_id=ASSISTANT_ID
             )
 
-            with st.spinner("🧬 Retrieving histology references..."):
+            with st.spinner("🔬 Searching for histology references..."):
                 while True:
-                    status = client.beta.threads.runs.retrieve(
+                    run_status = client.beta.threads.runs.retrieve(
                         thread_id=st.session_state.image_thread_id,
                         run_id=run.id
                     )
-                    if status.status in ("completed", "failed", "cancelled"):
+                    if run_status.status in ("completed", "failed", "cancelled"):
                         break
                     time.sleep(1)
 
-            if status.status == "completed":
+            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.strip()
+                        response_text = msg.content[0].text.value
                         st.session_state.image_response = response_text
 
-                        # ✅ Reliable parser for markdown blocks
-                        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 "Description:" in line:
-                                    card["desc"] = line.split("Description:")[1].strip()
-                                elif "Image URL:" in line:
-                                    card["url"] = line.split("Image URL:")[1].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
+                        # ✅ 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"❌ Assistant Error: {e}")
+            st.error(f"❌ Visual Assistant Error: {e}")
 
-    # 🖼️ Display images in vertical stack
-    if st.session_state.image_results:
-        st.markdown("### 🖼️ Image Results")
-        for idx, item in enumerate(st.session_state.image_results):
-            try:
-                img = Image.open(BytesIO(requests.get(item["url"], timeout=10).content))
-                st.image(img, use_container_width=True)
-            except:
-                st.image("https://via.placeholder.com/150?text=Image+Error", use_container_width=True)
+    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)
 
-            st.markdown(f"**{item.get('title', 'No Title')}**", unsafe_allow_html=True)
-            st.caption(item.get("desc", "")[:200] + "..." if item.get("desc") else "")
-            if st.button("🔍 View", key=f"img_{idx}"):
-                st.session_state.image_lightbox = item["url"]
+    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)  # ✅ Use as-is
+                        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.")
 
-    # 🔬 Full image lightbox
     if st.session_state.image_lightbox:
         st.markdown("### 🔬 Full Image View")
         try:
-            full_img = Image.open(BytesIO(requests.get(st.session_state.image_lightbox).content))
-            st.image(full_img, use_container_width=True)
+            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.error(f"⚠️ Could not load full image: {e}")
-        if st.button("❌ Close Image View"):
+            st.warning("⚠️ Could not load full image.")
+            st.error(str(e))
+        if st.button("❌ Close Preview"):
             st.session_state.image_lightbox = None
             st.rerun()