Update eb_agent_module.py

eb_agent_module.py

@@ -351,7 +351,9 @@ class EmployerBrandingAgent:
         if not self.all_dataframes or not GEMINI_API_KEY:
             logging.warning("Cannot initialize PandasAI agent: missing dataframes or API key")
             return
-    
+            
+        self._preprocess_dataframes_for_pandas_ai()
+
         try:
             # Configure LiteLLM with Gemini
             llm = LiteLLM(
@@ -408,8 +410,11 @@ class EmployerBrandingAgent:
             description_parts.append("Key columns: " + "; ".join(column_descriptions))
         
         # Add specific context for employer branding
+        # Special handling for follower_stats
         if name.lower() in ['follower_stats', 'followers']:
-            description_parts.append("This data tracks LinkedIn company page follower growth and demographics for employer branding analysis.")
+            description_parts.append("This data tracks LinkedIn company page follower growth and demographics. For monthly growth data, use the 'extracted_date' column for date-based queries instead of trying to cast 'category_name' as a date.")
+            if 'extracted_date' in df.columns:
+                description_parts.append("The 'extracted_date' column contains properly formatted dates (YYYY-MM-DD) extracted from category_name for follower_gains_monthly records.")
         elif name.lower() in ['posts', 'post_stats']:
             description_parts.append("This data contains LinkedIn post performance metrics for employer branding content analysis.")
         elif name.lower() in ['mentions', 'brand_mentions']:
@@ -445,6 +450,37 @@ class EmployerBrandingAgent:
 
     def _get_dataframes_summary(self) -> str:
         return get_all_schemas_representation(self.all_dataframes)
+
+    def _preprocess_dataframes_for_pandas_ai(self):
+        """Preprocess dataframes to handle date casting issues before PandasAI analysis"""
+        if not self.all_dataframes:
+            return
+        
+        for name, df in self.all_dataframes.items():
+            if name.lower() in ['follower_stats', 'followers']:
+                # Create a copy to avoid modifying original data
+                df_copy = df.copy()
+                
+                # Handle category_name column that contains dates for follower_gains_monthly
+                if 'category_name' in df_copy.columns and 'follower_count_type' in df_copy.columns:
+                    # Create a proper date column for date-based queries
+                    def extract_date_from_category(row):
+                        if row.get('follower_count_type') == 'follower_gains_monthly':
+                            category_name = str(row.get('category_name', ''))
+                            # Check if it matches YYYY-MM-DD format
+                            import re
+                            date_pattern = r'^\d{4}-\d{2}-\d{2}$'
+                            if re.match(date_pattern, category_name):
+                                return category_name
+                        return None
+                    
+                    # Add extracted_date column for cleaner date operations
+                    df_copy['extracted_date'] = df_copy.apply(extract_date_from_category, axis=1)
+                    
+                    # Update the dataframe in our collection
+                    self.all_dataframes[name] = df_copy
+                    
+                logging.info(f"Preprocessed {name} dataframe for date handling")
     
     def _build_system_prompt(self) -> str:
         """Enhanced system prompt that works with PandasAI integration"""
@@ -544,7 +580,7 @@ class EmployerBrandingAgent:
             # For multi-df queries, you'd use pai.chat(query, df1, df2, ...)
             if len(self.pandas_dfs) == 1:
                 df = list(self.pandas_dfs.values())[0]
-                logging.info(f"Using single DataFrame for query: {df}")
+                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