Update app.py

app.py

@@ -52,41 +52,51 @@ with col1:
 
 # ------------------ Chat Panel (Right) ------------------
 with col2:
-    # 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"])
+    # Pair user + assistant messages
+    paired_messages = []
+    buffer = []
 
-    # 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"])
+    for msg in st.session_state.messages:
+        buffer.append(msg)
+        if msg["role"] == "assistant" and len(buffer) == 2:
+            paired_messages.append(buffer.copy())
+            buffer.clear()
+
+    # If last message is from user and no assistant reply yet
+    if buffer:
+        paired_messages.append(buffer.copy())
+
+    # Show pairs in reverse order (latest on top)
+    for pair in reversed(paired_messages):
+        for msg in pair:
+            st.chat_message(msg["role"]).write(msg["content"])
 
     # ------------------ Chat Input Handling ------------------
     if prompt := st.chat_input("Type your question about the document..."):
         st.session_state.messages.append({"role": "user", "content": prompt})
 
         try:
-            # Create thread if it doesn't exist
+            # Create thread if not exists
             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 prompt
+            # Send user message to assistant
             client.beta.threads.messages.create(
                 thread_id=thread_id,
                 role="user",
                 content=prompt
             )
 
-            # Trigger run
+            # Run assistant
             run = client.beta.threads.runs.create(
                 thread_id=thread_id,
                 assistant_id=ASSISTANT_ID
             )
 
-            # Wait for completion
+            # Wait for assistant to respond
             with st.spinner("Assistant is thinking..."):
                 while True:
                     run_status = client.beta.threads.runs.retrieve(
@@ -97,7 +107,7 @@ with col2:
                         break
                     time.sleep(1)
 
-            # Get assistant response
+            # Retrieve assistant response
             messages = client.beta.threads.messages.list(thread_id=thread_id)
             assistant_message = None
             for message in reversed(messages.data):
@@ -107,7 +117,7 @@ with col2:
 
             st.session_state.messages.append({"role": "assistant", "content": assistant_message})
 
-            # Handle image from response if exists
+            # Extract image URL if present
             image_match = re.search(
                 r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png',
                 assistant_message