Update app.py

app.py

@@ -219,20 +219,25 @@ with tab2:
                         response_text = msg.content[0].text.value
                         st.session_state.image_response = response_text
 
-                        # ✅ FIXED: Robustly parse "Image URL" from markdown-formatted response
+                        # ✅ Robust URL extraction from markdown-formatted response
                         lines = response_text.splitlines()
                         image_urls = []
                         for i, line in enumerate(lines):
-                            clean_line = re.sub(r"\*\*", "", line).strip()
-                            if "Image URL:" in clean_line:
-                                url = clean_line.split("Image URL:")[-1].strip()
-                                if not url.startswith("http") and i + 1 < len(lines):
-                                    next_line = re.sub(r"\*\*", "", lines[i + 1].strip())
+                            clean_line = re.sub(r"[*`]", "", line).strip().lower()
+                            if "image url:" in clean_line:
+                                url_candidate = line.split("Image URL:")[-1].strip()
+                                if not url_candidate.startswith("http") and i + 1 < len(lines):
+                                    next_line = re.sub(r"[*`]", "", lines[i + 1].strip())
                                     if next_line.startswith("http"):
-                                        url = next_line
-                                if url.startswith("http"):
-                                    image_urls.append(url)
+                                        url_candidate = next_line
+                                if url_candidate.startswith("http"):
+                                    image_urls.append(url_candidate)
+
                         st.session_state.image_results = [{"image": url} for url in image_urls]
+
+                        # Optional: auto-preview first image
+                        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}")
@@ -243,6 +248,8 @@ with tab2:
         if st.session_state.image_response:
             st.markdown("### 🧠 Assistant Response")
             st.markdown(st.session_state.image_response, unsafe_allow_html=True)
+            # Optional: Debug raw output
+            # st.code(st.session_state.image_response, language="markdown")
 
     with image_col:
         if st.session_state.image_results: