Update main.py

main.py

@@ -2,7 +2,7 @@ from fastapi import FastAPI, UploadFile, File, Form, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse
 from transformers import pipeline
-from typing import Tuple
+from typing import Tuple, Optional
 import io
 import fitz  # PyMuPDF
 from PIL import Image
@@ -22,8 +22,9 @@ import seaborn as sns
 import tempfile
 import base64
 from io import BytesIO
-from typing import Optional
 from pydantic import BaseModel
+import traceback
+import ast
 
 # Initialize rate limiter
 limiter = Limiter(key_func=get_remote_address)
@@ -154,6 +155,151 @@ def extract_text(content: bytes, file_ext: str) -> str:
         logger.error(f"Text extraction failed for {file_ext}: {str(e)}")
         raise HTTPException(422, f"Failed to extract text from {file_ext} file")
 
+# Visualization Models
+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"
+    filters: Optional[dict] = None
+
+class NaturalLanguageRequest(BaseModel):
+    prompt: str
+    style: str = "seaborn"
+
+def generate_visualization_code(df: pd.DataFrame, request: VisualizationRequest) -> str:
+    """Generate Python code for visualization based on request parameters"""
+    code_lines = [
+        "import matplotlib.pyplot as plt",
+        "import seaborn as sns",
+        "import pandas as pd",
+        "",
+        "# Data preparation",
+        f"df = pd.DataFrame({df.to_dict(orient='list')})",
+    ]
+    
+    # Apply filters if specified
+    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']})")
+                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}])")
+            else:
+                filter_conditions.append(f"df['{column}'] == {repr(condition)}")
+        
+        if filter_conditions:
+            code_lines.extend([
+                "",
+                "# Apply filters",
+                f"df = df[{' & '.join(filter_conditions)}]"
+            ])
+    
+    code_lines.extend([
+        "",
+        "# Visualization",
+        f"plt.style.use('{request.style}')",
+        f"plt.figure(figsize=(10, 6))"
+    ])
+    
+    # Chart type specific code
+    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}')")
+        else:
+            code_lines.append(f"plt.plot(df['{request.x_column}'], df['{request.y_column}'])")
+    elif request.chart_type == "bar":
+        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":
+        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.hist(df['{request.x_column}'], bins=20)")
+    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}')")
+        else:
+            code_lines.append(f"sns.boxplot(data=df, x='{request.x_column}', y='{request.y_column}')")
+    elif request.chart_type == "heatmap":
+        code_lines.append(f"corr = df.corr()")
+        code_lines.append(f"sns.heatmap(corr, annot=True, cmap='coolwarm')")
+    else:
+        raise ValueError(f"Unsupported chart type: {request.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.extend([
+        "plt.tight_layout()",
+        "plt.show()"
+    ])
+    
+    return "\n".join(code_lines)
+
+def interpret_natural_language(prompt: str, df_columns: list) -> VisualizationRequest:
+    """Convert natural language prompt to visualization parameters"""
+    # Simple keyword-based interpretation (could be enhanced with NLP)
+    prompt = prompt.lower()
+    
+    # Determine chart type
+    chart_type = "bar"
+    if "line" in prompt:
+        chart_type = "line"
+    elif "scatter" in prompt:
+        chart_type = "scatter"
+    elif "histogram" in prompt:
+        chart_type = "histogram"
+    elif "box" in prompt:
+        chart_type = "boxplot"
+    elif "heatmap" in prompt or "correlation" in prompt:
+        chart_type = "heatmap"
+    
+    # Try to detect columns
+    x_col = None
+    y_col = None
+    hue_col = None
+    
+    for col in df_columns:
+        if col.lower() in prompt:
+            if not x_col:
+                x_col = col
+            elif not y_col:
+                y_col = col
+            else:
+                hue_col = col
+    
+    # Default to first columns if not detected
+    if not x_col and len(df_columns) > 0:
+        x_col = df_columns[0]
+    if not y_col and len(df_columns) > 1:
+        y_col = df_columns[1]
+    
+    return VisualizationRequest(
+        chart_type=chart_type,
+        x_column=x_col,
+        y_column=y_col,
+        hue_column=hue_col,
+        title="Generated from: " + prompt[:50] + ("..." if len(prompt) > 50 else ""),
+        style="seaborn"
+    )
+
 @app.post("/summarize")
 @limiter.limit("5/minute")
 async def summarize_document(request: Request, file: UploadFile = File(...)):
