Update eb_agent_module.py

eb_agent_module.py

@@ -649,113 +649,114 @@ class EmployerBrandingAgent:
         else:
             return "general"
 
-# Replace the _generate_pandas_response method and everything after it with this properly indented code:
-
-async def _generate_pandas_response(self, query: str) -> tuple[str, bool]:
-    """Generate response using PandasAI with enhanced error handling"""
-    if not self.pandas_agent or not hasattr(self, 'pandas_dfs'):
-        return "Data analysis not available - PandasAI not initialized.", False
+    # Replace the _generate_pandas_response method and everything after it with this properly indented code:
     
-    try:
-        logging.info(f"Processing data query with PandasAI: {query[:100]}...")
-        
-        # Clear any existing matplotlib figures
-        import matplotlib.pyplot as plt
-        plt.clf()
-        plt.close('all')
-        
-        # Enhanced query processing based on content
-        processed_query = query
-        
-        # Add helpful context for common chart requests
-        if any(word in query.lower() for word in ['chart', 'graph', 'plot', 'visualize']):
-            if 'monthly' in query.lower() and 'follower' in query.lower():
-                processed_query += """. 
-                Use the monthly gains data (follower_count_type='follower_gains_monthly') 
-                and use the extracted_date or month_name column for the x-axis. 
-                Make sure to filter out any null dates and sort by date.
-                Create a clear line chart showing the trend over time."""
-            elif 'cumulative' in query.lower() and 'follower' in query.lower():
-                processed_query += """. 
-                Use the cumulative data (follower_count_type='follower_count_cumulative') 
-                and create a chart showing the total follower growth over time."""
-        
-        # Execute the query
-        if len(self.pandas_dfs) == 1:
-            df = list(self.pandas_dfs.values())[0]
-            logging.info(f"Using single DataFrame for query with shape: {df.df.shape}")
-            pandas_response = df.chat(processed_query)
-        else:
-            dfs = list(self.pandas_dfs.values())
-            pandas_response = pai.chat(processed_query, *dfs)
-        
-        # Enhanced response processing
-        response_text = ""
-        chart_info = ""
-        
-        # Check for chart generation
-        chart_path = None
-        
-        # Method 1: Direct path response
-        if isinstance(pandas_response, str) and pandas_response.endswith(('.png', '.jpg', '.jpeg', '.svg')):
-            chart_path = pandas_response
-            response_text = "Analysis completed with visualization"
+    async def _generate_pandas_response(self, query: str) -> tuple[str, bool]:
+        """Generate response using PandasAI with enhanced error handling"""
+        if not self.pandas_agent or not hasattr(self, 'pandas_dfs'):
+            return "Data analysis not available - PandasAI not initialized.", False
         
-        # Method 2: Response object with plot path
-        elif hasattr(pandas_response, 'plot_path') and pandas_response.plot_path:
-            chart_path = pandas_response.plot_path
-            response_text = getattr(pandas_response, 'text', str(pandas_response))
-        
-        # Method 3: Check charts directory for new files
-        else:
-            if os.path.exists(self.charts_dir):
-                # Get all chart files sorted by modification time
-                chart_files = []
-                for f in os.listdir(self.charts_dir):
-                    if f.endswith(('.png', '.jpg', '.jpeg', '.svg')):
-                        full_path = os.path.join(self.charts_dir, f)
-                        chart_files.append((full_path, os.path.getmtime(full_path)))
-                
-                if chart_files:
-                    # Sort by modification time (newest first)
-                    chart_files.sort(key=lambda x: x[1], reverse=True)
-                    latest_chart_path, latest_time = chart_files[0]
-                    
-                    # Check if created in last 60 seconds
-                    import time
-                    if time.time() - latest_time < 60:
-                        chart_path = latest_chart_path
-                        logging.info(f"Found recent chart: {chart_path}")
+        try:
+            logging.info(f"Processing data query with PandasAI: {query[:100]}...")
             
-            # Handle text response
-            if pandas_response and str(pandas_response).strip():
-                response_text = str(pandas_response).strip()
+            # Clear any existing matplotlib figures
+            import matplotlib.pyplot as plt
+            plt.clf()
+            plt.close('all')
+            
+            # Enhanced query processing based on content
+            processed_query = query
+            
+            # Add helpful context for common chart requests
+            if any(word in query.lower() for word in ['chart', 'graph', 'plot', 'visualize']):
+                if 'monthly' in query.lower() and 'follower' in query.lower():
+                    processed_query += """. 
+                    Use the monthly gains data (follower_count_type='follower_gains_monthly') 
+                    and use the extracted_date or month_name column for the x-axis. 
+                    Make sure to filter out any null dates and sort by date.
+                    Create a clear line chart showing the trend over time."""
+                elif 'cumulative' in query.lower() and 'follower' in query.lower():
+                    processed_query += """. 
+                    Use the cumulative data (follower_count_type='follower_count_cumulative') 
+                    and create a chart showing the total follower growth over time."""
+            
+            # Execute the query
+            if len(self.pandas_dfs) == 1:
+                df = list(self.pandas_dfs.values())[0]
+                logging.info(f"Using single DataFrame for query with shape: {df.df.shape}")
+                pandas_response = df.chat(processed_query)
             else:
