Update main.py

main.py

@@ -30,7 +30,45 @@ from fastapi import Request
 from pathlib import Path
 from fastapi.staticfiles import StaticFiles
 
+# main.py
 
+# Standard library imports
+import io
+import re
+import logging
+import tempfile
+import base64
+import warnings
+from typing import Tuple, Optional
+from pathlib import Path
+
+# Third-party imports
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, HTMLResponse
+from transformers import pipeline
+import fitz  # PyMuPDF
+from PIL import Image
+import pandas as pd
+import uvicorn
+from docx import Document
+from pptx import Presentation
+import pytesseract
+from slowapi import Limiter
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+from slowapi.middleware import SlowAPIMiddleware
+import matplotlib.pyplot as plt
+import seaborn as sns
+from pydantic import BaseModel
+import traceback
+import ast
+from openpyxl import Workbook
+
+# Suppress openpyxl warnings
+warnings.filterwarnings("ignore", category=UserWarning, module="openpyxl")
+
+# Rest of your code (app setup, routes, etc.)...
 # Initialize rate limiter
 limiter = Limiter(key_func=get_remote_address)
 
@@ -122,20 +160,37 @@ async def process_uploaded_file(file: UploadFile) -> Tuple[str, bytes]:
     return file_ext, content
 
 def extract_text(content: bytes, file_ext: str) -> str:
-    """Extract text from various file formats with enhanced support"""
+    """Extract text from various file formats with enhanced Excel support"""
     try:
         if file_ext == "docx":
             doc = Document(io.BytesIO(content))
             return "\n".join(para.text for para in doc.paragraphs if para.text.strip())
         
         elif file_ext in {"xlsx", "xls"}:
-            df = pd.read_excel(io.BytesIO(content), sheet_name=None)
+            # Improved Excel handling with better NaN and date support
+            df = pd.read_excel(
+                io.BytesIO(content),
+                sheet_name=None,
+                engine='openpyxl',
+                na_values=['', 'NA', 'N/A', 'NaN', 'null'],
+                keep_default_na=False,
+                parse_dates=True
+            )
+            
             all_text = []
             for sheet_name, sheet_data in df.items():
                 sheet_text = []
+                # Convert all data to string and handle special types
                 for column in sheet_data.columns:
-                    sheet_text.extend(sheet_data[column].dropna().astype(str).tolist())
+                    # Handle datetime columns
+                    if pd.api.types.is_datetime64_any_dtype(sheet_data[column]):
+                        sheet_data[column] = sheet_data[column].dt.strftime('%Y-%m-%d %H:%M:%S')
+                    # Convert to string and clean
+                    col_text = sheet_data[column].astype(str).replace(['nan', 'None', 'NaT'], '').tolist()
+                    sheet_text.extend([x for x in col_text if x.strip()])
+                
                 all_text.append(f"Sheet: {sheet_name}\n" + "\n".join(sheet_text))
+            
             return "\n\n".join(all_text)
         
         elif file_ext == "pptx":
@@ -168,8 +223,8 @@ def extract_text(content: bytes, file_ext: str) -> str:
                 raise ValueError("Could not extract text or caption from image")
         
     except Exception as e:
-        logger.error(f"Text extraction failed for {file_ext}: {str(e)}")
-        raise HTTPException(422, f"Failed to extract text from {file_ext} file")
+        logger.error(f"Text extraction failed for {file_ext}: {str(e)}", exc_info=True)
+        raise HTTPException(422, f"Failed to extract text from {file_ext} file: {str(e)}")
 
 # Visualization Models
 class VisualizationRequest(BaseModel):
@@ -213,47 +268,82 @@ def validate_matplotlib_style(style: str) -> str:
 
     
 def generate_visualization_code(df: pd.DataFrame, request: VisualizationRequest) -> str:
-    """Generate Python code for visualization based on request parameters"""
+    """Generate Python code for visualization with enhanced NaN handling and type safety"""
     # Validate style
     valid_style = validate_matplotlib_style(request.style)
     
+    # Convert DataFrame to dict with proper NaN handling
+    df_dict = df.where(pd.notnull(df), None).to_dict(orient='list')
+    
     code_lines = [
         "import matplotlib.pyplot as plt",
         "import seaborn as sns",
         "import pandas as pd",
+        "import numpy as np",
+        "",
+        "# Data preparation with NaN handling and type conversion",
+        f"raw_data = {df_dict}",
+        "df = pd.DataFrame(raw_data)",
         "",
-        "# Data preparation",
-        f"df = pd.DataFrame({df.to_dict(orient='list')})",
+        "# Automatic type conversion and cleaning",
+        "for col in df.columns:",
+        "    # Convert strings that should be numeric",
+        "    if pd.api.types.is_string_dtype(df[col]):",
+        "        try:",
+        "            df[col] = pd.to_numeric(df[col])",
+        "            continue",
+        "        except (ValueError, TypeError):",
+        "            pass",
+        "    ",
+        "    # Convert string dates to datetime",
+        "    try:",
+        "        df[col] = pd.to_datetime(df[col])",
+        "        continue",
+        "    except (ValueError, TypeError):",
+        "        pass",
+        "    ",
+        "    # Clean remaining None/NaN values",
+        "    df[col] = df[col].where(pd.notnull(df[col]), None)",
     ]
     