@@ -248,85 +394,9 @@ async def question_answering(
         logger.error(f"QA processing failed: {str(e)}")
         raise HTTPException(500, detail=f"Analysis failed: {str(e)}")
 
-@app.exception_handler(RateLimitExceeded)
-async def rate_limit_exceeded_handler(request: Request, exc: RateLimitExceeded):
-    return JSONResponse(
-        status_code=429,
-        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")
+@app.post("/visualize/code")
 @limiter.limit("5/minute")
-async def generate_visualization_from_excel(
+async def visualize_with_code(
     request: Request,
     file: UploadFile = File(...),
     chart_type: str = Form(...),
@@ -336,18 +406,29 @@ async def generate_visualization_from_excel(
     title: Optional[str] = Form(None),
     x_label: Optional[str] = Form(None),
     y_label: Optional[str] = Form(None),
-    style: str = Form("seaborn")
+    style: str = Form("seaborn"),
+    filters: Optional[str] = Form(None)
 ):
     try:
         # Validate file
-        file_ext, content = await validate_file(file)
+        file_ext, content = await process_uploaded_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
+        # Parse filters if provided
+        filter_dict = {}
+        if filters:
+            try:
+                filter_dict = ast.literal_eval(filters)
+                if not isinstance(filter_dict, dict):
+                    filter_dict = {}
+            except:
+                filter_dict = {}
+        
+        # Create visualization request
         vis_request = VisualizationRequest(
             chart_type=chart_type,
             x_column=x_column,
@@ -356,12 +437,17 @@ async def generate_visualization_from_excel(
             title=title,
             x_label=x_label,
             y_label=y_label,
-            style=style
+            style=style,
+            filters=filter_dict
         )
         
-        # Generate and execute the visualization code
+        # Generate visualization code
+        visualization_code = generate_visualization_code(df, vis_request)
+        
+        # Execute the code to generate the plot
         plt.figure()
-        exec(generate_visualization(df, vis_request), globals(), locals())
+        local_vars = {}
+        exec(visualization_code, globals(), local_vars)
         
         # Save the plot to a temporary file
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmpfile:
@@ -378,7 +464,8 @@ async def generate_visualization_from_excel(
         return {
             "status": "success",
             "image": f"data:image/png;base64,{image_base64}",
-            "code": generate_visualization(df, vis_request)
+            "code": visualization_code,
+            "data_preview": df.head().to_dict(orient='records')
         }
         
     except HTTPException:
@@ -387,7 +474,61 @@ async def generate_visualization_from_excel(
         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("/visualize/natural")
+@limiter.limit("5/minute")
+async def visualize_with_natural_language(
+    request: Request,
+    file: UploadFile = File(...),
+    prompt: str = Form(...),
+    style: str = Form("seaborn")
+):
+    try:
+        # Validate file
+        file_ext, content = await process_uploaded_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))
+        
+        # Convert natural language to visualization parameters
+        nl_request = NaturalLanguageRequest(prompt=prompt, style=style)
+        vis_request = interpret_natural_language(nl_request.prompt, df.columns.tolist())
+        
+        # Generate visualization code
+        visualization_code = generate_visualization_code(df, vis_request)
+        
+        # Execute the code to generate the plot
+        plt.figure()
+        local_vars = {}
+        exec(visualization_code, globals(), local_vars)
+        
+        # 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": visualization_code,
+            "interpreted_parameters": vis_request.dict(),
+            "data_preview": df.head().to_dict(orient='records')
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        logger.error(f"Natural language visualization failed: {str(e)}\n{traceback.format_exc()}")
+        raise HTTPException(500, detail=f"Visualization failed: {str(e)}")
+
 @app.post("/get_columns")
 @limiter.limit("10/minute")
 async def get_excel_columns(
@@ -395,21 +536,26 @@ async def get_excel_columns(
     file: UploadFile = File(...)
 ):
     try:
-        file_ext, content = await validate_file(file)
+        file_ext, content = await process_uploaded_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')
+            "sample_data": df.head().to_dict(orient='records'),
+            "statistics": df.describe().to_dict() if len(df.select_dtypes(include=['number']).columns) > 0 else None
         }
     except Exception as e:
         logger.error(f"Column extraction failed: {str(e)}")
         raise HTTPException(500, detail="Failed to extract columns from Excel file")
 
-
-        
+@app.exception_handler(RateLimitExceeded)
+async def rate_limit_exceeded_handler(request: Request, exc: RateLimitExceeded):
+    return JSONResponse(
+        status_code=429,
+        content={"detail": "Too many requests. Please try again later."}
+    )
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)