Update eb_agent_module.py

eb_agent_module.py

@@ -580,74 +580,74 @@ class EmployerBrandingAgent:
         else:
             return "general"
 
-async def _generate_pandas_response(self, query: str) -> tuple[str, bool]:
-    """Generate response using PandasAI for data queries"""
-    if not self.pandas_agent or not hasattr(self, 'pandas_dfs'):
-        return "Data analysis not available - PandasAI not initialized.", False
-    
-    try:
-        logging.info(f"Processing data query with PandasAI: {query[:100]}...")
-        
-        # Clear any existing matplotlib figures to avoid conflicts
-        import matplotlib.pyplot as plt
-        plt.clf()
-        plt.close('all')
-        
-        # Use the first available dataframe for single-df queries
-        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(query)
-        else:
-            # For multiple dataframes, use pai.chat with all dfs
-            dfs = list(self.pandas_dfs.values())
-            pandas_response = pai.chat(query, *dfs)
-        
-        # Handle different response types
-        response_text = ""
-        chart_info = ""
-        
-        # Check if response is a plot path or contains plot information
-        if isinstance(pandas_response, str) and pandas_response.endswith(('.png', '.jpg', '.jpeg', '.svg')):
-            # Response is a chart path
-            chart_info = f"\n\n📊 **Chart Generated**: {os.path.basename(pandas_response)}\nChart saved at: {pandas_response}"
-            response_text = "Analysis completed with visualization"
-            logging.info(f"Chart generated: {pandas_response}")
-        elif hasattr(pandas_response, 'plot_path') and pandas_response.plot_path:
-            # Response object has plot path
-            chart_info = f"\n\n📊 **Chart Generated**: {os.path.basename(pandas_response.plot_path)}\nChart saved at: {pandas_response.plot_path}"
-            response_text = getattr(pandas_response, 'text', str(pandas_response))
-            logging.info(f"Chart generated: {pandas_response.plot_path}")
-        else:
-            # Check for any new chart files in the charts directory
-            if os.path.exists(self.charts_dir):
-                chart_files = [f for f in os.listdir(self.charts_dir) if f.endswith(('.png', '.jpg', '.jpeg', '.svg'))]
-                if chart_files:
-                    # Get the most recent chart file
-                    chart_files.sort(key=lambda x: os.path.getmtime(os.path.join(self.charts_dir, x)), reverse=True)
-                    latest_chart = chart_files[0]
-                    chart_path = os.path.join(self.charts_dir, latest_chart)
-                    # Check if this chart was created in the last 30 seconds (likely from this query)
-                    import time
-                    if time.time() - os.path.getmtime(chart_path) < 30:
-                        chart_info = f"\n\n📊 **Chart Generated**: {latest_chart}\nChart saved at: {chart_path}"
-                        logging.info(f"Chart generated: {chart_path}")
-            
-            # Handle text response
-            if pandas_response and str(pandas_response).strip():
-                response_text = str(pandas_response).strip()
+    async def _generate_pandas_response(self, query: str) -> tuple[str, bool]:
+        """Generate response using PandasAI for data queries"""
+        if not self.pandas_agent or not hasattr(self, 'pandas_dfs'):
+            return "Data analysis not available - PandasAI not initialized.", False
+        
+        try:
+            logging.info(f"Processing data query with PandasAI: {query[:100]}...")
+            
+            # Clear any existing matplotlib figures to avoid conflicts
+            import matplotlib.pyplot as plt
+            plt.clf()
+            plt.close('all')
+            
+            # Use the first available dataframe for single-df queries
+            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(query)
             else:
-                response_text = "Analysis completed"
+                # For multiple dataframes, use pai.chat with all dfs
+                dfs = list(self.pandas_dfs.values())
+                pandas_response = pai.chat(query, *dfs)
+            
+            # Handle different response types
+            response_text = ""
+            chart_info = ""
+            
+            # Check if response is a plot path or contains plot information
+            if isinstance(pandas_response, str) and pandas_response.endswith(('.png', '.jpg', '.jpeg', '.svg')):
+                # Response is a chart path
+                chart_info = f"\n\n📊 **Chart Generated**: {os.path.basename(pandas_response)}\nChart saved at: {pandas_response}"
+                response_text = "Analysis completed with visualization"
+                logging.info(f"Chart generated: {pandas_response}")
+            elif hasattr(pandas_response, 'plot_path') and pandas_response.plot_path:
+                # Response object has plot path
+                chart_info = f"\n\n📊 **Chart Generated**: {os.path.basename(pandas_response.plot_path)}\nChart saved at: {pandas_response.plot_path}"
+                response_text = getattr(pandas_response, 'text', str(pandas_response))
+                logging.info(f"Chart generated: {pandas_response.plot_path}")
+            else:
+                # Check for any new chart files in the charts directory
+                if os.path.exists(self.charts_dir):
+                    chart_files = [f for f in os.listdir(self.charts_dir) if f.endswith(('.png', '.jpg', '.jpeg', '.svg'))]
+                    if chart_files:
+                        # Get the most recent chart file
+                        chart_files.sort(key=lambda x: os.path.getmtime(os.path.join(self.charts_dir, x)), reverse=True)
+                        latest_chart = chart_files[0]
+                        chart_path = os.path.join(self.charts_dir, latest_chart)
+                        # Check if this chart was created in the last 30 seconds (likely from this query)
+                        import time
+                        if time.time() - os.path.getmtime(chart_path) < 30:
+                            chart_info = f"\n\n📊 **Chart Generated**: {latest_chart}\nChart saved at: {chart_path}"
+                            logging.info(f"Chart generated: {chart_path}")
+                
+                # Handle text response
+                if pandas_response and str(pandas_response).strip():
+                    response_text = str(pandas_response).strip()
+                else:
+                    response_text = "Analysis completed"
+            
+            final_response = response_text + chart_info
+            return final_response, True
         
-        final_response = response_text + chart_info
-        return final_response, True
+        except Exception as e:
+            logging.error(f"Error in PandasAI processing: {e}", exc_info=True)
+            # Try to provide a more helpful error message
+            if "Invalid output" in str(e) and "plot save path" in str(e):
+                return "I tried to create a visualization but encountered a formatting issue. Please try rephrasing your request or ask for specific data without requesting a chart.", False
     
-    except Exception as e:
-        logging.error(f"Error in PandasAI processing: {e}", exc_info=True)
-        # Try to provide a more helpful error message
-        if "Invalid output" in str(e) and "plot save path" in str(e):
-            return "I tried to create a visualization but encountered a formatting issue. Please try rephrasing your request or ask for specific data without requesting a chart.", 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"""