Update main.py

main.py

@@ -17,6 +17,13 @@ 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
+import tempfile
+import base64
+from io import BytesIO
+from typing import Optional
+from pydantic import BaseModel
 
 # Initialize rate limiter
 limiter = Limiter(key_func=get_remote_address)
@@ -248,5 +255,161 @@ async def rate_limit_exceeded_handler(request: Request, exc: RateLimitExceeded):
         content={"detail": "Too many requests. Please try again later."}
     )
 
+
+
+
+
+    # Add this new Pydantic model for visualization requests
+class VisualizationRequest(BaseModel):
+    chart_type: str
+    x_column: Optional[str] = None
+    y_column: Optional[str] = None
+    hue_column: Optional[str] = None
+    title: Optional[str] = None
+    x_label: Optional[str] = None
+    y_label: Optional[str] = None
+    style: str = "seaborn"  # seaborn or matplotlib
+
+# Add this new function for visualization code generation
+def generate_visualization(df: pd.DataFrame, request: VisualizationRequest) -> str:
+    """Generate and execute visualization code based on request"""
+    plt.style.use(request.style)
+    
+    code_lines = [
+        "import matplotlib.pyplot as plt",
+        "import seaborn as sns",
+        "import pandas as pd",
+        "",
+        "# Data preparation",
+        f"df = pd.DataFrame({df.head().to_dict()})",  # Simplified for demo
+        "",
+        "# Visualization code"
+    ]
+    
+    if request.chart_type == "line":
+        code_lines.append(f"plt.figure(figsize=(10, 6))")
+        if request.hue_column:
+            code_lines.append(f"sns.lineplot(data=df, x='{request.x_column}', y='{request.y_column}', hue='{request.hue_column}')")
+        else:
+            code_lines.append(f"plt.plot(df['{request.x_column}'], df['{request.y_column}'])")
+    elif request.chart_type == "bar":
+        code_lines.append(f"plt.figure(figsize=(10, 6))")
+        if request.hue_column:
+            code_lines.append(f"sns.barplot(data=df, x='{request.x_column}', y='{request.y_column}', hue='{request.hue_column}')")
+        else:
+            code_lines.append(f"plt.bar(df['{request.x_column}'], df['{request.y_column}'])")
+    elif request.chart_type == "scatter":
+        code_lines.append(f"plt.figure(figsize=(10, 6))")
+        if request.hue_column:
+            code_lines.append(f"sns.scatterplot(data=df, x='{request.x_column}', y='{request.y_column}', hue='{request.hue_column}')")
+        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.figure(figsize=(10, 6))")
+        code_lines.append(f"plt.hist(df['{request.x_column}'], bins=20)")
+    else:
+        raise ValueError("Unsupported chart type")
+
+    # Add labels and title
+    if request.title:
+        code_lines.append(f"plt.title('{request.title}')")
+    if request.x_label:
+        code_lines.append(f"plt.xlabel('{request.x_label}')")
+    if request.y_label:
+        code_lines.append(f"plt.ylabel('{request.y_label}')")
+    
+    code_lines.append("plt.tight_layout()")
+    code_lines.append("plt.show()")
+    
+    return "\n".join(code_lines)
+
+# Add this new endpoint for visualization
+@app.post("/visualize")
+@limiter.limit("5/minute")
+async def generate_visualization_from_excel(
+    request: Request,
+    file: UploadFile = File(...),
+    chart_type: str = Form(...),
+    x_column: Optional[str] = Form(None),
+    y_column: Optional[str] = Form(None),
+    hue_column: Optional[str] = Form(None),
+    title: Optional[str] = Form(None),
+    x_label: Optional[str] = Form(None),
+    y_label: Optional[str] = Form(None),
+    style: str = Form("seaborn")
+):
+    try:
+        # Validate file
+        file_ext, content = await validate_file(file)
+        if file_ext not in {"xlsx", "xls"}:
+            raise HTTPException(400, "Only Excel files are supported for visualization")
+        
+        # Read Excel file
+        df = pd.read_excel(io.BytesIO(content))
+        
+        # Generate visualization request
+        vis_request = VisualizationRequest(
+            chart_type=chart_type,
+            x_column=x_column,
+            y_column=y_column,
+            hue_column=hue_column,
+            title=title,
+            x_label=x_label,
+            y_label=y_label,
+            style=style
+        )
+        
+        # Generate and execute the visualization code
+        plt.figure()
+        exec(generate_visualization(df, vis_request), globals(), locals())
+        
+        # Save the plot to a temporary file
+        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
+            plt.savefig(tmpfile.name, format='png', dpi=300)
+            plt.close()
+            
+            # Read the image back as bytes
+            with open(tmpfile.name, "rb") as f:
+                image_bytes = f.read()
+        
+        # Encode image as base64
+        image_base64 = base64.b64encode(image_bytes).decode('utf-8')
+        
+        return {
+            "status": "success",
+            "image": f"data:image/png;base64,{image_base64}",
+            "code": generate_visualization(df, vis_request)
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Visualization failed: {str(e)}\n{traceback.format_exc()}")
+        raise HTTPException(500, detail=f"Visualization failed: {str(e)}")
+
+# Add this new endpoint for getting column names
+@app.post("/get_columns")
+@limiter.limit("10/minute")
+async def get_excel_columns(
+    request: Request,
+    file: UploadFile = File(...)
+):
+    try:
+        file_ext, content = await validate_file(file)
+        if file_ext not in {"xlsx", "xls"}:
+            raise HTTPException(400, "Only Excel files are supported")
+        
+        df = pd.read_excel(io.BytesIO(content))
+        return {
+            "columns": list(df.columns),
+            "sample_data": df.head().to_dict(orient='records')
+        }
+    except Exception as e:
+        logger.error(f"Column extraction failed: {str(e)}")
+        raise HTTPException(500, detail="Failed to extract columns from Excel file")
+
+
+        
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)