Update app.py

app.py

@@ -29,24 +29,28 @@ def login():
 
 if "authenticated" not in st.session_state:
     st.session_state.authenticated = False
+
 if not st.session_state.authenticated:
     login()
     st.stop()
 
-# ------------------ Configuration ------------------
+# ------------------ App Configuration ------------------
 st.set_page_config(page_title="AI Pathology Assistant", layout="wide", initial_sidebar_state="collapsed")
 st.title("🧬 AI Pathology Assistant")
 st.caption("AI-powered exploration of pathology, anatomy, and histology documents via OCR + GPT")
 
+# ------------------ Load OpenAI API Key ------------------
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
 if not OPENAI_API_KEY:
     st.error("❌ Missing OPENAI_API_KEY environment variable.")
     st.stop()
 
 client = OpenAI(api_key=OPENAI_API_KEY)
-ASSISTANT_ID = "asst_jXDSjCG8LI4HEaFEcjFVq8KB"
 
-# ------------------ State Setup ------------------
+# ------------------ Assistant Configuration ------------------
+ASSISTANT_ID = "asst_jXDSjCG8LI4HEaFEcjFVq8KB"  # Replace with your actual assistant ID
+
+# ------------------ Session State ------------------
 if "messages" not in st.session_state:
     st.session_state.messages = []
 if "thread_id" not in st.session_state:
@@ -66,7 +70,7 @@ with st.sidebar:
         st.session_state.pending_prompt = None
         st.rerun()
 
-    show_image = st.toggle("📸 Show Images", 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}"
@@ -87,7 +91,7 @@ with st.sidebar:
     if action != actions[0]:
         st.session_state.pending_prompt = action
 
-# ------------------ Chat UI ------------------
+# ------------------ Main Chat UI ------------------
 chat_col, image_col = st.columns([2, 1])
 
 with chat_col:
@@ -127,20 +131,15 @@ with chat_col:
                 messages = client.beta.threads.messages.list(thread_id=st.session_state.thread_id)
                 for m in reversed(messages.data):
                     if m.role == "assistant":
-                        raw_reply = m.content[0].text.value
+                        reply = m.content[0].text.value
+                        st.session_state.messages.append({"role": "assistant", "content": reply})
 
+                        # ✅ Extract GitHub raw image URLs
                         image_matches = re.findall(
-                            r'https://raw\.githubusercontent\.com/AndrewLORTech/witspathologai/main/[^"]+?\.png',
-                            raw_reply
+                            r'https://raw\.githubusercontent\.com/AndrewLORTech/witspathologai/main/[^\s\n]+\.png',
+                            reply
                         )
                         st.session_state.image_urls = image_matches
-
-                        reply_cleaned = re.sub(
-                            r'https://raw\.githubusercontent\.com/AndrewLORTech/witspathologai/main/[^"]+?\.png',
-                            '[Image reference available in Slide Previews →]',
-                            raw_reply
-                        )
-                        st.session_state.messages.append({"role": "assistant", "content": reply_cleaned})
                         break
             else:
                 st.error("❌ Assistant failed to respond.")
@@ -152,19 +151,25 @@ with chat_col:
         with st.chat_message(msg["role"]):
             st.markdown(msg["content"], unsafe_allow_html=True)
 
-# ------------------ Image Display (Vertical Stack) ------------------
+# ------------------ Image Display with Encoding Fix ------------------
 with image_col:
     if show_image and st.session_state.image_urls:
-        st.markdown("### 🖼️ Image(s)")
+        st.markdown("### 🖼️ Slide Previews")
         for raw_url in 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)
                 segments = raw_path.strip().split("/")
                 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)
+
             except Exception as e:
-                st.error(f"❌ Failed to load image: {e}")
+                st.error(f"❌ Failed to load image from {raw_url}: {e}")