Update agents/analytics_agent.py

agents/analytics_agent.py

@@ -2,6 +2,7 @@ import os
 import json
 import pandas as pd
 import numpy as np
+import sys
 from typing import Dict, List, Any, Tuple, Optional
 from pydantic import BaseModel, Field
 from langchain_anthropic import ChatAnthropic
@@ -11,6 +12,11 @@ import re
 import matplotlib.pyplot as plt
 import seaborn as sns
 from io import StringIO
+import logging
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger("analytics_agent")
 
 class AnalysisRequest(BaseModel):
     """Structure for an analysis request"""
@@ -125,6 +131,8 @@ Please generate Python code to perform this analysis.
         
         # In-memory storage for analysis artifacts
         self.analysis_artifacts = {}
+        
+        logger.info("Analytics Agent initialized successfully")
     
     def extract_python_from_response(self, response: str) -> str:
         """Extract Python code from LLM response"""
@@ -139,6 +147,7 @@ Please generate Python code to perform this analysis.
             return python_match.group(1).strip()
         
         # If all else fails, return empty string
+        logger.warning("No code block found in response")
         return ""
     
     def extract_insights_from_code_output(self, output: Dict[str, Any]) -> Tuple[List[Dict[str, Any]], Dict[str, float], float]:
@@ -149,104 +158,207 @@ Please generate Python code to perform this analysis.
         
         return insights, attribution, confidence
     
-    def perform_analysis(self, request: AnalysisRequest, data_sources: Dict[str, Any]) -> AnalysisResult:
-        """Perform analysis based on request and return results"""
-        print(f"Analytics Agent: Performing {request.analysis_type} analysis - {request.description}")
-        
-        # Format data sources description for the prompt
-        data_sources_desc = ""
-        for source_id, source in data_sources.items():
-            df = source.content
-            data_sources_desc += f"Data source '{source_id}' ({source.name}):\n"
-            data_sources_desc += f"- Shape: {df.shape[0]} rows, {df.shape[1]} columns\n"
-            data_sources_desc += f"- Columns: {', '.join(df.columns)}\n"
-            data_sources_desc += f"- Sample data:\n{df.head(3).to_string()}\n\n"
+    def generate_default_analysis(self, request: AnalysisRequest, data_sources: Dict[str, Any]) -> Dict[str, Any]:
+        """Generate a default analysis output when code execution fails"""
+        logger.info(f"Generating default analysis for {request.description}")
         
-        # Format the request for the prompt
-        request_data = {
-            "description": request.description,
-            "data_sources": data_sources_desc,
-            "analysis_type": request.analysis_type,
-            "parameters": json.dumps(request.parameters, indent=2) if request.parameters else "None",
-            "purpose": request.purpose
-        }
+        # Create default insights based on request description
+        insights = [
+            {
+                "finding": f"Analysis of {request.description}",
+                "details": "Default analysis created due to execution issues",
+                "impact": "Recommend manual investigation"
+            }
+        ]
         
-        # Generate analysis code
-        response = self.analysis_chain.invoke(request_data)
+        # Create default attribution
+        attribution = {"unknown_factors": 1.0}
         
-        # Extract Python code
-        python_code = self.extract_python_from_response(response)
+        # Default metrics
+        metrics = {"analysis_completion": 0.0}
         
-        # Execute analysis (with safety checks)
-        insights = []
-        attribution = {}
-        confidence = 0.0
-        visualizations = []
-        metrics = {}
+        return {
+            "insights": insights,
+            "attribution": attribution,
+            "metrics": metrics,
+            "visualizations": [],
+            "confidence": 0.5
+        }
+    
+    def perform_analysis(self, request: AnalysisRequest, data_sources: Dict[str, Any]) -> AnalysisResult:
+        """Perform analysis based on request and return results"""
+        logger.info(f"Analytics Agent: Performing {request.analysis_type} analysis - {request.description}")
         
-        if not python_code:
-            print("Warning: No analysis code generated.")
-        else:
-            try:
-                # Prepare data sources for the analysis
-                analysis_data_sources = {src_id: src.content for src_id, src in data_sources.items()}
-                
-                # Create a local namespace with access to pandas, numpy, etc.
-                local_namespace = {
-                    "pd": pd,
-                    "np": np,
-                    "plt": plt,
-                    "sns": sns,
-                    "data_sources": analysis_data_sources
-                }
+        try:
+            # Format data sources description for the prompt
+            data_sources_desc = ""
+            for source_id, source in data_sources.items():
+                if not hasattr(source, 'content') or source.content is None:
+                    logger.warning(f"Data source {source_id} has no content attribute or content is None")
+                    continue
                 
