Update services/chart_analysis.py

services/chart_analysis.py

@@ -88,13 +88,15 @@ class ChartAnalysisService:
             if previous_analysis and not st.session_state.conversation_context:
                 st.session_state.conversation_context.append({
                     'role': 'assistant',
-                    'content': previous_analysis
+                    'content': previous_analysis,
+                    'timestamp': datetime.now().isoformat()
                 })
             
             # Add current question to context
             st.session_state.conversation_context.append({
                 'role': 'user',
-                'content': question
+                'content': question,
+                'timestamp': datetime.now().isoformat()
             })
             
             # Create a conversation-aware prompt
@@ -111,20 +113,36 @@ class ChartAnalysisService:
                 # Add response to context
                 st.session_state.conversation_context.append({
                     'role': 'assistant',
-                    'content': response
+                    'content': response,
+                    'timestamp': datetime.now().isoformat()
+                })
+                
+                # Add to followups for UI display
+                if 'followups' not in st.session_state:
+                    st.session_state.followups = []
+                    
+                st.session_state.followups.append({
+                    'question': question,
+                    'response': response,
+                    'timestamp': datetime.now().isoformat()
                 })
                 
                 return {
                     'timestamp': datetime.now().isoformat(),
                     'question': question,
                     'analysis': response,
-                    'analysis_type': 'Follow-up'
+                    'analysis_type': 'Follow-up',
+                    'context_length': len(st.session_state.conversation_context)
                 }
             return None
         except Exception as e:
             st.error(f"Error handling follow-up question: {str(e)}")
             return None
 
+    def get_conversation_context(self):
+        """Get the current conversation context"""
+        return st.session_state.get('conversation_context', [])
+
     def create_comparison_prompt(self, individual_analyses, comparison_type, expertise_level="Novice"):
         """Create a prompt for comparing multiple charts"""
         expertise_adjustments = {