Update app.py

app.py

@@ -41,58 +41,58 @@ if st.sidebar.button("🔄 Clear Chat"):
 
 show_image = st.sidebar.checkbox("📖 Show Document Image", value=True)
 
-# ------------------ Split Layout ------------------
+# ------------------ Layout: Image (Left) and Chat (Right) ------------------
 col1, col2 = st.columns([1, 2])
 
-# ------------------ Image Panel (Left) ------------------
+# ------------------ Left Panel: Document Image ------------------
 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
 
-# ------------------ Chat Panel (Right) ------------------
+# ------------------ Right Panel: Chat Interface ------------------
 with col2:
-    # Pair user + assistant messages
+    # -- Chat Input (Always at the Top)
+    prompt = st.chat_input("Type your question about the document...")
+
+    # -- Pair messages (user + assistant) for display
     paired_messages = []
     buffer = []
-
     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"])
+    # -- Display chat history (latest first, older scroll down)
+    with st.container():
+        for pair in reversed(paired_messages):
+            for msg in pair:
+                with st.chat_message(msg["role"]):
+                    st.write(msg["content"])
 
-    # ------------------ Chat Input Handling ------------------
-    if prompt := st.chat_input("Type your question about the document..."):
+    # -- Process new prompt if entered
+    if prompt:
         st.session_state.messages.append({"role": "user", "content": prompt})
 
         try:
-            # Create thread if not exists
+            # Initialize thread if not already done
             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 user message to assistant
             client.beta.threads.messages.create(
-                thread_id=thread_id,
+                thread_id=st.session_state.thread_id,
                 role="user",
                 content=prompt
             )
 
-            # Run assistant
+            # Trigger assistant run
             run = client.beta.threads.runs.create(
-                thread_id=thread_id,
+                thread_id=st.session_state.thread_id,
                 assistant_id=ASSISTANT_ID
             )
 
@@ -100,15 +100,15 @@ with col2:
             with st.spinner("Assistant is thinking..."):
                 while True:
                     run_status = client.beta.threads.runs.retrieve(
-                        thread_id=thread_id,
+                        thread_id=st.session_state.thread_id,
                         run_id=run.id
                     )
                     if run_status.status == "completed":
                         break
                     time.sleep(1)
 
-            # Retrieve assistant response
-            messages = client.beta.threads.messages.list(thread_id=thread_id)
+            # Retrieve latest assistant message
+            messages = client.beta.threads.messages.list(thread_id=st.session_state.thread_id)
             assistant_message = None
             for message in reversed(messages.data):
                 if message.role == "assistant":
@@ -117,7 +117,7 @@ with col2:
 
             st.session_state.messages.append({"role": "assistant", "content": assistant_message})
 
-            # Extract image URL if present
+            # -- Extract GitHub-hosted image URL if present
             image_match = re.search(
                 r'https://raw\.githubusercontent\.com/AndrewLORTech/surgical-pathology-manual/main/[\w\-/]*\.png',
                 assistant_message
@@ -126,7 +126,7 @@ with col2:
                 st.session_state.image_url = image_match.group(0)
                 st.session_state.image_updated = True
 
-            st.rerun()
+            st.rerun()  # Rerun to show updated image + message at top
 
         except Exception as e:
             st.error(f"❌ Error: {str(e)}")