-                response_text = "Analysis completed"
-        
-        # Format final response
-        if chart_path and os.path.exists(chart_path):
-            chart_info = f"\n\n📊 **Chart Generated**: {os.path.basename(chart_path)}\nChart saved at: {chart_path}"
-            logging.info(f"Chart successfully generated: {chart_path}")
-        
-        final_response = response_text + chart_info
-        success = True
-        
-        return final_response, success
-    
-    except Exception as e:
-        logging.error(f"Error in PandasAI processing: {e}", exc_info=True)
+                dfs = list(self.pandas_dfs.values())
+                pandas_response = pai.chat(processed_query, *dfs)
+            
+            # Enhanced response processing
+            response_text = ""
+            chart_info = ""
+            
+            # Check for chart generation
+            chart_path = None
+            
+            # Method 1: Direct path response
+            if isinstance(pandas_response, str) and pandas_response.endswith(('.png', '.jpg', '.jpeg', '.svg')):
+                chart_path = pandas_response
+                response_text = "Analysis completed with visualization"
+            
+            # Method 2: Response object with plot path
+            elif hasattr(pandas_response, 'plot_path') and pandas_response.plot_path:
+                chart_path = pandas_response.plot_path
+                response_text = getattr(pandas_response, 'text', str(pandas_response))
+            
+            # Method 3: Check charts directory for new files
+            else:
+                if os.path.exists(self.charts_dir):
+                    # Get all chart files sorted by modification time
+                    chart_files = []
+                    for f in os.listdir(self.charts_dir):
+                        if f.endswith(('.png', '.jpg', '.jpeg', '.svg')):
+                            full_path = os.path.join(self.charts_dir, f)
+                            chart_files.append((full_path, os.path.getmtime(full_path)))
+                    
+                    if chart_files:
+                        # Sort by modification time (newest first)
+                        chart_files.sort(key=lambda x: x[1], reverse=True)
+                        latest_chart_path, latest_time = chart_files[0]
+                        
+                        # Check if created in last 60 seconds
+                        import time
+                        if time.time() - latest_time < 60:
+                            chart_path = latest_chart_path
+                            logging.info(f"Found recent chart: {chart_path}")
+                
+                # Handle text response
+                if pandas_response and str(pandas_response).strip():
+                    response_text = str(pandas_response).strip()
+                else:
+                    response_text = "Analysis completed"
+            
+            # Format final response
+            if chart_path and os.path.exists(chart_path):
+                chart_info = f"\n\n📊 **Chart Generated**: {os.path.basename(chart_path)}\nChart saved at: {chart_path}"
+                logging.info(f"Chart successfully generated: {chart_path}")
+            
+            final_response = response_text + chart_info
+            success = True
+            
+            return final_response, success
         
-        # Enhanced error handling
-        error_str = str(e).lower()
-        if "matplotlib" in error_str and "none" in error_str:
-            return "I encountered a data visualization error. This might be due to missing or null values in your data. Please try asking for the raw data first, or specify which specific columns you'd like to analyze.", False
-        elif "strftime" in error_str:
-            return "I encountered a date formatting issue. Please try asking for the data without specific date formatting, or ask me to show the raw data structure first.", False
-        elif "ambiguous" in error_str:
-            return "I encountered an ambiguous data type issue. Please try being more specific about which data you'd like to analyze (e.g., 'show monthly follower gains' vs 'show cumulative followers').", False
-        else:
-            return f"Error processing data query: {str(e)}", False
+        except Exception as e:
+            logging.error(f"Error in PandasAI processing: {e}", exc_info=True)
+            
+            # Enhanced error handling
+            error_str = str(e).lower()
+            if "matplotlib" in error_str and "none" in error_str:
+                return "I encountered a data visualization error. This might be due to missing or null values in your data. Please try asking for the raw data first, or specify which specific columns you'd like to analyze.", False
+            elif "strftime" in error_str:
+                return "I encountered a date formatting issue. Please try asking for the data without specific date formatting, or ask me to show the raw data structure first.", False
+            elif "ambiguous" in error_str:
+                return "I encountered an ambiguous data type issue. Please try being more specific about which data you'd like to analyze (e.g., 'show monthly follower gains' vs 'show cumulative followers').", False
+            else:
+                return f"Error processing data query: {str(e)}", False
+
 
     async def _generate_enhanced_response(self, query: str, pandas_result: str = "", query_type: str = "general") -> str:
         """Generate enhanced response combining PandasAI results with RAG context"""