Update main.py

main.py

@@ -101,12 +101,58 @@ app.add_middleware(SlowAPIMiddleware)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
+    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Constants
+
+UPLOAD_FOLDER = "uploads"
+OUTPUT_FOLDER = "static"
+os.makedirs(UPLOAD_FOLDER, exist_ok=True)
+os.makedirs(OUTPUT_FOLDER, exist_ok=True)
+
+# Lightweight model configuration
+MODEL_NAME = "distilgpt2"
+TIMEOUT = 10  # seconds
+MAX_ROWS = 100
+MAX_COLUMNS = 5
 MAX_FILE_SIZE = 10 * 1024 * 1024  # 10MB
+try:
+    visualization_model = pipeline(
+        "text-generation",
+        model=MODEL_NAME,
+        device=-1,  # CPU
+        framework="pt"
+    )
+except Exception as e:
+    print(f"Model loading failed: {str(e)}")
+    visualization_model = None
+
+executor = ThreadPoolExecutor(max_workers=2)
+
+def safe_read_file(file_content, file_ext):
+    """Robust file reading with size limits"""
+    file_like = io.BytesIO(file_content)
+    if file_ext == 'csv':
+        return pd.read_csv(file_like, nrows=MAX_ROWS)
+    return pd.read_excel(file_like, nrows=MAX_ROWS)
+
+def generate_simple_plot(df, chart_type):
+    """Fallback plotting function"""
+    plt.figure(figsize=(8, 5))
+    numeric_cols = df.select_dtypes(include='number').columns
+    
+    if len(numeric_cols) >= 2:
+        df[numeric_cols[:2]].plot(kind=chart_type if chart_type in ['bar', 'line', 'scatter'] else 'bar')
+    elif len(numeric_cols) == 1:
+        df[numeric_cols[0]].plot(kind='bar')
+    else:
+        df.iloc[:, 0].value_counts().plot(kind='bar')
+    
+    plt.tight_layout()
+
+    
 SUPPORTED_FILE_TYPES = {
     "docx", "xlsx", "pptx", "pdf", "jpg", "jpeg", "png", "txt" 
 }
@@ -498,6 +544,146 @@ def validate_french_response(text: str) -> str:
     return text.capitalize()
 
 
+
+
+@app.post("/generate-visualization")
+async def generate_visualization(
+    file: UploadFile = File(...),
+    request: str = Form(...),
+    chart_type: Optional[str] = Form("auto")
+):
+    try:
+        # 1. Validate input
+        file_ext = file.filename.split('.')[-1].lower()
+        if file_ext not in ['csv', 'xlsx', 'xls']:
+            raise HTTPException(400, "Only CSV/Excel files accepted")
+
+        file_content = await file.read()
+        if len(file_content) > MAX_FILE_SIZE:
+            raise HTTPException(400, f"File size exceeds {MAX_FILE_SIZE//1024}KB limit")
+
+        # 2. Process data
+        df = await asyncio.get_event_loop().run_in_executor(
+            executor, 
+            lambda: safe_read_file(file_content, file_ext)
+        )
+        
+        # Simplify dataframe
+        df = df.iloc[:, :MAX_COLUMNS].dropna(how='all')
+        if df.empty:
+            raise HTTPException(400, "No plottable data found")
+
+        # 3. Generate visualization
+        plt.switch_backend('Agg')
+        generated_code = None
+        
+        if visualization_model:
+            try:
+                prompt = f"Create {chart_type} chart for {list(df.columns)}. Python code only:"
+                code = visualization_model(
+                    prompt,
+                    max_length=300,
+                    num_return_sequences=1,
+                    temperature=0.3
+                )[0]['generated_text'].split("```python")[-1].split("```")[0].strip()
+                
+                if code:
+                    generated_code = code
+                    exec(code, {'df': df, 'plt': plt})
+            except Exception as e:
+                print(f"Model failed, using fallback: {e}")
+                generate_simple_plot(df, chart_type)
+                numeric_cols = df.select_dtypes(include='number').columns.tolist()
+                if len(numeric_cols) >= 2:
+                    cols = numeric_cols[:2]
+                    generated_code = f"""
+import pandas as pd
+import matplotlib.pyplot as plt
+data = {df[cols].to_dict()}
+df = pd.DataFrame(data)
+df.plot(kind='{chart_type if chart_type in ['bar', 'line', 'scatter'] else 'bar'}')
+plt.tight_layout()
+plt.show()
+"""
+                elif len(numeric_cols) == 1:
+                    generated_code = f"""
+import pandas as pd
+import matplotlib.pyplot as plt
+data = {df[numeric_cols[0]].to_dict()}
+df = pd.DataFrame(data)
+df.plot(kind='bar')
+plt.tight_layout()
+plt.show()
+"""
+                else:
+                    generated_code = f"""
+import pandas as pd
+import matplotlib.pyplot as plt
+data = {df.iloc[:, 0].value_counts().to_dict()}
+df = pd.DataFrame(list(data.items()), columns=['Category', 'Count'])
+df.plot(x='Category', y='Count', kind='bar')
+plt.tight_layout()
+plt.show()
+"""
+        else:
+            generate_simple_plot(df, chart_type)
+            numeric_cols = df.select_dtypes(include='number').columns.tolist()
+            if len(numeric_cols) >= 2:
+                cols = numeric_cols[:2]
+                generated_code = f"""
+import pandas as pd
+import matplotlib.pyplot as plt
+data = {df[cols].to_dict()}
+df = pd.DataFrame(data)
+df.plot(kind='{chart_type if chart_type in ['bar', 'line', 'scatter'] else 'bar'}')
+plt.tight_layout()
+plt.show()
+"""
+            elif len(numeric_cols) == 1:
+                generated_code = f"""
+import pandas as pd
+import matplotlib.pyplot as plt
+data = {df[numeric_cols[0]].to_dict()}
+df = pd.DataFrame(data)
+df.plot(kind='bar')
+plt.tight_layout()
+plt.show()
+"""
+            else:
+                generated_code = f"""
+import pandas as pd
+import matplotlib.pyplot as plt
+data = {df.iloc[:, 0].value_counts().to_dict()}
+df = pd.DataFrame(list(data.items()), columns=['Category', 'Count'])
+df.plot(x='Category', y='Count', kind='bar')
+plt.tight_layout()
+plt.show()
+"""
+
+        # 4. Save output
+        output_id = uuid.uuid4().hex[:8]
+        image_path = f"{OUTPUT_FOLDER}/plot_{output_id}.png"
+        plt.savefig(image_path, bbox_inches='tight', dpi=80)
+        plt.close()
+
+        return JSONResponse({
+            "image_url": f"/static/plot_{output_id}.png",
+            "python_code": generated_code,
+            "columns": list(df.columns),
+            "note": "Visualization generated successfully"
+        })
+
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(500, f"Processing error: {str(e)}")
+
+@app.get("/static/{filename}")
+async def serve_static(filename: str):
+    file_path = f"{OUTPUT_FOLDER}/{filename}"
+    if not os.path.exists(file_path):
+        raise HTTPException(404, "Image not found")
+    return FileResponse(file_path)