Update app.py

app.py

@@ -42,48 +42,51 @@ if st.sidebar.button("🔄 Clear Chat"):
 show_image = st.sidebar.checkbox("📖 Show Document Image", value=True)
 
 # ------------------ Split Layout ------------------
-col1, col2 = st.columns([1, 2])  # Adjust ratio as needed
+col1, col2 = st.columns([1, 2])
 
 # ------------------ Image Panel (Left) ------------------
 with col1:
     if show_image and st.session_state.image_url:
         st.image(st.session_state.image_url, caption="📑 Extracted Page", use_container_width=True)
-        st.session_state.image_updated = False  # Reset flag after rendering
+        st.session_state.image_updated = False
 
 # ------------------ Chat Panel (Right) ------------------
 with col2:
-    # Show messages in reverse order (latest at top)
-    for message in reversed(st.session_state.messages):
-        role, content = message["role"], message["content"]
-        st.chat_message(role).write(content)
+    # Display latest message block separately at the top (if available)
+    if st.session_state.messages:
+        latest_msg = st.session_state.messages[-1]
+        st.chat_message(latest_msg["role"]).write(latest_msg["content"])
 
-    # Chat input box at the bottom
+    # Show the rest of the chat history (excluding latest)
+    for message in st.session_state.messages[:-1]:
+        st.chat_message(message["role"]).write(message["content"])
+
+    # ------------------ Chat Input Handling ------------------
     if prompt := st.chat_input("Type your question about the document..."):
         st.session_state.messages.append({"role": "user", "content": prompt})
-        st.chat_message("user").write(prompt)
 
         try:
-            # Initialize thread if needed
+            # Create thread if it doesn't exist
             if st.session_state.thread_id is None:
                 thread = client.beta.threads.create()
                 st.session_state.thread_id = thread.id
 
             thread_id = st.session_state.thread_id
 
-            # Send message to assistant
+            # Send prompt
             client.beta.threads.messages.create(
                 thread_id=thread_id,
                 role="user",
                 content=prompt
             )
 
-            # Run assistant
+            # Trigger run
             run = client.beta.threads.runs.create(
                 thread_id=thread_id,
                 assistant_id=ASSISTANT_ID
             )
 
-            # Wait for assistant response
+            # Wait for completion
             with st.spinner("Assistant is thinking..."):
                 while True:
                     run_status = client.beta.threads.runs.retrieve(
@@ -102,10 +105,9 @@ with col2:
                     assistant_message = message.content[0].text.value
                     break
 
-            st.chat_message("assistant").write(assistant_message)
             st.session_state.messages.append({"role": "assistant", "content": assistant_message})
 
-            # Extract GitHub image from response if available
+            # Handle image from response if exists
             image_match = re.search(
                 r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png',
                 assistant_message
@@ -113,7 +115,8 @@ with col2:
             if image_match:
                 st.session_state.image_url = image_match.group(0)
                 st.session_state.image_updated = True
-                st.rerun()  # Trigger rerun to refresh image display
+
+            st.rerun()
 
         except Exception as e:
             st.error(f"❌ Error: {str(e)}")