Update app.py

app.py

@@ -54,7 +54,13 @@ for key in ["messages", "thread_id", "image_urls", "pending_prompt", "image_url"
         st.session_state[key] = [] if key.endswith("s") else None if "url" in key else False
 
 # ------------------ Tabs ------------------
-tab1 = st.tabs(["💬 Chat Assistant"])
+show_tab2 = False  # Toggle this to True to activate Tab 2
+
+if show_tab2:
+    tab1, tab2 = st.tabs(["💬 Chat Assistant", "📷 Visual Reference Search"])
+else:
+    tab1, = st.tabs(["💬 Chat Assistant"])
+
 
 # ------------------ Tab 1: Chat Assistant ------------------
 with tab1:
@@ -159,120 +165,121 @@ with tab1:
 
 
 # ------------------ Tab 2: Visual Reference Search ------------------
-with tab2:
-    ASSISTANT_ID = "asst_9v09zgizdcuuhNdcFQpRo9RO"
-
-    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 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,
-                role="user",
-                content=image_input
-            )
-
-            run = client.beta.threads.runs.create(
-                thread_id=st.session_state.image_thread_id,
-                assistant_id=ASSISTANT_ID
-            )
-
-            with st.spinner("🔬 Searching for histology references..."):
-                while True:
-                    run_status = client.beta.threads.runs.retrieve(
-                        thread_id=st.session_state.image_thread_id,
-                        run_id=run.id
-                    )
-                    if run_status.status in ("completed", "failed", "cancelled"):
-                        break
-                    time.sleep(1)
-
-            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
-                        st.session_state.image_response = response_text
-
-                        # ✅ 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"❌ Visual Assistant Error: {e}")
-
-    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)
+if show_tab2:
+    with tab2:
+        ASSISTANT_ID = "asst_9v09zgizdcuuhNdcFQpRo9RO"
+
+        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
 
-    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.")
-
-    if st.session_state.image_lightbox:
-        st.markdown("### 🔬 Full Image View")
-        try:
-            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.warning("⚠️ Could not load full image.")
-            st.error(str(e))
-        if st.button("❌ Close Preview"):
+        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
-            st.rerun()
+
+            try:
+                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,
+                    role="user",
+                    content=image_input
+                )
+
+                run = client.beta.threads.runs.create(
+                    thread_id=st.session_state.image_thread_id,
+                    assistant_id=ASSISTANT_ID
+                )
+
+                with st.spinner("🔬 Searching for histology references..."):
+                    while True:
+                        run_status = client.beta.threads.runs.retrieve(
+                            thread_id=st.session_state.image_thread_id,
+                            run_id=run.id
+                        )
+                        if run_status.status in ("completed", "failed", "cancelled"):
+                            break
+                        time.sleep(1)
+
+                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
+                            st.session_state.image_response = response_text
+
+                            # ✅ 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"❌ Visual Assistant Error: {e}")
+
+        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)
+
+        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)
+                            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.")
+
+        if st.session_state.image_lightbox:
+            st.markdown("### 🔬 Full Image View")
+            try:
+                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.warning("⚠️ Could not load full image.")
+                st.error(str(e))
+            if st.button("❌ Close Preview"):
+                st.session_state.image_lightbox = None
+                st.rerun()