Update sozo_gen.py

sozo_gen.py

@@ -310,11 +310,11 @@ def perform_autonomous_data_analysis(df: pd.DataFrame, user_ctx: str, filename:
     """
     logging.info("Performing autonomous data analysis...")
     
-    # Basic data profiling
+    # Basic data profiling with JSON-safe types
     basic_info = {
         "shape": df.shape,
         "columns": list(df.columns),
-        "dtypes": df.dtypes.to_dict(),
+        "dtypes": {col: str(dtype) for col, dtype in df.dtypes.items()},
         "filename": filename,
         "user_context": user_ctx
     }
@@ -388,11 +388,11 @@ def classify_dataset_domain(df: pd.DataFrame, filename: str) -> Dict[str, Any]:
     
     return {
         "primary_domain": primary_domain,
-        "domain_confidence": domain_scores.get(primary_domain, 0),
-        "domain_scores": domain_scores,
+        "domain_confidence": int(domain_scores.get(primary_domain, 0)),
+        "domain_scores": {k: int(v) for k, v in domain_scores.items()},
         "data_characteristics": {
-            "numeric_ratio": numeric_ratio,
-            "categorical_ratio": categorical_ratio,
+            "numeric_ratio": float(numeric_ratio),
+            "categorical_ratio": float(categorical_ratio),
             "is_time_series": detect_time_series(df),
             "is_transactional": detect_transactional_data(df),
             "is_experimental": detect_experimental_data(df)
@@ -412,12 +412,20 @@ def generate_statistical_profile(df: pd.DataFrame) -> Dict[str, Any]:
         "missing_data": {}
     }
     
-    # Summary statistics for numeric columns
+    # Summary statistics for numeric columns with JSON-safe conversion
     numeric_cols = df.select_dtypes(include=[np.number]).columns
     if len(numeric_cols) > 0:
-        profile["summary_stats"] = df[numeric_cols].describe().to_dict()
+        desc_stats = df[numeric_cols].describe()
+        # Convert to JSON-safe format
+        profile["summary_stats"] = {
+            col: {
+                stat: float(val) if pd.notna(val) else None 
+                for stat, val in desc_stats[col].items()
+            }
+            for col in desc_stats.columns
+        }
         
-        # Correlation analysis
+        # Correlation analysis with JSON-safe conversion
         if len(numeric_cols) > 1:
             corr_matrix = df[numeric_cols].corr()
             # Find strong correlations
@@ -425,11 +433,11 @@ def generate_statistical_profile(df: pd.DataFrame) -> Dict[str, Any]:
             for i in range(len(corr_matrix.columns)):
                 for j in range(i+1, len(corr_matrix.columns)):
                     corr_val = corr_matrix.iloc[i, j]
-                    if abs(corr_val) > 0.7:  # Strong correlation threshold
+                    if abs(corr_val) > 0.7 and pd.notna(corr_val):  # Strong correlation threshold
                         strong_corrs.append({
                             "var1": corr_matrix.columns[i],
                             "var2": corr_matrix.columns[j],
-                            "correlation": corr_val
+                            "correlation": float(corr_val)
                         })
             profile["correlations"] = {"strong_correlations": strong_corrs}
     
@@ -438,16 +446,22 @@ def generate_statistical_profile(df: pd.DataFrame) -> Dict[str, Any]:
     if len(categorical_cols) > 0:
         profile["categorical_analysis"] = {}
         for col in categorical_cols:
+            value_counts = df[col].value_counts().head(5)
             profile["categorical_analysis"][col] = {
-                "unique_count": df[col].nunique(),
-                "top_values": df[col].value_counts().head(5).to_dict()
+                "unique_count": int(df[col].nunique()),
+                "top_values": {str(k): int(v) for k, v in value_counts.items()}
             }
     
-    # Missing data analysis
+    # Missing data analysis with JSON-safe conversion
     missing_data = df.isnull().sum()
+    missing_dict = {}
+    for col, missing_count in missing_data.items():
+        if missing_count > 0:
+            missing_dict[col] = int(missing_count)
+    
     profile["missing_data"] = {
-        "columns_with_missing": missing_data[missing_data > 0].to_dict(),
-        "total_missing_percentage": (df.isnull().sum().sum() / (len(df) * len(df.columns))) * 100
+        "columns_with_missing": missing_dict,
+        "total_missing_percentage": float((df.isnull().sum().sum() / (len(df) * len(df.columns))) * 100)
     }
     
     return profile
@@ -472,11 +486,11 @@ def discover_data_relationships(df: pd.DataFrame) -> Dict[str, Any]:
             for col2 in numeric_cols:
                 if col1 != col2:
                     correlation = df[col1].corr(df[col2])
-                    if abs(correlation) > 0.5:  # Moderate to strong correlation
+                    if abs(correlation) > 0.5 and pd.notna(correlation):  # Moderate to strong correlation
                         relationships["key_relationships"].append({
                             "variable1": col1,
                             "variable2": col2,
-                            "relationship_strength": correlation,
+                            "relationship_strength": float(correlation),
                             "relationship_type": "positive" if correlation > 0 else "negative"
                         })
     
@@ -489,8 +503,8 @@ def discover_data_relationships(df: pd.DataFrame) -> Dict[str, Any]:
                 relationships["patterns"].append({
                     "column": col,
                     "pattern_type": "categorical_distribution",
-                    "dominant_category": value_counts.index[0],
-                    "dominance_percentage": (value_counts.iloc[0] / len(df)) * 100
+                    "dominant_category": str(value_counts.index[0]),
+                    "dominance_percentage": float((value_counts.iloc[0] / len(df)) * 100)
                 })
     
     return relationships
@@ -526,7 +540,7 @@ def analyze_temporal_patterns(df: pd.DataFrame) -> Dict[str, Any]:
                 "start": df_temp[primary_date_col].min().strftime('%Y-%m-%d'),
                 "end": df_temp[primary_date_col].max().strftime('%Y-%m-%d')
             },
-            "time_span_days": (df_temp[primary_date_col].max() - df_temp[primary_date_col].min()).days,
+            "time_span_days": int((df_temp[primary_date_col].max() - df_temp[primary_date_col].min()).days),
             "frequency": detect_temporal_frequency(df_temp[primary_date_col])
         }
     
@@ -544,16 +558,16 @@ def assess_data_quality(df: pd.DataFrame) -> Dict[str, Any]:
         "data_consistency": {}
     }
     
-    # Completeness assessment
-    completeness = (1 - df.isnull().sum().sum() / (len(df) * len(df.columns))) * 100
+    # Completeness assessment with JSON-safe conversion
+    completeness = float((1 - df.isnull().sum().sum() / (len(df) * len(df.columns))) * 100)
     quality_metrics["data_completeness"] = completeness
     
     # Identify quality issues
     if completeness < 95:
         quality_metrics["quality_issues"].append("Missing data detected")
     
-    # Check for duplicates
-    duplicate_rows = df.duplicated().sum()
+    # Check for duplicates with JSON-safe conversion
+    duplicate_rows = int(df.duplicated().sum())
     if duplicate_rows > 0:
         quality_metrics["quality_issues"].append(f"{duplicate_rows} duplicate rows found")
     
@@ -563,11 +577,11 @@ def assess_data_quality(df: pd.DataFrame) -> Dict[str, Any]:
             if df[col].str.isnumeric().any() and not df[col].str.isnumeric().all():
                 quality_metrics["quality_issues"].append(f"Inconsistent data types in {col}")
     
-    # Calculate overall quality score
-    base_score = 100
+    # Calculate overall quality score with JSON-safe conversion
+    base_score = 100.0
     base_score -= (100 - completeness) * 0.5  # Penalize missing data
     base_score -= len(quality_metrics["quality_issues"]) * 5  # Penalize each quality issue
-    quality_metrics["overall_quality_score"] = max(0, base_score)
+    quality_metrics["overall_quality_score"] = float(max(0, base_score))
     
     return quality_metrics
 
@@ -839,6 +853,235 @@ def generate_original_report(df: pd.DataFrame, llm, ctx: str, uid: str, project_
     return {"raw_md": md, "chartUrls": chart_urls}
 
 
+def generate_fallback_report(autonomous_context: Dict[str, Any]) -> str:
+    """
+    Generates a basic fallback report when enhanced generation fails.
+    """
+    basic_info = autonomous_context["basic_info"]
+    domain = autonomous_context["domain"]["primary_domain"]
+    
+    return f"""
+# What This Data Reveals
+
+Looking at this {domain} dataset with {basic_info['shape'][0]} records, there are several key insights worth highlighting.
+
+## The Numbers Tell a Story
+
+This dataset contains {basic_info['shape'][1]} different variables, suggesting a comprehensive view of the underlying processes or behaviors being measured.
+
+<generate_chart: "bar | Data overview showing key metrics">
+
+## What You Should Know
+
+The data structure and patterns suggest this is worth deeper investigation. The variety of data types and relationships indicate multiple analytical opportunities.
+
+## Next Steps
+
+Based on this initial analysis, I recommend diving deeper into the specific patterns and relationships within the data to unlock more actionable insights.
+
+*Note: This is a simplified analysis. Enhanced storytelling temporarily unavailable.*
+"""
+# Removed - no longer needed since we're letting AI decide everything organically
+
+
+def generate_autonomous_charts(llm, df: pd.DataFrame, report_md: str, uid: str, project_id: str, bucket) -> Dict[str, str]:
+    """
+    Generates charts autonomously based on the report content and data characteristics.
+    """
+    # Extract chart descriptions from the enhanced report
+    chart_descs = extract_chart_tags(report_md)[:MAX_CHARTS]
+    chart_urls = {}
+    
+    if not chart_descs:
+        # If no charts specified, generate intelligent defaults
+        chart_descs = generate_intelligent_chart_suggestions(df, llm)
+    
+    chart_generator = ChartGenerator(llm, df)
+    
+    for desc in chart_descs:
+        try:
+            # Create a safe key for Firebase
+            safe_desc = sanitize_for_firebase_key(desc)
+            
+            # Replace chart tags in markdown
+            report_md = report_md.replace(f'<generate_chart: "{desc}">', f'<generate_chart: "{safe_desc}">')
+            report_md = report_md.replace(f'<generate_chart: {desc}>', f'<generate_chart: "{safe_desc}">')
+            
+            # Generate chart
+            with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+                img_path = Path(temp_file.name)
+                try:
+                    chart_spec = chart_generator.generate_chart_spec(desc)
+                    if execute_chart_spec(chart_spec, df, img_path):
+                        blob_name = f"sozo_projects/{uid}/{project_id}/charts/{uuid.uuid4().hex}.png"
+                        blob = bucket.blob(blob_name)
+                        blob.upload_from_filename(str(img_path))
+                        
+                        chart_urls[safe_desc] = blob.public_url
+                        logging.info(f"Generated autonomous chart: {safe_desc}")
+                finally:
+                    if os.path.exists(img_path):
+                        os.unlink(img_path)
+                        
+        except Exception as e:
+            logging.error(f"Failed to generate chart '{desc}': {str(e)}")
+            continue
+    
+    return chart_urls
+
+
+def generate_intelligent_chart_suggestions(df: pd.DataFrame, llm) -> List[str]:
+    """
+    Generates intelligent chart suggestions based on data characteristics.
+    """
+    numeric_cols = df.select_dtypes(include=[np.number]).columns
+    categorical_cols = df.select_dtypes(include=['object']).columns
+    
+    suggestions = []
+    
+    # Time series chart if temporal data exists
+    if detect_time_series(df):
+        suggestions.append("line | Time series trend analysis | Show temporal patterns")
+    
+    # Distribution chart for numeric data
+    if len(numeric_cols) > 0:
+        main_numeric = numeric_cols[0]
+        suggestions.append(f"hist | Distribution of {main_numeric} | Understand data distribution")
+    
+    # Correlation analysis if multiple numeric columns
+    if len(numeric_cols) > 1:
+        suggestions.append("scatter | Correlation analysis | Identify relationships between variables")
+    
+    # Categorical breakdown
+    if len(categorical_cols) > 0:
+        main_categorical = categorical_cols[0]
+        suggestions.append(f"bar | {main_categorical} breakdown | Show categorical distribution")
+    
+    return suggestions[:MAX_CHARTS]
+
+
+# Helper functions (preserve existing functionality)
+def detect_time_series(df: pd.DataFrame) -> bool:
+    """Detect if dataset contains time series data."""
+    for col in df.columns:
+        if 'date' in col.lower() or 'time' in col.lower():
+            return True
+        try:
+            pd.to_datetime(df[col])
+            return True
+        except:
+            continue
+    return False
+
+
+def detect_transactional_data(df: pd.DataFrame) -> bool:
+    """Detect if dataset contains transactional data."""
+    transaction_indicators = ['transaction', 'payment', 'order', 'invoice', 'amount', 'quantity']
+    columns_lower = [col.lower() for col in df.columns]
+    return any(indicator in col for col in columns_lower for indicator in transaction_indicators)
+
+
+def detect_experimental_data(df: pd.DataFrame) -> bool:
+    """Detect if dataset contains experimental data."""
+    experimental_indicators = ['test', 'experiment', 'trial', 'group', 'treatment', 'control']
+    columns_lower = [col.lower() for col in df.columns]
+    return any(indicator in col for col in columns_lower for indicator in experimental_indicators)
+
+
+def detect_temporal_frequency(date_series: pd.Series) -> str:
+    """Detect the frequency of temporal data."""
+    if len(date_series) < 2:
+        return "insufficient_data"
+    
+    # Calculate time differences
+    time_diffs = date_series.sort_values().diff().dropna()
+    median_diff = time_diffs.median()
+    
+    if median_diff <= pd.Timedelta(days=1):
+        return "daily"
+    elif median_diff <= pd.Timedelta(days=7):
+        return "weekly"
+    elif median_diff <= pd.Timedelta(days=31):
+        return "monthly"
+    else:
+        return "irregular"
+
+
+def determine_analysis_complexity(df: pd.DataFrame, domain_analysis: Dict[str, Any]) -> str:
+    """Determine the complexity level of analysis required."""
+    complexity_factors = 0
+    
+    # Data size factor
+    if len(df) > 10000:
+        complexity_factors += 1
+    if len(df.columns) > 20:
+        complexity_factors += 1
+    
+    # Data type diversity
+    if len(df.select_dtypes(include=[np.number]).columns) > 5:
+        complexity_factors += 1
+    if len(df.select_dtypes(include=['object']).columns) > 5:
+        complexity_factors += 1
+    
+    # Domain complexity
+    if domain_analysis["primary_domain"] in ["scientific", "financial"]:
+        complexity_factors += 1
+    
+    if complexity_factors >= 3:
+        return "high"
+    elif complexity_factors >= 2:
+        return "medium"
+    else:
+        return "low"
+
+
+def generate_original_report(df: pd.DataFrame, llm, ctx: str, uid: str, project_id: str, bucket) -> Dict[str, str]:
+    """
+    Fallback to original report generation logic if enhanced version fails.
+    """
+    logging.info("Using fallback report generation")
+    
+    # Original logic preserved
+    ctx_dict = {"shape": df.shape, "columns": list(df.columns), "user_ctx": ctx}
+    enhanced_ctx = enhance_data_context(df, ctx_dict)
+    
+    report_prompt = f"""
+    You are a senior data analyst and business intelligence expert. Analyze the provided dataset and write a comprehensive executive-level Markdown report.
+    **Dataset Analysis Context:** {json.dumps(enhanced_ctx, indent=2)}
+    **Instructions:**
+    1. **Executive Summary**: Start with a high-level summary of key findings.
+    2. **Key Insights**: Provide 3-5 key insights, each with its own chart tag.
+    3. **Visual Support**: Insert chart tags like: `<generate_chart: "chart_type | specific description">`.
+       Valid chart types: bar, pie, line, scatter, hist.
+    Generate insights that would be valuable to C-level executives.
+    """
+    
+    md = llm.invoke(report_prompt).content
+    chart_descs = extract_chart_tags(md)[:MAX_CHARTS]
+    chart_urls = {}
+    chart_generator = ChartGenerator(llm, df)
+
+    for desc in chart_descs:
+        safe_desc = sanitize_for_firebase_key(desc)
+        md = md.replace(f'<generate_chart: "{desc}">', f'<generate_chart: "{safe_desc}">')
+        md = md.replace(f'<generate_chart: {desc}>', f'<generate_chart: "{safe_desc}">')
+
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as temp_file:
+            img_path = Path(temp_file.name)
+            try:
+                chart_spec = chart_generator.generate_chart_spec(desc)
+                if execute_chart_spec(chart_spec, df, img_path):
+                    blob_name = f"sozo_projects/{uid}/{project_id}/charts/{uuid.uuid4().hex}.png"
+                    blob = bucket.blob(blob_name)
+                    blob.upload_from_filename(str(img_path))
+                    chart_urls[safe_desc] = blob.public_url
+            finally:
+                if os.path.exists(img_path):
+                    os.unlink(img_path)
+
+    return {"raw_md": md, "chartUrls": chart_urls}
+
+
 def generate_fallback_report(autonomous_context: Dict[str, Any]) -> str:
     """
     Generates a basic fallback report when enhanced generation fails.