Update sozo_gen.py

sozo_gen.py

@@ -1,4 +1,5 @@
 # sozo_gen.py
+# sozo_gen.py
 
 import os
 import re
@@ -77,7 +78,7 @@ def clean_narration(txt: str) -> str:
 def placeholder_img() -> Image.Image: return Image.new("RGB", (WIDTH, HEIGHT), (230, 230, 230))
 
 def generate_image_from_prompt(prompt: str) -> Image.Image:
-    model_main = "gemini-1.5-flash-latest";
+    model_main = "gemini-2.0-flash-exp";
     full_prompt = "A clean business-presentation illustration: " + prompt
     try:
         model = genai.GenerativeModel(model_main)
@@ -309,13 +310,35 @@ def analyze_data_intelligence(df: pd.DataFrame, ctx_dict: Dict) -> Dict[str, Any
     # Statistical Profile
     statistical_summary = {}
     if numeric_cols:
-        statistical_summary = {
-            'correlations': df[numeric_cols].corr().abs().max().to_dict(),
-            'distributions': {col: 'normal' if stats.normaltest(df[col].dropna())[1] > 0.05 else 'non_normal' 
-                           for col in numeric_cols if len(df[col].dropna()) > 8},
-            'outliers': {col: len(df[col][np.abs(stats.zscore(df[col].dropna())) > 3]) 
-                        for col in numeric_cols if len(df[col].dropna()) > 0}
-        }
+        try:
+            correlations = df[numeric_cols].corr().abs().max()
+            correlations_dict = {k: float(v) if pd.notna(v) else 0.0 for k, v in correlations.to_dict().items()}
+            
+            distributions = {}
+            for col in numeric_cols:
+                if len(df[col].dropna()) > 8:
+                    try:
+                        p_value = stats.normaltest(df[col].dropna())[1]
+                        distributions[col] = 'normal' if p_value > 0.05 else 'non_normal'
+                    except:
+                        distributions[col] = 'unknown'
+            
+            outliers = {}
+            for col in numeric_cols:
+                if len(df[col].dropna()) > 0:
+                    try:
+                        z_scores = np.abs(stats.zscore(df[col].dropna()))
+                        outliers[col] = int(len(df[col][z_scores > 3]))
+                    except:
+                        outliers[col] = 0
+            
+            statistical_summary = {
+                'correlations': correlations_dict,
+                'distributions': distributions,
+                'outliers': outliers
+            }
+        except Exception as e:
+            statistical_summary = {'error': 'Could not compute statistical summary'}
     
     # Pattern Detection
     patterns = {
@@ -394,6 +417,23 @@ def get_narrative_suggestions(domain: str, opportunities: List[str], patterns: D
     
     return narrative_frameworks.get(domain, narrative_frameworks['general'])
 
+def json_serializable(obj):
+    """Convert objects to JSON-serializable format"""
+    if isinstance(obj, (np.integer, np.floating)):
+        return float(obj)
+    elif isinstance(obj, np.ndarray):
+        return obj.tolist()
+    elif isinstance(obj, (np.bool_, bool)):
+        return bool(obj)
+    elif isinstance(obj, dict):
+        return {k: json_serializable(v) for k, v in obj.items()}
+    elif isinstance(obj, (list, tuple)):
+        return [json_serializable(item) for item in obj]
+    elif pd.isna(obj):
+        return None
+    else:
+        return obj
+
 def create_autonomous_prompt(df: pd.DataFrame, enhanced_ctx: Dict, intelligence: Dict) -> str:
     """
     Generate a dynamic, intelligence-driven prompt that creates compelling narratives
@@ -407,15 +447,18 @@ def create_autonomous_prompt(df: pd.DataFrame, enhanced_ctx: Dict, intelligence:
     # Dynamic chart strategy based on data characteristics
     chart_strategy = generate_chart_strategy(intelligence)
     
+    # Make context JSON serializable
+    serializable_ctx = json_serializable(enhanced_ctx)
+    
     prompt = f"""You are an elite data storyteller with deep expertise in {domain} analytics. Your mission is to uncover the compelling narrative hidden in this dataset and present it as a captivating story that drives action.
 
 **THE DATA'S STORY CONTEXT:**
-{json.dumps(enhanced_ctx, indent=2)}
+{json.dumps(serializable_ctx, indent=2)}
 
 **INTELLIGENCE ANALYSIS:**
 - Primary Domain: {domain}
 - Key Opportunities: {', '.join(opportunities)}
-- Data Characteristics: {intelligence['data_structure']}
+- Data Characteristics: {json_serializable(intelligence['data_structure'])}
 - Narrative Framework: {narrative['structure']}
 
 **YOUR STORYTELLING MISSION:**
@@ -488,20 +531,38 @@ def enhance_data_context(df: pd.DataFrame, ctx_dict: Dict) -> Dict[str, Any]:
     if not df.empty:
         numeric_cols = df.select_dtypes(include=[np.number]).columns
         if len(numeric_cols) > 0:
+            key_metrics = {}
+            for col in numeric_cols[:3]:  # Top 3 numeric columns
+                try:
+                    mean_val = df[col].mean()
+                    std_val = df[col].std()
+                    key_metrics[col] = {
+                        'mean': float(mean_val) if pd.notna(mean_val) else 0.0,
+                        'std': float(std_val) if pd.notna(std_val) else 0.0
+                    }
+                except:
+                    key_metrics[col] = {'mean': 0.0, 'std': 0.0}
+            
             enhanced['statistical_summary'] = {
-                'numeric_columns': len(numeric_cols),
-                'total_records': len(df),
-                'missing_data_percentage': (df.isnull().sum().sum() / (len(df) * len(df.columns))) * 100,
-                'key_metrics': {col: {'mean': df[col].mean(), 'std': df[col].std()} 
-                               for col in numeric_cols[:3]}  # Top 3 numeric columns
+                'numeric_columns': int(len(numeric_cols)),
+                'total_records': int(len(df)),
+                'missing_data_percentage': float((df.isnull().sum().sum() / (len(df) * len(df.columns))) * 100),
+                'key_metrics': key_metrics
             }
     
     # Add categorical context
     categorical_cols = df.select_dtypes(include=['object', 'category']).columns
     if len(categorical_cols) > 0:
+        unique_values = {}
+        for col in categorical_cols[:3]:
+            try:
+                unique_values[col] = int(df[col].nunique())
+            except:
+                unique_values[col] = 0
+        
         enhanced['categorical_summary'] = {
-            'categorical_columns': len(categorical_cols),
-            'unique_values': {col: df[col].nunique() for col in categorical_cols[:3]}
+            'categorical_columns': int(len(categorical_cols)),
+            'unique_values': unique_values
         }
     
     # Merge with intelligence analysis
@@ -614,237 +675,6 @@ def identify_key_metrics(df: pd.DataFrame, domain: str) -> List[str]:
         key_metrics = variances.head(3).index.tolist()
     
     return key_metrics[:5]  # Return top 5 key metrics
-
-
-# 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.
-    """
-    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