Update ui/components.py

ui/components.py

@@ -70,29 +70,63 @@ def show_analysis_section(uploaded_files):
     analyze_btn = st.button("Analyze Charts", key="main_analyze_btn")
     return analyze_btn
 
-def show_chat_history(chat_history):
-    """Display chat history"""
-    st.subheader("Chat History")
-    for idx, chat in enumerate(chat_history):
-        timestamp = chat.get('timestamp', 'No date')
-        analysis_type = chat.get('analysis_type', 'Analysis')
-        label = f"{analysis_type} from {timestamp[:16]}"
-        
-        expander = st.expander(label)
-        with expander:
-            st.write(chat.get('analysis', ''))
-            if chat.get('question'):
-                st.write(f"Follow-up: {chat['question']}")
-
 def show_follow_up_section(key_suffix=""):
     """Show the follow-up question section"""
-    st.subheader("Continue Analysis")
     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}")
-    return follow_up if send_btn else None
+    
+    # Store the input in session state
+    input_key = f"followup_input_{key_suffix}"
+    if input_key not in st.session_state:
+        st.session_state[input_key] = ""
+    
+    # Create a unique key for the send button
+    button_key = f"followup_send_btn_{key_suffix}"
+    
+    # Use button click with session state to maintain conversation
+    if st.button("Send", key=button_key):
+        if follow_up:
+            return follow_up
+    return None
+
+def show_chat_history(chat_history):
+    """Display chat history"""
+    st.subheader("Chat History")
+    
+    # Group chat messages by conversation
+    current_conversation = []
+    for idx, chat in enumerate(chat_history):
+        # Start a new conversation if this is an initial analysis
+        if chat.get('analysis_type') == 'Individual' or chat.get('analysis_type') == 'Correlated Analysis':
+            if current_conversation:
+                display_conversation_group(current_conversation)
+            current_conversation = [chat]
+        else:
+            current_conversation.append(chat)
+    
+    # Display the last conversation group
+    if current_conversation:
+        display_conversation_group(current_conversation)
+
+def display_conversation_group(conversation):
+    """Display a group of related chat messages"""
+    if not conversation:
+        return
+        
+    # Use the first message timestamp for the expander label
+    timestamp = conversation[0].get('timestamp', 'No date')
+    analysis_type = conversation[0].get('analysis_type', 'Analysis')
+    
+    with st.expander(f"{analysis_type} from {timestamp[:16]}", expanded=True):
+        for message in conversation:
+            if message.get('analysis_type') in ['Individual', 'Correlated Analysis']:
+                st.markdown("**Initial Analysis:**")
+            elif message.get('question'):
+                st.markdown(f"**Follow-up Question:** {message['question']}")
+            
+            st.markdown(message.get('analysis', ''))
 
 def show_save_options():
     """Show save chat options"""