Update app.py

app.py

@@ -146,11 +146,10 @@ for message in st.session_state.chat_history:
         st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {message['content']}</div>", unsafe_allow_html=True)
 st.markdown("</div>", unsafe_allow_html=True)
 
-# --- Input & Button Handling ---
-# --- Input Handling with "Enter to Send" ---
-user_input = st.chat_input("Type your message here...")  # Removes "Send" button
+# --- Single Input Box for Both Initial and Follow-Up Messages ---
+user_input = st.chat_input("Type your message here...")  # Single input box
 
-if user_input and not st.session_state.response_ready:
+if user_input:
     response, follow_up, st.session_state.chat_history, image_url = get_response(
         system_message="You are a helpful AI assistant.",
         user_text=user_input,
@@ -162,9 +161,18 @@ if user_input and not st.session_state.response_ready:
     if image_url:
         st.image(image_url, caption="NASA Image of the Day")
 
+    # Store follow-up question in session state
     st.session_state.follow_up = follow_up
     st.session_state.response_ready = True  # Ensures response cycle continues
 
+# Display follow-up question inside chat if available
+if st.session_state.response_ready and st.session_state.follow_up:
+    st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {st.session_state.follow_up}</div>", unsafe_allow_html=True)
+
+    # Reset response state so user can type next input
+    st.session_state.response_ready = False
+
+
 # Conversational Follow-up with "Enter to Send"
 if st.session_state.response_ready and st.session_state.follow_up:
     st.markdown(f"<div class='assistant-msg'><strong>HAL:</strong> {st.session_state.follow_up}</div>", unsafe_allow_html=True)