app.py

import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import openai
from fastapi import FastAPI, UploadFile, File
import io

# Initialize FastAPI app
app = FastAPI()

# Function to generate Python visualization code using Hugging Face model
def generate_viz_code(prompt: str) -> str:
    """Generate Python code for visualization based on user prompt."""
    response = openai.ChatCompletion.create(
        model="mistralai/Mistral-7B",  # Replace with the actual Hugging Face model
        messages=[
            {"role": "system", "content": "You are an AI assistant for data visualization."},
            {"role": "user", "content": prompt}
        ]
    )
    return response["choices"][0]["message"]["content"]

# Function to handle file upload and visualization
@app.post("/visualize")
def visualize_data(file: UploadFile = File(...), prompt: str = ""):
    try:
        # Read the uploaded Excel file
        contents = file.file.read()
        df = pd.read_excel(io.BytesIO(contents))

        # Generate visualization code
        code = generate_viz_code(prompt)
        print("Generated Code:\n", code)  # Debug output

        # Execute the generated code
        exec_globals = {"plt": plt, "sns": sns, "pd": pd, "df": df}
        exec(code, exec_globals)

        # Save the generated plot
        img_path = "visualization.png"
        plt.savefig(img_path)
        plt.close()
        return {"image_path": img_path}
    except Exception as e:
        return {"error": str(e)}

# Uncomment below to run standalone FastAPI app
# if __name__ == "__main__":
#     import uvicorn
#     uvicorn.run(app, host="0.0.0.0", port=8000)