Update app.py

app.py

@@ -1,59 +1,133 @@
-import gradio as gr
+from fastapi import FastAPI, File, UploadFile, HTTPException, Form
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
 import pandas as pd
 import matplotlib.pyplot as plt
-import io
-import base64
-import google.generativeai as genai
-
-def process_file(api_key, file, instructions):
-    # Set up Gemini API
-    genai.configure(api_key=api_key)
-    model = genai.GenerativeModel('gemini-2.5-pro-preview-03-25')
-
-    # Read the file
-    if file.name.endswith('.csv'):
-        df = pd.read_csv(file.name)
-    else:
-        df = pd.read_excel(file.name)
-
-    # Analyze data and get visualization suggestions from Gemini
-    data_description = df.describe().to_string()
-    prompt = f"Given this data: {data_description}\n"
-    if instructions:
-        prompt += f"And these instructions: {instructions}\n"
-    prompt += "Suggest 3 ways to visualize this data."
-
-    response = model.generate_content(prompt)
-    suggestions = response.text.split('\n')
-
-    # Generate visualizations
-    visualizations = []
-    for i, suggestion in enumerate(suggestions[:3]):
-        plt.figure(figsize=(10, 6))
-        plt.title(f"Visualization {i+1}")
-        plt.text(0.5, 0.5, suggestion, ha='center', va='center', wrap=True)
-        buf = io.BytesIO()
-        plt.savefig(buf, format='png')
-        buf.seek(0)
-        img_str = base64.b64encode(buf.getvalue()).decode()
-        visualizations.append(f"data:image/png;base64,{img_str}")
+import seaborn as sns
+import os
+import logging
+from huggingface_hub import InferenceClient
+from dotenv import load_dotenv
+import hashlib
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+load_dotenv()
+
+app = FastAPI()
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+app.mount("/static", StaticFiles(directory="static"), name="static")
+
+API_TOKEN = os.getenv("HF_TOKEN")
+if not API_TOKEN:
+    raise ValueError("HF_TOKEN environment variable not set.")
+
+MODEL_NAME = "gemini-2.5-pro-preview-03-25"
+client = InferenceClient(model=MODEL_NAME, token=API_TOKEN)
+
+UPLOAD_DIR = "uploads"
+os.makedirs(UPLOAD_DIR, exist_ok=True)
+
+IMAGES_DIR = os.path.join("static", "images")
+os.makedirs(IMAGES_DIR, exist_ok=True)
+
+@app.post("/upload/")
+async def upload_file(file: UploadFile = File(...)):
+    if not file.filename.endswith((".xlsx", ".csv")):
+        raise HTTPException(status_code=400, detail="File must be an Excel (.xlsx) or CSV file")
+
+    file_path = os.path.join(UPLOAD_DIR, file.filename)
+    with open(file_path, "wb") as buffer:
+        buffer.write(await file.read())
+
+    logger.info(f"File uploaded: {file.filename}")
+    return {"filename": file.filename}
+
+@app.post("/generate-visualization/")
+async def generate_visualization(prompt: str = Form(...), filename: str = Form(...)):
+    file_path = os.path.join(UPLOAD_DIR, filename)
+
+    if not os.path.exists(file_path):
+        raise HTTPException(status_code=404, detail="File not found on server.")
+
+    try:
+        if filename.endswith('.csv'):
+            df = pd.read_csv(file_path)
+        else:
+            df = pd.read_excel(file_path)
+        if df.empty:
+            raise ValueError("File is empty.")
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error reading file: {str(e)}")
+
+    input_text = f"""
+    Given the DataFrame 'df' with columns {', '.join(df.columns)} and preview:
+    {df.head().to_string()}
+    Write Python code to: {prompt}
+    - Use ONLY 'df' (no external data loading).
+    - Use pandas (pd), matplotlib.pyplot (plt), or seaborn (sns).
+    - Include axis labels and a title.
+    - Output ONLY executable code (no comments, functions, print, or triple quotes).
+    """
+
+    try:
+        generated_code = client.text_generation(input_text, max_new_tokens=500)
+        logger.info(f"Generated code:\n{generated_code}")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error querying model: {str(e)}")
+
+    if not generated_code.strip():
+        raise HTTPException(status_code=500, detail="No code generated by the AI model.")
+
+    generated_code = generated_code.strip()
+    if generated_code.startswith('"""') or generated_code.startswith("'''"):
+        generated_code = generated_code.split('"""')[1] if '"""' in generated_code else generated_code.split("'''")[1]
+    if generated_code.endswith('"""') or generated_code.endswith("'''"):
+        generated_code = generated_code.rsplit('"""')[0] if '"""' in generated_code else generated_code.rsplit("'''")[0]
+    generated_code = generated_code.strip()
+
+    lines = generated_code.splitlines()
+    executable_code = "\n".join(
+        line.strip() for line in lines
+        if line.strip() and not line.strip().startswith(('#', 'def', 'class', '"""', "'''"))
+        and not any(kw in line for kw in ["pd.read_csv", "pd.read_excel", "http", "raise", "print"])
+    ).strip()
+
+    executable_code = executable_code.replace("plt.show()", "").strip()
+
+    logger.info(f"Executable code:\n{executable_code}")
+
+    plot_hash = hashlib.md5(f"{filename}_{prompt}".encode()).hexdigest()[:8]
+    plot_filename = f"plot_{plot_hash}.png"
+    plot_path = os.path.join(IMAGES_DIR, plot_filename)
+
+    try:
+        exec_globals = {"pd": pd, "plt": plt, "sns": sns, "df": df}
+        exec(executable_code, exec_globals)
+        plt.savefig(plot_path, bbox_inches="tight")
         plt.close()
+    except Exception as e:
+        logger.error(f"Error executing code:\n{executable_code}\nException: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error executing code: {str(e)}")
+
+    if not os.path.exists(plot_path):
+        raise HTTPException(status_code=500, detail="Plot file was not created.")
+
+    return {"plot_url": f"/static/images/{plot_filename}", "generated_code": generated_code}
 
-    return visualizations
-
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("Data Visualization with Gemini")
-    api_key = gr.Textbox(label="Enter Gemini API Key", type="password")
-    file = gr.File(label="Upload Excel or CSV file")
-    instructions = gr.Textbox(label="Optional visualization instructions")
-    submit = gr.Button("Generate Visualizations")
-    outputs = [gr.Image(label=f"Visualization {i+1}") for i in range(3)]
-    
-    submit.click(
-        fn=process_file,
-        inputs=[api_key, file, instructions],
-        outputs=outputs
-    )
-
-demo.launch(share=True)
+@app.get("/")
+async def serve_frontend():
+    with open("static/index.html", "r") as f:
+        return HTMLResponse(content=f.read())