-                # Capture print outputs
-                original_stdout = sys.stdout
-                sys.stdout = mystdout = StringIO()
-                
-                # Execute the code
-                exec(python_code, local_namespace)
-                
-                # Restore stdout
-                sys.stdout = original_stdout
-                print_output = mystdout.getvalue()
-                
-                # Look for a run_analysis function and execute it
-                if "run_analysis" in local_namespace:
-                    analysis_output = local_namespace["run_analysis"](analysis_data_sources)
-                    if isinstance(analysis_output, dict):
-                        insights = analysis_output.get("insights", [])
-                        attribution = analysis_output.get("attribution", {})
-                        confidence = analysis_output.get("confidence", 0.0)
-                        metrics = analysis_output.get("metrics", {})
-                        visualizations = analysis_output.get("visualizations", [])
+                df = source.content
+                data_sources_desc += f"Data source '{source_id}' ({source.name}):\n"
+                data_sources_desc += f"- Shape: {df.shape[0]} rows, {df.shape[1]} columns\n"
+                data_sources_desc += f"- Columns: {', '.join(df.columns)}\n"
+                data_sources_desc += f"- Sample data:\n{df.head(3).to_string()}\n\n"
+            
+            # Format the request for the prompt
+            request_data = {
+                "description": request.description,
+                "data_sources": data_sources_desc,
+                "analysis_type": request.analysis_type,
+                "parameters": json.dumps(request.parameters, indent=2) if request.parameters else "None",
+                "purpose": request.purpose
+            }
+            
+            # Generate analysis code
+            logger.info("Generating analysis code")
+            response = self.analysis_chain.invoke(request_data)
+            
+            # Extract Python code
+            python_code = self.extract_python_from_response(response)
+            
+            # Initialize default values
+            insights = []
+            attribution = {}
+            confidence = 0.0
+            visualizations = []
+            metrics = {}
+            
+            if not python_code:
+                logger.warning("No analysis code generated. Using default analysis.")
+                default_analysis = self.generate_default_analysis(request, data_sources)
+                insights = default_analysis["insights"]
+                attribution = default_analysis["attribution"]
+                confidence = default_analysis["confidence"]
+                metrics = default_analysis["metrics"]
+            else:
+                try:
+                    # Prepare data sources for the analysis
+                    analysis_data_sources = {}
+                    for src_id, src in data_sources.items():
+                        if hasattr(src, 'content') and src.content is not None:
+                            analysis_data_sources[src_id] = src.content
+                    
+                    if not analysis_data_sources:
+                        logger.warning("No valid data sources found for analysis")
+                        default_analysis = self.generate_default_analysis(request, data_sources)
+                        insights = default_analysis["insights"]
+                        attribution = default_analysis["attribution"]
+                        confidence = default_analysis["confidence"]
+                        metrics = default_analysis["metrics"]
+                    else:
+                        # Create a local namespace with access to pandas, numpy, etc.
+                        local_namespace = {
+                            "pd": pd,
+                            "np": np,
+                            "plt": plt,
+                            "sns": sns,
+                            "data_sources": analysis_data_sources
+                        }
                         