-    # Apply filters if specified
+    # Apply filters if specified (with enhanced safety)
     if request.filters:
         filter_conditions = []
         for column, condition in request.filters.items():
             if isinstance(condition, dict):
                 if 'min' in condition and 'max' in condition:
-                    filter_conditions.append(f"(df['{column}'] >= {condition['min']}) & (df['{column}'] <= {condition['max']})")
+                    filter_conditions.append(
+                        f"(pd.notna(df['{column}']) & "
+                        f"(df['{column}'] >= {condition['min']}) & "
+                        f"(df['{column}'] <= {condition['max']})"
+                    )
                 elif 'values' in condition:
                     values = ', '.join([f"'{v}'" if isinstance(v, str) else str(v) for v in condition['values']])
-                    filter_conditions.append(f"df['{column}'].isin([{values}])")
+                    filter_conditions.append(
+                        f"(pd.notna(df['{column}'])) & "
+                        f"(df['{column}'].isin([{values}]))"
+                    )
             else:
-                filter_conditions.append(f"df['{column}'] == {repr(condition)}")
+                filter_conditions.append(
+                    f"(pd.notna(df['{column}'])) & "
+                    f"(df['{column}'] == {repr(condition)})"
+                )
         
         if filter_conditions:
             code_lines.extend([
                 "",
-                "# Apply filters",
-                f"df = df[{' & '.join(filter_conditions)}]"
+                "# Apply filters with NaN checking",
+                f"df = df[{' & '.join(filter_conditions)}].copy()"
             ])
     
     code_lines.extend([
         "",
-        "# Visualization",
+        "# Visualization setup",
         f"plt.style.use('{valid_style}')",
         f"plt.figure(figsize=(10, 6))"
     ])
     
-    # Chart type specific code
+    # Chart type specific code (unchanged from your original)
     if request.chart_type == "line":
         if request.hue_column:
             code_lines.append(f"sns.lineplot(data=df, x='{request.x_column}', y='{request.y_column}', hue='{request.hue_column}')")
@@ -270,15 +360,16 @@ def generate_visualization_code(df: pd.DataFrame, request: VisualizationRequest)
         else:
             code_lines.append(f"plt.scatter(df['{request.x_column}'], df['{request.y_column}'])")
     elif request.chart_type == "histogram":
-        code_lines.append(f"plt.hist(df['{request.x_column}'], bins=20)")
+        code_lines.append(f"plt.hist(df['{request.x_column}'].dropna(), bins=20)")  # Added dropna()
     elif request.chart_type == "boxplot":
         if request.hue_column:
-            code_lines.append(f"sns.boxplot(data=df, x='{request.x_column}', y='{request.y_column}', hue='{request.hue_column}')")
+            code_lines.append(f"sns.boxplot(data=df.dropna(), x='{request.x_column}', y='{request.y_column}', hue='{request.hue_column}')")  # Added dropna()
         else:
-            code_lines.append(f"sns.boxplot(data=df, x='{request.x_column}', y='{request.y_column}')")
+            code_lines.append(f"sns.boxplot(data=df.dropna(), x='{request.x_column}', y='{request.y_column}')")  # Added dropna()
     elif request.chart_type == "heatmap":
-        code_lines.append(f"corr = df.corr()")
-        code_lines.append(f"sns.heatmap(corr, annot=True, cmap='coolwarm')")
+        code_lines.append("numeric_df = df.select_dtypes(include=[np.number])")  # Filter numeric only
+        code_lines.append("corr = numeric_df.corr()")
+        code_lines.append("sns.heatmap(corr, annot=True, cmap='coolwarm')")
     else:
         raise ValueError(f"Unsupported chart type: {request.chart_type}")
     
@@ -296,7 +387,6 @@ def generate_visualization_code(df: pd.DataFrame, request: VisualizationRequest)
     ])
     
     return "\n".join(code_lines)
-
 def interpret_natural_language(prompt: str, df_columns: list) -> VisualizationRequest:
     """Convert natural language prompt to visualization parameters"""
     prompt = prompt.lower()