Update ui/components.py

ui/components.py

@@ -101,14 +101,34 @@ def show_chat_history(chat_history):
             st.markdown("---")
 
 def show_follow_up_section(key_suffix=""):
-    """Show the follow-up question section"""
-    follow_up = st.text_input(
-        "Ask a follow-up question:",
-        key=f"followup_input_{key_suffix}"
-    )
-    send_btn = st.button("Send", key=f"followup_send_btn_{key_suffix}")
+    """Show the follow-up question section with enhanced UI"""
+    st.markdown("---")  # Visual separator
+    
+    # Chat container
+    chat_container = st.container()
+    with chat_container:
+        # Show conversation context if exists
+        if 'conversation_context' in st.session_state:
+            for msg in st.session_state.conversation_context:
+                if msg['role'] == 'user':
+                    st.markdown(f"**You:** {msg['content']}")
+                else:
+                    st.markdown(f"**Assistant:** {msg['content']}")
+                st.markdown("---")
+    
+    # Input section
+    col1, col2 = st.columns([4, 1])
+    with col1:
+        follow_up = st.text_input(
+            "Ask a follow-up question:",
+            key=f"followup_input_{key_suffix}",
+            placeholder="Type your question here..."
+        )
+    with col2:
+        send_btn = st.button("Send", key=f"followup_send_btn_{key_suffix}")
+    
     return follow_up if send_btn else None
-
+    
 def show_save_options():
     """Show save chat options"""
     st.subheader("Save Analysis")