-                        # Store any figures in the local namespace as base64 encoded images
-                        for var_name, var_value in local_namespace.items():
-                            if isinstance(var_value, plt.Figure):
-                                fig_filename = f"figure_{request.request_id}_{var_name}.png"
-                                var_value.savefig(fig_filename)
-                                self.analysis_artifacts[fig_filename] = fig_filename
-                                visualizations.append(fig_filename)
-                
-            except Exception as e:
-                print(f"Analysis execution error: {e}")
-        
-        # Create analysis result
-        result = AnalysisResult(
-            result_id=f"analysis_{request.request_id}",
-            name=f"Analysis of {request.description}",
-            description=request.description,
-            analysis_type=request.analysis_type,
-            code=python_code,
-            visualizations=visualizations,
-            insights=insights,
-            metrics=metrics,
-            attribution=attribution,
-            confidence=confidence
-        )
-        
-        return result
+                        # Capture print outputs
+                        stdout_backup = sys.stdout
+                        sys.stdout = mystdout = StringIO()
+                        
+                        # Execute the code
+                        logger.info("Executing analysis code")
+                        exec(python_code, local_namespace)
+                        
+                        # Restore stdout
+                        sys.stdout = stdout_backup
+                        print_output = mystdout.getvalue()
+                        logger.debug(f"Code execution output: {print_output}")
+                        
+                        # Look for a run_analysis function and execute it
+                        if "run_analysis" in local_namespace:
+                            logger.info("Running analysis function")
+                            analysis_output = local_namespace["run_analysis"](analysis_data_sources)
+                            
+                            if isinstance(analysis_output, dict):
+                                insights = analysis_output.get("insights", [])
+                                attribution = analysis_output.get("attribution", {})
+                                confidence = analysis_output.get("confidence", 0.0)
+                                metrics = analysis_output.get("metrics", {})
+                                visualizations = analysis_output.get("visualizations", [])
+                                
+                                # Store any figures in the local namespace as base64 encoded images
+                                for var_name, var_value in local_namespace.items():
+                                    if isinstance(var_value, plt.Figure):
+                                        fig_filename = f"figure_{request.request_id}_{var_name}.png"
+                                        var_value.savefig(fig_filename)
+                                        self.analysis_artifacts[fig_filename] = fig_filename
+                                        visualizations.append(fig_filename)
+                            else:
+                                logger.warning(f"run_analysis returned non-dict type: {type(analysis_output)}")
+                                default_analysis = self.generate_default_analysis(request, data_sources)
+                                insights = default_analysis["insights"]
+                                attribution = default_analysis["attribution"]
+                                confidence = default_analysis["confidence"]
+                                metrics = default_analysis["metrics"]
+                        else:
+                            logger.warning("No run_analysis function found in generated code")
+                            # Generate a minimal default analysis
+                            default_analysis = self.generate_default_analysis(request, data_sources)
+                            insights = default_analysis["insights"]
+                            attribution = default_analysis["attribution"]
+                            confidence = default_analysis["confidence"]
+                            metrics = default_analysis["metrics"]
+                            
+                except Exception as e:
+                    logger.error(f"Analysis execution error: {e}", exc_info=True)
+                    logger.error(f"Python code that failed: {python_code}")
+                    
+                    # Generate a minimal default analysis on execution failure
+                    default_analysis = self.generate_default_analysis(request, data_sources)
+                    insights = default_analysis["insights"]
+                    attribution = default_analysis["attribution"]
+                    confidence = default_analysis["confidence"]
+                    metrics = default_analysis["metrics"]
+            
+            # Ensure we have at least one insight
+            if not insights:
+                insights = [{"finding": "No specific insights found", "details": "Analysis completed but no significant patterns were identified", "impact": "No immediate action required"}]
+            
+            # Ensure we have attribution
+            if not attribution:
+                attribution = {"unattributed_factors": 1.0}
+            
+            # Create analysis result
+            result = AnalysisResult(
+                result_id=f"analysis_{request.request_id}",
+                name=f"Analysis of {request.description}",
+                description=request.description,
+                analysis_type=request.analysis_type,
+                code=python_code,
+                visualizations=visualizations,
+                insights=insights,
+                metrics=metrics,
+                attribution=attribution,
+                confidence=confidence
+            )
+            
+            logger.info(f"Analysis for {request.description} completed successfully")
+            return result
+            
+        except Exception as e:
+            logger.error(f"Error in perform_analysis: {e}", exc_info=True)
+            
+            # Create a fallback analysis result on error
+            default_analysis = self.generate_default_analysis(request, data_sources)
+            
+            return AnalysisResult(
+                result_id=f"analysis_{request.request_id}",
+                name=f"Analysis of {request.description} (Error)",
+                description=request.description,
+                analysis_type=request.analysis_type,
+                code="# Error during analysis",
+                insights=default_analysis["insights"],
+                metrics=default_analysis["metrics"],
+                attribution=default_analysis["attribution"],
+                confidence=default_analysis["confidence"]
+            )
 
 # For testing
 if __name__ == "__main__":