Update app.py

app.py

@@ -48,17 +48,15 @@ if not OPENAI_API_KEY:
 client = OpenAI(api_key=OPENAI_API_KEY)
 
 # ------------------ Assistant Configuration ------------------
-ASSISTANT_ID = "asst_jXDSjCG8LI4HEaFEcjFVq8KB"  # 🔁 Replace with your actual OpenAI assistant ID
+ASSISTANT_ID = "asst_jXDSjCG8LI4HEaFEcjFVq8KB"  # Replace with your Assistant ID
 
 # ------------------ Session State ------------------
 if "messages" not in st.session_state:
     st.session_state.messages = []
 if "thread_id" not in st.session_state:
     st.session_state.thread_id = None
-if "image_url" not in st.session_state:
-    st.session_state.image_url = None
-if "image_updated" not in st.session_state:
-    st.session_state.image_updated = False
+if "image_urls" not in st.session_state:
+    st.session_state.image_urls = []
 if "pending_prompt" not in st.session_state:
     st.session_state.pending_prompt = None
 
@@ -68,12 +66,11 @@ with st.sidebar:
     if st.button("🧹 Clear Chat"):
         st.session_state.messages = []
         st.session_state.thread_id = None
-        st.session_state.image_url = None
-        st.session_state.image_updated = False
+        st.session_state.image_urls = []
         st.session_state.pending_prompt = None
         st.rerun()
 
-    show_image = st.toggle("📸 Show Slide Image", value=True)
+    show_image = st.toggle("📸 Show Slide Images", value=True)
     keyword = st.text_input("Keyword Search", placeholder="e.g. mitosis, carcinoma")
     if st.button("🔎 Search") and keyword:
         st.session_state.pending_prompt = f"Find clauses or references related to: {keyword}"
@@ -136,13 +133,13 @@ with chat_col:
                     if m.role == "assistant":
                         reply = m.content[0].text.value
                         st.session_state.messages.append({"role": "assistant", "content": reply})
-                        match = re.search(r'Document Reference:\s*(.*?),\s*Page\s*(\d+)', reply)
-                        if match:
-                            doc, page = match.group(1).strip(), int(match.group(2))
-                            folder = quote(doc)
-                            img_url = f"https://raw.githubusercontent.com/AndrewLORTech/witspathologai/main/{folder}/{folder}_page_{page:04d}.png"
-                            st.session_state.image_url = img_url
-                            st.session_state.image_updated = True
+
+                        # 🧠 Extract ALL image URLs from assistant reply
+                        image_matches = re.findall(
+                            r'https://raw\.githubusercontent\.com/AndrewLORTech/witspathologai/main/[^ \n]+\.png',
+                            reply
+                        )
+                        st.session_state.image_urls = image_matches
                         break
             else:
                 st.error("❌ Assistant failed to respond.")
@@ -154,13 +151,15 @@ with chat_col:
         with st.chat_message(msg["role"]):
             st.markdown(msg["content"], unsafe_allow_html=True)
 
-# ------------------ Right Column: OCR Page Image ------------------
+# ------------------ Right Column: Image Previews ------------------
 with image_col:
-    if show_image and st.session_state.image_url:
-        try:
-            r = requests.get(st.session_state.image_url)
-            r.raise_for_status()
-            img = Image.open(BytesIO(r.content))
-            st.image(img, caption="🧾 OCR Page Image", use_container_width=True)
-        except Exception as e:
-            st.error(f"🖼️ Failed to load image: {e}")
+    if show_image and st.session_state.image_urls:
+        st.markdown("### 🖼️ Slide Previews")
+        for url in st.session_state.image_urls:
+            try:
+                r = requests.get(url)
+                r.raise_for_status()
+                img = Image.open(BytesIO(r.content))
+                st.image(img, caption=f"📷 {url.split('/')[-1]}", use_container_width=True)
+            except Exception as e:
+                st.error(f"❌ Failed to load image from {url}: {e}")