Update app.py

app.py

@@ -132,9 +132,17 @@ with chat_col:
                 for m in reversed(messages.data):
                     if m.role == "assistant":
                         reply = m.content[0].text.value
-                        st.session_state.messages.append({"role": "assistant", "content": reply})
 
-                        # ✅ Extract GitHub raw image URLs
+                        # ✂️ Split output into main and references (if structured this way)
+                        split_marker = "---"
+                        if split_marker in reply:
+                            main_answer, references = reply.split(split_marker, 1)
+                            st.session_state.messages.append({"role": "assistant", "content": main_answer.strip()})
+                            st.session_state.messages.append({"role": "references", "content": references.strip()})
+                        else:
+                            st.session_state.messages.append({"role": "assistant", "content": reply})
+
+                        # Extract GitHub image URLs
                         image_matches = re.findall(
                             r'https://raw\.githubusercontent\.com/AndrewLORTech/witspathologai/main/[^\s\n]+\.png',
                             reply
@@ -148,16 +156,15 @@ with chat_col:
             st.error(f"❌ Error: {e}")
 
     for msg in st.session_state.messages:
-        with st.chat_message(msg["role"]):
+        with st.chat_message("assistant" if msg["role"] != "references" else "📄"):
             st.markdown(msg["content"], unsafe_allow_html=True)
 
-# ------------------ Image Display with Encoding Fix ------------------
+# ------------------ Image Viewer (Scrollable Format) ------------------
 with image_col:
     if show_image and st.session_state.image_urls:
         st.markdown("### 🖼️ Slide Previews")
-        for raw_url in st.session_state.image_urls:
+        for i, raw_url in enumerate(st.session_state.image_urls):
             try:
-                # Extract everything after the domain and encode each path segment
                 if "githubusercontent.com" not in raw_url:
                     continue
                 _, raw_path = raw_url.split("githubusercontent.com/", 1)
@@ -165,11 +172,6 @@ with image_col:
                 encoded_segments = [quote(seg) for seg in segments]
                 encoded_url = "https://raw.githubusercontent.com/" + "/".join(encoded_segments)
 
-                # Load and display
-                r = requests.get(encoded_url)
-                r.raise_for_status()
-                img = Image.open(BytesIO(r.content))
-                st.image(img, caption=f"📷 {encoded_segments[-1]}", use_container_width=True)
-
+                st.image(encoded_url, caption=f"Slide {i + 1}", use_column_width=True)
             except Exception as e:
-                st.error(f"❌ Failed to load image from {raw_url}: {e}")
+                st.error(f"❌ Failed to load image: {e}")